Repository: deepset-ai/haystack

Issue:
enhancement: Agent - refactor/simplify main agent logic
### Why
Over time, the `run`/`run_async` methods of the Agent component have become long, complex and difficult to read.

This makes the code harder to maintain and intimidating to modify.

### What 
I would focus on making the `run`/`run_async` methods shorter and more readable, by extracting some parts into helper methods or functions.

Possible candidates to convert into helper functions: all breakpoint handling, components execution, others...

This would make the main logic shorter and easier to follow, while hiding the complexity of breakpoint settings. It would also improve maintainability and testability.

**Heads-up**: This issue is intended for internal work, so community contributors do not need to submit pull requests for it.

Example ⬇️ 


Discussion:
### Example (untested)

Still a bit long but 85 lines vs 170 lines of the [current `run` method](https://github.com/deepset-ai/haystack/blob/f8d3a82997d48f4f033387d06ec77e78582746f0/haystack/components/agents/agent.py#L241)

```python
def run(self, messages: list[ChatMessage], streaming_callback: Optional[StreamingCallbackT] = None,
         *, break_point: Optional[AgentBreakpoint] = None, snapshot: Optional[AgentSnapshot] = None,
         **kwargs: Any) -> dict[str, Any]:
    
    parent_snapshot = kwargs.pop("parent_snapshot", None)
    
    if not self._is_warmed_up and hasattr(self.chat_generator, "warm_up"):
        raise RuntimeError("The component Agent wasn't warmed up. Run 'warm_up()' before calling 'run()'.")
    
    if break_point and snapshot:
        raise ValueError("break_point and snapshot cannot be provided at the same time. The agent run will be aborted.")
    
    if break_point and isinstance(break_point.break_point, ToolBreakpoint):
        _validate_tool_breakpoint_is_valid(agent_breakpoint=break_point, tools=self.tools)
    
    # Initialize execution state
    if snapshot:
        component_visits, state_data, messages, skip_chat_generator, streaming_callback = self._initialize_from_snapshot(snapshot)
    else:
        component_visits, state_data, messages, skip_chat_generator, streaming_callback = self._initialize_new_execution(messages, **kwargs)
    
    state = State(schema=self.state_schema, data=state_data)
    state.set("messages", messages)
    
    streaming_callback = select_streaming_callback(
        init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=False
    )
    generator_inputs = self._prepare_generator_inputs(streaming_callback=streaming_callback)
    
    # Main execution loop
    with self._create_agent_span() as span:
        span.set_content_tag(
            "haystack.agent.input",
            _deepcopy_with_exceptions({"messages": messages, "streaming_callback": streaming_callback, **kwargs}),
        )
        counter = 0
        
        while counter < self.max_agent_steps:
            # Check for breakpoint before ChatGenerator
            self._handle_breakpoint_before_chat_generator(
                break_point, component_visits, messages, generator_inputs, 
                state, streaming_callback, parent_snapshot
            )
            
            # 1. Call the ChatGenerator
            if skip_chat_generator:
                llm_messages = state.get("messages", [])[-1:]
                skip_chat_generator = False
            else:
                llm_messages = self._execute_chat_generator(messages, generator_inputs, component_visits, state, span)
            
            # 2. Check if we should continue with tools
            if not any(msg.tool_call for msg in llm_messages) or self._tool_invoker is None:
                counter += 1
                break
            
            # Check for breakpoint before ToolInvoker
            self._handle_breakpoint_before_tool_invoker(
                break_point, component_visits, llm_messages, generator_inputs,
                state, streaming_callback, parent_snapshot
            )
            
            # 3. Call the ToolInvoker
            messages, state = self._execute_tool_invoker(
                llm_messages, state, streaming_callback, component_visits, span
            )
            
            # 4. Check exit conditions
            if self.exit_conditions != ["text"] and self._check_exit_conditions(llm_messages, messages):
                counter += 1
                break
            
            counter += 1
        
        # Finalize and return results
        if counter >= self.max_agent_steps:
            logger.warning(
                "Agent reached maximum agent steps of {max_agent_steps}, stopping.",
                max_agent_steps=self.max_agent_steps,
            )
        
        span.set_content_tag("haystack.agent.output", state.data)
        span.set_tag("haystack.agent.steps_taken", counter)
    
    return self._build_result(state)
```



Patch:
diff --git a/haystack/components/agents/agent.py b/haystack/components/agents/agent.py
index 8234a2546a..56a006018e 100644
--- a/haystack/components/agents/agent.py
+++ b/haystack/components/agents/agent.py
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import inspect
+from dataclasses import dataclass
 from typing import Any, Optional, Union
 
 from haystack import logging, tracing
@@ -11,16 +12,16 @@
 from haystack.core.component.component import component
 from haystack.core.pipeline.async_pipeline import AsyncPipeline
 from haystack.core.pipeline.breakpoint import (
-    _check_chat_generator_breakpoint,
-    _check_tool_invoker_breakpoint,
     _create_agent_snapshot,
+    _handle_tool_invoker_breakpoint,
+    _trigger_chat_generator_breakpoint,
     _validate_tool_breakpoint_is_valid,
 )
 from haystack.core.pipeline.pipeline import Pipeline
 from haystack.core.pipeline.utils import _deepcopy_with_exceptions
 from haystack.core.serialization import component_to_dict, default_from_dict, default_to_dict
 from haystack.dataclasses import ChatMessage, ChatRole
-from haystack.dataclasses.breakpoints import AgentBreakpoint, AgentSnapshot, ToolBreakpoint
+from haystack.dataclasses.breakpoints import AgentBreakpoint, AgentSnapshot, PipelineSnapshot, ToolBreakpoint
 from haystack.dataclasses.streaming_chunk import StreamingCallbackT, select_streaming_callback
 from haystack.tools import Tool, Toolset, deserialize_tools_or_toolset_inplace, serialize_tools_or_toolset
 from haystack.utils import _deserialize_value_with_schema
@@ -33,6 +34,28 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class _ExecutionContext:
+    """
+    Context for executing the agent.
+
+    :param state: The current state of the agent, including messages and any additional data.
+    :param component_visits: A dictionary tracking how many times each component has been visited.
+    :param chat_generator_inputs: Runtime inputs to be passed to the chat generator.
+    :param tool_invoker_inputs: Runtime inputs to be passed to the tool invoker.
+    :param counter: A counter to track the number of steps taken in the agent's run.
+    :param skip_chat_generator: A flag to indicate whether to skip the chat generator in the next iteration.
+        This is useful when resuming from a ToolBreakpoint where the ToolInvoker needs to be called first.
+    """
+
+    state: State
+    component_visits: dict
+    chat_generator_inputs: dict
+    tool_invoker_inputs: dict
+    counter: int = 0
+    skip_chat_generator: bool = False
+
+
 @component
 class Agent:
     """
@@ -178,11 +201,6 @@ def to_dict(self) -> dict[str, Any]:
 
         :return: Dictionary with serialized data
         """
-        if self.streaming_callback is not None:
-            streaming_callback = serialize_callable(self.streaming_callback)
-        else:
-            streaming_callback = None
-
         return default_to_dict(
             self,
             chat_generator=component_to_dict(obj=self.chat_generator, name="chat_generator"),
@@ -192,7 +210,7 @@ def to_dict(self) -> dict[str, Any]:
             # We serialize the original state schema, not the resolved one to reflect the original user input
             state_schema=_schema_to_dict(self._state_schema),
             max_agent_steps=self.max_agent_steps,
-            streaming_callback=streaming_callback,
+            streaming_callback=serialize_callable(self.streaming_callback) if self.streaming_callback else None,
             raise_on_tool_invocation_failure=self.raise_on_tool_invocation_failure,
             tool_invoker_kwargs=self.tool_invoker_kwargs,
         )
@@ -219,13 +237,6 @@ def from_dict(cls, data: dict[str, Any]) -> "Agent":
 
         return default_from_dict(cls, data)
 
-    def _prepare_generator_inputs(self, streaming_callback: Optional[StreamingCallbackT] = None) -> dict[str, Any]:
-        """Prepare inputs for the chat generator."""
-        generator_inputs: dict[str, Any] = {"tools": self.tools}
-        if streaming_callback is not None:
-            generator_inputs["streaming_callback"] = streaming_callback
-        return generator_inputs
-
     def _create_agent_span(self) -> Any:
         """Create a span for the agent run."""
         return tracing.tracer.trace(
@@ -238,7 +249,193 @@ def _create_agent_span(self) -> Any:
             },
         )
 
-    def run(  # noqa: PLR0915
+    def _initialize_fresh_execution(
+        self,
+        messages: list[ChatMessage],
+        streaming_callback: Optional[StreamingCallbackT],
+        requires_async: bool,
+        **kwargs,
+    ) -> _ExecutionContext:
+        """
+        Initialize execution context for a fresh run of the agent.
+
+        :param messages: List of ChatMessage objects to start the agent with.
+        :param streaming_callback: Optional callback for streaming responses.
+        :param requires_async: Whether the agent run requires asynchronous execution.
+        :param kwargs: Additional data to pass to the State used by the Agent.
+        """
+        if self.system_prompt is not None:
+            messages = [ChatMessage.from_system(self.system_prompt)] + messages
+
+        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.")
+
+        state = State(schema=self.state_schema, data=kwargs)
+        state.set("messages", messages)
+
+        streaming_callback = select_streaming_callback(  # type: ignore[call-overload]
+            init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=requires_async
+        )
+
+        tool_invoker_inputs: dict[str, Any] = {}
+        generator_inputs: dict[str, Any] = {"tools": self.tools}
+        if streaming_callback is not None:
+            tool_invoker_inputs["streaming_callback"] = streaming_callback
+            generator_inputs["streaming_callback"] = streaming_callback
+
+        return _ExecutionContext(
+            state=state,
+            component_visits=dict.fromkeys(["chat_generator", "tool_invoker"], 0),
+            chat_generator_inputs=generator_inputs,
+            tool_invoker_inputs=tool_invoker_inputs,
+        )
+
+    def _initialize_from_snapshot(
+        self, snapshot: AgentSnapshot, streaming_callback: Optional[StreamingCallbackT], requires_async: bool
+    ) -> _ExecutionContext:
+        """
+        Initialize execution context from an AgentSnapshot.
+
+        :param snapshot: An AgentSnapshot containing the state of a previously saved agent execution.
+        :param streaming_callback: Optional callback for streaming responses.
+        :param requires_async: Whether the agent run requires asynchronous execution.
+        """
+        component_visits = snapshot.component_visits
+        current_inputs = {
+            "chat_generator": _deserialize_value_with_schema(snapshot.component_inputs["chat_generator"]),
+            "tool_invoker": _deserialize_value_with_schema(snapshot.component_inputs["tool_invoker"]),
+        }
+
+        state_data = current_inputs["tool_invoker"]["state"].data
+        state = State(schema=self.state_schema, data=state_data)
+
+        if isinstance(snapshot.break_point.break_point, ToolBreakpoint):
+            messages = current_inputs["tool_invoker"]["messages"]
+            skip_chat_generator = True
+        else:
+            messages = current_inputs["chat_generator"]["messages"]
+            skip_chat_generator = False
+
+        state.set("messages", messages)
+
+        streaming_callback = current_inputs["chat_generator"].get("streaming_callback", streaming_callback)
+        streaming_callback = select_streaming_callback(  # type: ignore[call-overload]
+            init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=requires_async
+        )
+
+        tool_invoker_inputs: dict[str, Any] = {}
+        generator_inputs: dict[str, Any] = {"tools": self.tools}
+        if streaming_callback is not None:
+            tool_invoker_inputs["streaming_callback"] = streaming_callback
+            generator_inputs["streaming_callback"] = streaming_callback
+
+        return _ExecutionContext(
+            state=state,
+            component_visits=component_visits,
+            chat_generator_inputs=generator_inputs,
+            tool_invoker_inputs=tool_invoker_inputs,
+            counter=snapshot.break_point.break_point.visit_count,
+            skip_chat_generator=skip_chat_generator,
+        )
+
+    def _runtime_checks(self, break_point: Optional[AgentBreakpoint], snapshot: Optional[AgentSnapshot]) -> None:
+        """
+        Perform runtime checks before running the agent.
+
+        :param break_point: An AgentBreakpoint, can be a Breakpoint for the "chat_generator" or a ToolBreakpoint
+            for "tool_invoker".
+        :param snapshot: An AgentSnapshot containing the state of a previously saved agent execution.
+        :raises RuntimeError: If the Agent component wasn't warmed up before calling `run()`.
+        :raises ValueError: If both break_point and snapshot are provided, or if the break_point is invalid.
+        """
+        if not self._is_warmed_up and hasattr(self.chat_generator, "warm_up"):
+            raise RuntimeError("The component Agent wasn't warmed up. Run 'warm_up()' before calling 'run()'.")
+
+        if break_point and snapshot:
+            raise ValueError(
+                "break_point and snapshot cannot be provided at the same time. The agent run will be aborted."
+            )
+
+        if break_point and isinstance(break_point.break_point, ToolBreakpoint):
+            _validate_tool_breakpoint_is_valid(agent_breakpoint=break_point, tools=self.tools)
+
+    def _check_chat_generator_breakpoint(
+        self,
+        execution_context: _ExecutionContext,
+        break_point: Optional[AgentBreakpoint],
+        parent_snapshot: Optional[PipelineSnapshot],
+    ) -> None:
+        """
+        Check if the chat generator breakpoint should be triggered.
+
+        If the breakpoint should be triggered, create an agent snapshot and trigger the chat generator breakpoint.
+
+        :param execution_context: The current execution context of the agent.
+        :param break_point: An AgentBreakpoint, can be a Breakpoint for the "chat_generator" or a ToolBreakpoint
+            for "tool_invoker".
+        :param parent_snapshot: An optional parent snapshot for the agent execution.
+        """
+        if (
+            break_point
+            and break_point.break_point.component_name == "chat_generator"
+            and execution_context.component_visits["chat_generator"] == break_point.break_point.visit_count
+        ):
+            agent_snapshot = _create_agent_snapshot(
+                component_visits=execution_context.component_visits,
+                agent_breakpoint=break_point,
+                component_inputs={
+                    "chat_generator": {
+                        "messages": execution_context.state.data["messages"],
+                        **execution_context.chat_generator_inputs,
+                    },
+                    "tool_invoker": {
+                        "messages": [],
+                        "state": execution_context.state,
+                        **execution_context.tool_invoker_inputs,
+                    },
+                },
+            )
+            _trigger_chat_generator_breakpoint(agent_snapshot=agent_snapshot, parent_snapshot=parent_snapshot)
+
+    def _check_tool_invoker_breakpoint(
+        self,
+        execution_context: _ExecutionContext,
+        break_point: Optional[AgentBreakpoint],
+        parent_snapshot: Optional[PipelineSnapshot],
+    ) -> None:
+        """
+        Check if the tool invoker breakpoint should be triggered.
+
+        If the breakpoint should be triggered, create an agent snapshot and trigger the tool invoker breakpoint.
+
+        :param execution_context: The current execution context of the agent.
+        :param break_point: An AgentBreakpoint, can be a Breakpoint for the "chat_generator" or a ToolBreakpoint
+            for "tool_invoker".
+        :param parent_snapshot: An optional parent snapshot for the agent execution.
+        """
+        if (
+            break_point
+            and break_point.break_point.component_name == "tool_invoker"
+            and break_point.break_point.visit_count == execution_context.component_visits["tool_invoker"]
+        ):
+            messages = execution_context.state.data["messages"]
+            agent_snapshot = _create_agent_snapshot(
+                component_visits=execution_context.component_visits,
+                agent_breakpoint=break_point,
+                component_inputs={
+                    "chat_generator": {"messages": messages[:-1], **execution_context.chat_generator_inputs},
+                    "tool_invoker": {
+                        "messages": messages[:-1],
+                        "state": execution_context.state,
+                        **execution_context.tool_invoker_inputs,
+                    },
+                },
+            )
+            _handle_tool_invoker_breakpoint(
+                llm_messages=messages[-1:], agent_snapshot=agent_snapshot, parent_snapshot=parent_snapshot
+            )
+
+    def run(
         self,
         messages: list[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
@@ -251,13 +448,12 @@ def run(  # noqa: PLR0915
         Process messages and execute tools until an exit condition is met.
 
         :param messages: List of Haystack ChatMessage objects to process.
-            If a list of dictionaries is provided, each dictionary will be converted to a ChatMessage object.
         :param streaming_callback: A callback that will be invoked when a response is streamed from the LLM.
             The same callback can be configured to emit tool results when a tool is called.
         :param break_point: An AgentBreakpoint, can be a Breakpoint for the "chat_generator" or a ToolBreakpoint
-                           for "tool_invoker".
+            for "tool_invoker".
         :param snapshot: A dictionary containing a snapshot of a previously saved agent execution. The snapshot contains
-                         the relevant information to restart the Agent execution from where it left off.
+            the relevant information to restart the Agent execution from where it left off.
         :param kwargs: Additional data to pass to the State schema used by the Agent.
             The keys must match the schema defined in the Agent's `state_schema`.
         :returns:
@@ -267,169 +463,94 @@ def run(  # noqa: PLR0915
             - Any additional keys defined in the `state_schema`.
         :raises RuntimeError: If the Agent component wasn't warmed up before calling `run()`.
         :raises BreakpointException: If an agent breakpoint is triggered.
-
         """
-        # kwargs can contain the key parent_snapshot.
-        # We pop it here to avoid passing it into State. We explicitly pass it on if a break point is
-        # triggered.
+        # We pop parent_snapshot from kwargs to avoid passing it into State.
         parent_snapshot = kwargs.pop("parent_snapshot", None)
+        agent_inputs = {
+            "messages": messages,
+            "streaming_callback": streaming_callback,
+            "break_point": break_point,
+            "snapshot": snapshot,
+            **kwargs,
+        }
+        self._runtime_checks(break_point=break_point, snapshot=snapshot)
 
-        if not self._is_warmed_up and hasattr(self.chat_generator, "warm_up"):
-            raise RuntimeError("The component Agent wasn't warmed up. Run 'warm_up()' before calling 'run()'.")
-
-        if break_point and snapshot:
-            raise ValueError(
-                "break_point and snapshot cannot be provided at the same time. The agent run will be aborted."
-            )
-
-        # validate breakpoints
-        if break_point and isinstance(break_point.break_point, ToolBreakpoint):
-            _validate_tool_breakpoint_is_valid(agent_breakpoint=break_point, tools=self.tools)
-
-        # Handle agent snapshot if provided
         if snapshot:
-            component_visits = snapshot.component_visits
-            current_inputs = {
-                "chat_generator": _deserialize_value_with_schema(snapshot.component_inputs["chat_generator"]),
-                "tool_invoker": _deserialize_value_with_schema(snapshot.component_inputs["tool_invoker"]),
-            }
-            state_data = current_inputs["tool_invoker"]["state"].data
-            if isinstance(snapshot.break_point.break_point, ToolBreakpoint):
-                # If the break point is a ToolBreakpoint, we need to get the messages from the tool invoker inputs
-                messages = current_inputs["tool_invoker"]["messages"]
-                # Needed to know if we should start with the ToolInvoker or ChatGenerator
-                skip_chat_generator = True
-            else:
-                messages = current_inputs["chat_generator"]["messages"]
-                skip_chat_generator = False
-            # We also load the streaming_callback from the snapshot if it exists
-            streaming_callback = current_inputs["chat_generator"].get("streaming_callback", None)
+            exe_context = self._initialize_from_snapshot(
+                snapshot=snapshot, streaming_callback=streaming_callback, requires_async=False
+            )
         else:
-            skip_chat_generator = False
-            if self.system_prompt is not None:
-                messages = [ChatMessage.from_system(self.system_prompt)] + messages
-
-            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 as the "
-                    "Agent will not perform any actions specific to user input. Consider adding user messages to the "
-                    "input."
-                )
-            component_visits = dict.fromkeys(["chat_generator", "tool_invoker"], 0)
-            state_data = kwargs
-
-        state = State(schema=self.state_schema, data=state_data)
-        state.set("messages", messages)
+            exe_context = self._initialize_fresh_execution(
+                messages=messages, streaming_callback=streaming_callback, requires_async=False, **kwargs
+            )
 
-        streaming_callback = select_streaming_callback(
-            init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=False
-        )
-        generator_inputs = self._prepare_generator_inputs(streaming_callback=streaming_callback)
         with self._create_agent_span() as span:
-            span.set_content_tag(
-                "haystack.agent.input",
-                _deepcopy_with_exceptions({"messages": messages, "streaming_callback": streaming_callback, **kwargs}),
-            )
-            counter = 0
-
-            while counter < self.max_agent_steps:
-                # check for breakpoint before ChatGenerator
-                if (
-                    break_point
-                    and break_point.break_point.component_name == "chat_generator"
-                    and component_visits["chat_generator"] == break_point.break_point.visit_count
-                ):
-                    agent_snapshot = _create_agent_snapshot(
-                        component_visits=component_visits,
-                        agent_breakpoint=break_point,
-                        component_inputs={
-                            "chat_generator": {"messages": messages, **generator_inputs},
-                            "tool_invoker": {"messages": [], "state": state, "streaming_callback": streaming_callback},
-                        },
-                    )
-                    _check_chat_generator_breakpoint(agent_snapshot=agent_snapshot, parent_snapshot=parent_snapshot)
-
-                # 1. Call the ChatGenerator
-                # We skip the chat generator when restarting from a snapshot where we restart at the ToolInvoker.
-                if skip_chat_generator:
-                    llm_messages = state.get("messages", [])[-1:]
-                    # We set it to False to ensure that the next iteration will call the chat generator again
-                    skip_chat_generator = False
+            span.set_content_tag("haystack.agent.input", _deepcopy_with_exceptions(agent_inputs))
+
+            while exe_context.counter < self.max_agent_steps:
+                # Handle breakpoint and ChatGenerator call
+                self._check_chat_generator_breakpoint(
+                    execution_context=exe_context, break_point=break_point, parent_snapshot=parent_snapshot
+                )
+                # We skip the chat generator when restarting from a snapshot from a ToolBreakpoint
+                if exe_context.skip_chat_generator:
+                    llm_messages = exe_context.state.get("messages", [])[-1:]
+                    # Set to False so the next iteration will call the chat generator
+                    exe_context.skip_chat_generator = False
                 else:
                     result = Pipeline._run_component(
                         component_name="chat_generator",
                         component={"instance": self.chat_generator},
-                        inputs={"messages": messages, **generator_inputs},
-                        component_visits=component_visits,
+                        inputs={"messages": exe_context.state.data["messages"], **exe_context.chat_generator_inputs},
+                        component_visits=exe_context.component_visits,
                         parent_span=span,
                     )
                     llm_messages = result["replies"]
-                    state.set("messages", llm_messages)
+                    exe_context.state.set("messages", llm_messages)
 
-                # 2. Check if any of the LLM responses contain a tool call or if the LLM is not using tools
+                # Check if any of the LLM responses contain a tool call or if the LLM is not using tools
                 if not any(msg.tool_call for msg in llm_messages) or self._tool_invoker is None:
-                    counter += 1
+                    exe_context.counter += 1
                     break
 
-                # check for breakpoint before ToolInvoker
-                if (
-                    break_point
-                    and break_point.break_point.component_name == "tool_invoker"
-                    and break_point.break_point.visit_count == component_visits["tool_invoker"]
-                ):
-                    agent_snapshot = _create_agent_snapshot(
-                        component_visits=component_visits,
-                        agent_breakpoint=break_point,
-                        component_inputs={
-                            "chat_generator": {"messages": messages[:-1], **generator_inputs},
-                            "tool_invoker": {
-                                "messages": llm_messages,
-                                "state": state,
-                                "streaming_callback": streaming_callback,
-                            },
-                        },
-                    )
-                    _check_tool_invoker_breakpoint(
-                        llm_messages=llm_messages, agent_snapshot=agent_snapshot, parent_snapshot=parent_snapshot
-                    )
-
-                # 3. Call the ToolInvoker
+                # Handle breakpoint and ToolInvoker call
+                self._check_tool_invoker_breakpoint(
+                    execution_context=exe_context, break_point=break_point, parent_snapshot=parent_snapshot
+                )
                 # We only send the messages from the LLM to the tool invoker
                 tool_invoker_result = Pipeline._run_component(
                     component_name="tool_invoker",
                     component={"instance": self._tool_invoker},
-                    inputs={"messages": llm_messages, "state": state, "streaming_callback": streaming_callback},
-                    component_visits=component_visits,
+                    inputs={"messages": llm_messages, "state": exe_context.state, **exe_context.tool_invoker_inputs},
+                    component_visits=exe_context.component_visits,
                     parent_span=span,
                 )
                 tool_messages = tool_invoker_result["tool_messages"]
-                state = tool_invoker_result["state"]
-                state.set("messages", tool_messages)
+                exe_context.state = tool_invoker_result["state"]
+                exe_context.state.set("messages", tool_messages)
 
-                # 4. Check if any LLM message's tool call name matches an exit condition
+                # Check if any LLM message's tool call name matches an exit condition
                 if self.exit_conditions != ["text"] and self._check_exit_conditions(llm_messages, tool_messages):
-                    counter += 1
+                    exe_context.counter += 1
                     break
 
-                # 5. Fetch the combined messages and send them back to the LLM
-                messages = state.get("messages")
-                counter += 1
+                # Increment the step counter
+                exe_context.counter += 1
 
-            if counter >= self.max_agent_steps:
+            if exe_context.counter >= self.max_agent_steps:
                 logger.warning(
                     "Agent reached maximum agent steps of {max_agent_steps}, stopping.",
                     max_agent_steps=self.max_agent_steps,
                 )
-            span.set_content_tag("haystack.agent.output", state.data)
-            span.set_tag("haystack.agent.steps_taken", counter)
+            span.set_content_tag("haystack.agent.output", exe_context.state.data)
+            span.set_tag("haystack.agent.steps_taken", exe_context.counter)
 
-        result = {**state.data}
-        all_messages = state.get("messages")
-        if all_messages:
-            result.update({"last_message": all_messages[-1]})
+        result = {**exe_context.state.data}
+        if msgs := result.get("messages"):
+            result["last_message"] = msgs[-1]
         return result
 
-    async def run_async(  # noqa: PLR0915
+    async def run_async(
         self,
         messages: list[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
@@ -445,12 +566,13 @@ async def run_async(  # noqa: PLR0915
         asynchronous operations where possible, such as calling the `run_async` method of the ChatGenerator
         if available.
 
-        :param messages: List of chat messages to process
-        :param streaming_callback: An asynchronous callback that will be invoked when a response
-        is streamed from the LLM. The same callback can be configured to emit tool results when a tool is called.
+        :param messages: List of Haystack ChatMessage objects to process.
+        :param streaming_callback: An asynchronous callback that will be invoked when a response is streamed from the
+            LLM. The same callback can be configured to emit tool results when a tool is called.
         :param break_point: An AgentBreakpoint, can be a Breakpoint for the "chat_generator" or a ToolBreakpoint
-                           for "tool_invoker".
-        :param snapshot: A dictionary containing the state of a previously saved agent execution.
+            for "tool_invoker".
+        :param snapshot: A dictionary containing a snapshot of a previously saved agent execution. The snapshot contains
+            the relevant information to restart the Agent execution from where it left off.
         :param kwargs: Additional data to pass to the State schema used by the Agent.
             The keys must match the schema defined in the Agent's `state_schema`.
         :returns:
@@ -461,169 +583,97 @@ async def run_async(  # noqa: PLR0915
         :raises RuntimeError: If the Agent component wasn't warmed up before calling `run_async()`.
         :raises BreakpointException: If an agent breakpoint is triggered.
         """
-        # kwargs can contain the key parent_snapshot.
-        # We pop it here to avoid passing it into State. We explicitly handle it pass it on if a break point is
-        # triggered.
+        # We pop parent_snapshot from kwargs to avoid passing it into State.
         parent_snapshot = kwargs.pop("parent_snapshot", None)
+        agent_inputs = {
+            "messages": messages,
+            "streaming_callback": streaming_callback,
+            "break_point": break_point,
+            "snapshot": snapshot,
+            **kwargs,
+        }
+        self._runtime_checks(break_point=break_point, snapshot=snapshot)
 
-        if not self._is_warmed_up and hasattr(self.chat_generator, "warm_up"):
-            raise RuntimeError("The component Agent wasn't warmed up. Run 'warm_up()' before calling 'run_async()'.")
-
-        if break_point and snapshot:
-            raise ValueError(
-                "break_point and snapshot cannot be provided at the same time. The agent run will be aborted."
-            )
-
-        # validate breakpoints
-        if break_point and isinstance(break_point.break_point, ToolBreakpoint):
-            _validate_tool_breakpoint_is_valid(agent_breakpoint=break_point, tools=self.tools)
-
-        # Handle agent snapshot if provided
         if snapshot:
-            component_visits = snapshot.component_visits
-            current_inputs = {
-                "chat_generator": _deserialize_value_with_schema(snapshot.component_inputs["chat_generator"]),
-                "tool_invoker": _deserialize_value_with_schema(snapshot.component_inputs["tool_invoker"]),
-            }
-            state_data = current_inputs["tool_invoker"]["state"].data
-            if isinstance(snapshot.break_point.break_point, ToolBreakpoint):
-                # If the break point is a ToolBreakpoint, we need to get the messages from the tool invoker inputs
-                messages = current_inputs["tool_invoker"]["messages"]
-                # Needed to know if we should start with the ToolInvoker or ChatGenerator
-                skip_chat_generator = True
-            else:
-                messages = current_inputs["chat_generator"]["messages"]
-                skip_chat_generator = False
-            # We also load the streaming_callback from the snapshot if it exists
-            streaming_callback = current_inputs["chat_generator"].get("streaming_callback", None)
+            exe_context = self._initialize_from_snapshot(
+                snapshot=snapshot, streaming_callback=streaming_callback, requires_async=False
+            )
         else:
-            skip_chat_generator = False
-            if self.system_prompt is not None:
-                messages = [ChatMessage.from_system(self.system_prompt)] + messages
-
-            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 as the "
-                    "Agent will not perform any actions specific to user input. Consider adding user messages to the "
-                    "input."
-                )
-            component_visits = dict.fromkeys(["chat_generator", "tool_invoker"], 0)
-            state_data = kwargs
-
-        state = State(schema=self.state_schema, data=state_data)
-        state.set("messages", messages)
+            exe_context = self._initialize_fresh_execution(
+                messages=messages, streaming_callback=streaming_callback, requires_async=False, **kwargs
+            )
 
-        streaming_callback = select_streaming_callback(
-            init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=True
-        )
-        generator_inputs = self._prepare_generator_inputs(streaming_callback=streaming_callback)
         with self._create_agent_span() as span:
-            span.set_content_tag(
-                "haystack.agent.input",
-                _deepcopy_with_exceptions({"messages": messages, "streaming_callback": streaming_callback, **kwargs}),
-            )
-            counter = 0
-
-            while counter < self.max_agent_steps:
-                # check for breakpoint before ChatGenerator
-                if (
-                    break_point
-                    and break_point.break_point.component_name == "chat_generator"
-                    and component_visits["chat_generator"] == break_point.break_point.visit_count
-                ):
-                    agent_snapshot = _create_agent_snapshot(
-                        component_visits=component_visits,
-                        agent_breakpoint=break_point,
-                        component_inputs={
-                            "chat_generator": {"messages": messages, **generator_inputs},
-                            "tool_invoker": {"messages": [], "state": state, "streaming_callback": streaming_callback},
-                        },
-                    )
-                    _check_chat_generator_breakpoint(agent_snapshot=agent_snapshot, parent_snapshot=parent_snapshot)
-
-                # 1. Call the ChatGenerator
-                # If skip_chat_generator is True, we skip the chat generator and use the messages from the state
-                # This is useful when the agent is resumed from a snapshot where the chat generator already ran.
-                if skip_chat_generator:
-                    llm_messages = state.get("messages", [])[-1:]  # Get the last message from the state
-                    # We set it to False to ensure that the next iteration will call the chat generator again
-                    skip_chat_generator = False
+            span.set_content_tag("haystack.agent.input", _deepcopy_with_exceptions(agent_inputs))
+
+            while exe_context.counter < self.max_agent_steps:
+                # Handle breakpoint and ChatGenerator call
+                self._check_chat_generator_breakpoint(
+                    execution_context=exe_context, break_point=break_point, parent_snapshot=parent_snapshot
+                )
+                # We skip the chat generator when restarting from a snapshot from a ToolBreakpoint
+                if exe_context.skip_chat_generator:
+                    llm_messages = exe_context.state.get("messages", [])[-1:]
+                    # Set to False so the next iteration will call the chat generator
+                    exe_context.skip_chat_generator = False
                 else:
                     result = await AsyncPipeline._run_component_async(
                         component_name="chat_generator",
                         component={"instance": self.chat_generator},
-                        component_inputs={"messages": messages, **generator_inputs},
-                        component_visits=component_visits,
+                        component_inputs={
+                            "messages": exe_context.state.data["messages"],
+                            **exe_context.chat_generator_inputs,
+                        },
+                        component_visits=exe_context.component_visits,
                         parent_span=span,
                     )
                     llm_messages = result["replies"]
-                    state.set("messages", llm_messages)
+                    exe_context.state.set("messages", llm_messages)
 
-                # 2. Check if any of the LLM responses contain a tool call or if the LLM is not using tools
+                # Check if any of the LLM responses contain a tool call or if the LLM is not using tools
                 if not any(msg.tool_call for msg in llm_messages) or self._tool_invoker is None:
-                    counter += 1
+                    exe_context.counter += 1
                     break
 
-                # Check for breakpoint before ToolInvoker
-                if (
-                    break_point
-                    and break_point.break_point.component_name == "tool_invoker"
-                    and break_point.break_point.visit_count == component_visits["tool_invoker"]
-                ):
-                    agent_snapshot = _create_agent_snapshot(
-                        component_visits=component_visits,
-                        agent_breakpoint=break_point,
-                        component_inputs={
-                            "chat_generator": {"messages": messages[:-1], **generator_inputs},
-                            "tool_invoker": {
-                                "messages": llm_messages,
-                                "state": state,
-                                "streaming_callback": streaming_callback,
-                            },
-                        },
-                    )
-                    _check_tool_invoker_breakpoint(
-                        llm_messages=llm_messages, agent_snapshot=agent_snapshot, parent_snapshot=parent_snapshot
-                    )
-
-                # 3. Call the ToolInvoker
+                # Handle breakpoint and ToolInvoker call
+                self._check_tool_invoker_breakpoint(
+                    execution_context=exe_context, break_point=break_point, parent_snapshot=parent_snapshot
+                )
                 # We only send the messages from the LLM to the tool invoker
                 tool_invoker_result = await AsyncPipeline._run_component_async(
                     component_name="tool_invoker",
                     component={"instance": self._tool_invoker},
                     component_inputs={
                         "messages": llm_messages,
-                        "state": state,
-                        "streaming_callback": streaming_callback,
+                        "state": exe_context.state,
+                        **exe_context.tool_invoker_inputs,
                     },
-                    component_visits=component_visits,
+                    component_visits=exe_context.component_visits,
                     parent_span=span,
                 )
                 tool_messages = tool_invoker_result["tool_messages"]
-                state = tool_invoker_result["state"]
-                state.set("messages", tool_messages)
+                exe_context.state = tool_invoker_result["state"]
+                exe_context.state.set("messages", tool_messages)
 
-                # 4. Check if any LLM message's tool call name matches an exit condition
+                # Check if any LLM message's tool call name matches an exit condition
                 if self.exit_conditions != ["text"] and self._check_exit_conditions(llm_messages, tool_messages):
-                    counter += 1
+                    exe_context.counter += 1
                     break
 
-                # 5. Fetch the combined messages and send them back to the LLM
-                messages = state.get("messages")
-                counter += 1
+                # Increment the step counter
+                exe_context.counter += 1
 
-            if counter >= self.max_agent_steps:
+            if exe_context.counter >= self.max_agent_steps:
                 logger.warning(
                     "Agent reached maximum agent steps of {max_agent_steps}, stopping.",
                     max_agent_steps=self.max_agent_steps,
                 )
-            span.set_content_tag("haystack.agent.output", state.data)
-            span.set_tag("haystack.agent.steps_taken", counter)
+            span.set_content_tag("haystack.agent.output", exe_context.state.data)
+            span.set_tag("haystack.agent.steps_taken", exe_context.counter)
 
-        result = {**state.data}
-        all_messages = state.get("messages")
-        if all_messages:
-            result.update({"last_message": all_messages[-1]})
+        result = {**exe_context.state.data}
+        if msgs := result.get("messages"):
+            result["last_message"] = msgs[-1]
         return result
 
     def _check_exit_conditions(self, llm_messages: list[ChatMessage], tool_messages: list[ChatMessage]) -> bool:
diff --git a/haystack/core/pipeline/breakpoint.py b/haystack/core/pipeline/breakpoint.py
index e9d6971c6f..22a604dae8 100644
--- a/haystack/core/pipeline/breakpoint.py
+++ b/haystack/core/pipeline/breakpoint.py
@@ -340,15 +340,15 @@ def _validate_tool_breakpoint_is_valid(agent_breakpoint: AgentBreakpoint, tools:
         raise ValueError(f"Tool '{tool_breakpoint.tool_name}' is not available in the agent's tools")
 
 
-def _check_chat_generator_breakpoint(
+def _trigger_chat_generator_breakpoint(
     *, agent_snapshot: AgentSnapshot, parent_snapshot: Optional[PipelineSnapshot]
 ) -> None:
     """
-    Check for breakpoint before calling the ChatGenerator.
+    Trigger a breakpoint before ChatGenerator execution in Agent.
 
     :param agent_snapshot: AgentSnapshot object containing the agent's state and breakpoints
     :param parent_snapshot: Optional parent snapshot containing the state of the pipeline that houses the agent.
-    :raises BreakpointException: If a breakpoint is triggered
+    :raises BreakpointException: Always raised when this function is called, indicating a breakpoint has been triggered.
     """
 
     break_point = agent_snapshot.break_point.break_point
@@ -381,16 +381,17 @@ def _check_chat_generator_breakpoint(
     )
 
 
-def _check_tool_invoker_breakpoint(
+def _handle_tool_invoker_breakpoint(
     *, llm_messages: list[ChatMessage], agent_snapshot: AgentSnapshot, parent_snapshot: Optional[PipelineSnapshot]
 ) -> None:
     """
-    Check for breakpoint before calling the ToolInvoker.
+    Check if a tool call breakpoint should be triggered before executing the tool invoker.
 
     :param llm_messages: List of ChatMessage objects containing potential tool calls.
     :param agent_snapshot: AgentSnapshot object containing the agent's state and breakpoints.
     :param parent_snapshot: Optional parent snapshot containing the state of the pipeline that houses the agent.
-    :raises BreakpointException: If a breakpoint is triggered
+    :raises BreakpointException: If the breakpoint is triggered, indicating a breakpoint has been reached for a tool
+        call.
     """
     if not isinstance(agent_snapshot.break_point.break_point, ToolBreakpoint):
         return
diff --git a/releasenotes/notes/refactor-agent-b8d3b690f52c3bb5.yaml b/releasenotes/notes/refactor-agent-b8d3b690f52c3bb5.yaml
new file mode 100644
index 0000000000..dccc10e5f9
--- /dev/null
+++ b/releasenotes/notes/refactor-agent-b8d3b690f52c3bb5.yaml
@@ -0,0 +1,4 @@
+---
+enhancements:
+  - |
+    The internal Agent logic was refactored to help with readability and maintanability. This should help developers understand and extend the internal Agent logic moving forward.


Test patch:
diff --git a/test/components/agents/test_agent.py b/test/components/agents/test_agent.py
index 430074e489..9bd6dc701e 100644
--- a/test/components/agents/test_agent.py
+++ b/test/components/agents/test_agent.py
@@ -12,13 +12,12 @@
 from openai import Stream
 from openai.types.chat import ChatCompletionChunk, chat_completion_chunk
 
-from haystack import Pipeline, tracing
+from haystack import Pipeline, component, tracing
 from haystack.components.agents import Agent
 from haystack.components.agents.state import merge_lists
 from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
 from haystack.components.builders.prompt_builder import PromptBuilder
 from haystack.components.generators.chat.openai import OpenAIChatGenerator
-from haystack.components.generators.chat.types import ChatGenerator
 from haystack.core.component.types import OutputSocket
 from haystack.dataclasses import ChatMessage, ToolCall
 from haystack.dataclasses.chat_message import ChatRole, TextContent
@@ -101,12 +100,10 @@ def openai_mock_chat_completion_chunk():
         yield mock_chat_completion_create
 
 
-class MockChatGeneratorWithoutTools(ChatGenerator):
+@component
+class MockChatGeneratorWithoutTools:
     """A mock chat generator that implements ChatGenerator protocol but doesn't support tools."""
 
-    __haystack_input__ = MagicMock(_sockets_dict={})
-    __haystack_output__ = MagicMock(_sockets_dict={})
-
     def to_dict(self) -> dict[str, Any]:
         return {"type": "MockChatGeneratorWithoutTools", "data": {}}
 
@@ -114,16 +111,15 @@ def to_dict(self) -> dict[str, Any]:
     def from_dict(cls, data: dict[str, Any]) -> "MockChatGeneratorWithoutTools":
         return cls()
 
+    @component.output_types(replies=list[ChatMessage])
     def run(self, messages: list[ChatMessage]) -> dict[str, Any]:
         return {"replies": [ChatMessage.from_assistant("Hello")]}
 
 
-class MockChatGeneratorWithoutRunAsync(ChatGenerator):
+@component
+class MockChatGeneratorWithoutRunAsync:
     """A mock chat generator that implements ChatGenerator protocol but doesn't have run_async method."""
 
-    __haystack_input__ = MagicMock(_sockets_dict={})
-    __haystack_output__ = MagicMock(_sockets_dict={})
-
     def to_dict(self) -> dict[str, Any]:
         return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}}
 
@@ -131,17 +127,15 @@ def to_dict(self) -> dict[str, Any]:
     def from_dict(cls, data: dict[str, Any]) -> "MockChatGeneratorWithoutRunAsync":
         return cls()
 
+    @component.output_types(replies=list[ChatMessage])
     def run(
         self, messages: list[ChatMessage], tools: Optional[Union[list[Tool], Toolset]] = None, **kwargs
     ) -> dict[str, Any]:
         return {"replies": [ChatMessage.from_assistant("Hello")]}
 
 
-class MockChatGeneratorWithRunAsync(ChatGenerator):
-    __haystack_supports_async__ = True
-    __haystack_input__ = MagicMock(_sockets_dict={})
-    __haystack_output__ = MagicMock(_sockets_dict={})
-
+@component
+class MockChatGenerator:
     def to_dict(self) -> dict[str, Any]:
         return {"type": "MockChatGeneratorWithoutRunAsync", "data": {}}
 
@@ -149,11 +143,13 @@ def to_dict(self) -> dict[str, Any]:
     def from_dict(cls, data: dict[str, Any]) -> "MockChatGeneratorWithoutRunAsync":
         return cls()
 
+    @component.output_types(replies=list[ChatMessage])
     def run(
         self, messages: list[ChatMessage], tools: Optional[Union[list[Tool], Toolset]] = 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: Optional[Union[list[Tool], Toolset]] = None, **kwargs
     ) -> dict[str, Any]:
@@ -818,7 +814,7 @@ async def test_run_async_falls_back_to_run_when_chat_generator_has_no_run_async(
 
     @pytest.mark.asyncio
     async def test_run_async_uses_chat_generator_run_async_when_available(self, weather_tool):
-        chat_generator = MockChatGeneratorWithRunAsync()
+        chat_generator = MockChatGenerator()
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
         agent.warm_up()
 
@@ -904,8 +900,8 @@ def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool):
             "chat_generator",
             "MockChatGeneratorWithoutRunAsync",
             '{"messages": "list", "tools": "list"}',
-            "{}",
-            "{}",
+            '{"messages": {"type": "list", "senders": []}, "tools": {"type": "typing.Union[list[haystack.tools.tool.Tool], haystack.tools.toolset.Toolset, NoneType]", "senders": []}}',  # noqa: E501
+            '{"replies": {"type": "list", "receivers": []}}',
             '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]}',  # noqa: E501
             1,
             '{"replies": [{"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello"}]}]}',
@@ -913,7 +909,7 @@ def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool):
             '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]',  # noqa: E501
             '["text"]',
             '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}}',  # noqa: E501
-            '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}',  # noqa: E501
+            '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null, "break_point": null, "snapshot": null}',  # noqa: E501
             '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello"}]}]}',  # noqa: E501
             1,
         ]
@@ -927,7 +923,7 @@ def test_agent_tracing_span_run(self, caplog, monkeypatch, weather_tool):
 
     @pytest.mark.asyncio
     async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_tool):
-        chat_generator = MockChatGeneratorWithRunAsync()
+        chat_generator = MockChatGenerator()
         agent = Agent(chat_generator=chat_generator, tools=[weather_tool])
 
         tracing.tracer.is_content_tracing_enabled = True
@@ -962,10 +958,10 @@ async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_t
 
         expected_tag_values = [
             "chat_generator",
-            "MockChatGeneratorWithRunAsync",
+            "MockChatGenerator",
             '{"messages": "list", "tools": "list"}',
-            "{}",
-            "{}",
+            '{"messages": {"type": "list", "senders": []}, "tools": {"type": "typing.Union[list[haystack.tools.tool.Tool], haystack.tools.toolset.Toolset, NoneType]", "senders": []}}',  # noqa: E501
+            '{"replies": {"type": "list", "receivers": []}}',
             '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "tools": [{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]}',  # noqa: E501
             1,
             '{"replies": [{"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}]}',  # noqa: E501
@@ -973,7 +969,7 @@ async def test_agent_tracing_span_async_run(self, caplog, monkeypatch, weather_t
             '[{"type": "haystack.tools.tool.Tool", "data": {"name": "weather_tool", "description": "Provides weather information for a given location.", "parameters": {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}, "function": "test_agent.weather_function", "outputs_to_string": null, "inputs_from_state": null, "outputs_to_state": null}}]',  # noqa: E501
             '["text"]',
             '{"messages": {"type": "list[haystack.dataclasses.chat_message.ChatMessage]", "handler": "haystack.components.agents.state.state_utils.merge_lists"}}',  # noqa: E501
-            '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null}',  # noqa: E501
+            '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}], "streaming_callback": null, "break_point": null, "snapshot": null}',  # noqa: E501
             '{"messages": [{"role": "user", "meta": {}, "name": null, "content": [{"text": "What\'s the weather in Paris?"}]}, {"role": "assistant", "meta": {}, "name": null, "content": [{"text": "Hello from run_async"}]}]}',  # noqa: E501
             1,
         ]
diff --git a/test/components/agents/test_agent_breakpoints.py b/test/components/agents/test_agent_breakpoints.py
new file mode 100644
index 0000000000..764fe2643a
--- /dev/null
+++ b/test/components/agents/test_agent_breakpoints.py
@@ -0,0 +1,270 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from haystack.components.agents import Agent
+from haystack.core.errors import BreakpointException
+from haystack.core.pipeline.breakpoint import load_pipeline_snapshot
+from haystack.dataclasses import ChatMessage, ToolCall
+from haystack.dataclasses.breakpoints import AgentBreakpoint, AgentSnapshot, Breakpoint, ToolBreakpoint
+from haystack.tools import Tool
+from test.components.agents.test_agent import MockChatGenerator, weather_function
+
+
+@pytest.fixture
+def weather_tool():
+    return Tool(
+        name="weather_tool",
+        description="Provides weather information for a given location.",
+        parameters={"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]},
+        function=weather_function,
+    )
+
+
+@pytest.fixture
+def mock_chat_generator():
+    generator = MockChatGenerator()
+    reply = {
+        "replies": [
+            ChatMessage.from_assistant(
+                "I'll help you check the weather.",
+                tool_calls=[{"tool_name": "weather_tool", "tool_args": {"location": "Berlin"}}],
+            )
+        ]
+    }
+
+    def run(messages, tools=None, **kwargs):
+        return reply
+
+    async def run_async(messages, tools=None, **kwargs):
+        return reply
+
+    generator.run = run
+    generator.run_async = run_async
+    return generator
+
+
+@pytest.fixture
+def agent(mock_chat_generator, weather_tool):
+    return Agent(
+        chat_generator=mock_chat_generator,
+        tools=[weather_tool],
+        system_prompt="You are a helpful assistant that can use tools to help users.",
+        max_agent_steps=10,  # Increase max steps to allow breakpoints to trigger
+    )
+
+
+@pytest.fixture
+def mock_agent_with_tool_calls(monkeypatch, weather_tool):
+    monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
+    generator = MockChatGenerator()
+    mock_messages = [
+        ChatMessage.from_assistant("First response"),
+        ChatMessage.from_assistant(tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]),
+    ]
+    agent = Agent(chat_generator=generator, tools=[weather_tool], max_agent_steps=10)  # Increase max steps
+    agent.warm_up()
+    agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages})
+    agent.chat_generator.run_async = AsyncMock(return_value={"replies": mock_messages})
+    return agent
+
+
+class TestAgentBreakpoints:
+    def test_run_with_chat_generator_breakpoint(self, agent):
+        messages = [ChatMessage.from_user("What's the weather in Berlin?")]
+        chat_generator_bp = Breakpoint(component_name="chat_generator", visit_count=0)
+        agent_breakpoint = AgentBreakpoint(break_point=chat_generator_bp, agent_name="test_agent")
+        with pytest.raises(BreakpointException) as exc_info:
+            agent.run(messages=messages, break_point=agent_breakpoint)
+        assert exc_info.value.component == "chat_generator"
+        assert "messages" in exc_info.value.inputs["chat_generator"]["serialized_data"]
+
+    def test_run_with_tool_invoker_breakpoint(self, mock_agent_with_tool_calls):
+        messages = [ChatMessage.from_user("What's the weather in Berlin?")]
+        tool_bp = ToolBreakpoint(component_name="tool_invoker", visit_count=0, tool_name="weather_tool")
+        agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test_agent")
+        with pytest.raises(BreakpointException) as exc_info:
+            mock_agent_with_tool_calls.run(messages=messages, break_point=agent_breakpoint)
+
+        assert exc_info.value.component == "tool_invoker"
+        assert {"chat_generator", "tool_invoker"} == set(exc_info.value.inputs.keys())
+        assert "serialization_schema" in exc_info.value.inputs["chat_generator"]
+        assert "serialized_data" in exc_info.value.inputs["chat_generator"]
+        assert "messages" in exc_info.value.inputs["chat_generator"]["serialized_data"]
+
+    def test_resume_from_chat_generator(self, agent, tmp_path):
+        debug_path = str(tmp_path / "debug_snapshots")
+        chat_generator_bp = Breakpoint(component_name="chat_generator", snapshot_file_path=debug_path)
+        agent_breakpoint = AgentBreakpoint(break_point=chat_generator_bp, agent_name="test_agent")
+
+        try:
+            agent.run(messages=[ChatMessage.from_user("What's the weather in Berlin?")], break_point=agent_breakpoint)
+        except BreakpointException:
+            pass
+
+        snapshot_files = list(Path(debug_path).glob("test_agent_chat_generator_*.json"))
+        assert len(snapshot_files) > 0
+        latest_snapshot_file = str(max(snapshot_files, key=os.path.getctime))
+
+        result = agent.run(
+            messages=[ChatMessage.from_user("Continue from where we left off.")],
+            snapshot=load_pipeline_snapshot(latest_snapshot_file).agent_snapshot,
+        )
+
+        assert "messages" in result
+        assert "last_message" in result
+        assert len(result["messages"]) > 0
+
+    def test_resume_from_tool_invoker(self, mock_agent_with_tool_calls, tmp_path):
+        messages = [ChatMessage.from_user("What's the weather in Berlin?")]
+        debug_path = str(tmp_path / "debug_snapshots")
+        tool_bp = ToolBreakpoint(component_name="tool_invoker", snapshot_file_path=debug_path)
+        agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test_agent")
+
+        try:
+            mock_agent_with_tool_calls.run(messages=messages, break_point=agent_breakpoint)
+        except BreakpointException:
+            pass
+
+        snapshot_files = list(Path(debug_path).glob("test_agent_tool_invoker_*.json"))
+        assert len(snapshot_files) > 0
+        latest_snapshot_file = str(max(snapshot_files, key=os.path.getctime))
+
+        result = mock_agent_with_tool_calls.run(
+            messages=[ChatMessage.from_user("Continue from where we left off.")],
+            snapshot=load_pipeline_snapshot(latest_snapshot_file).agent_snapshot,
+        )
+
+        assert "messages" in result
+        assert "last_message" in result
+        assert len(result["messages"]) > 0
+
+    def test_invalid_combination_breakpoint_and_pipeline_snapshot(self, mock_agent_with_tool_calls):
+        messages = [ChatMessage.from_user("What's the weather in Berlin?")]
+        tool_bp = ToolBreakpoint(component_name="tool_invoker", tool_name="weather_tool")
+        agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test_agent")
+        with pytest.raises(ValueError, match="break_point and snapshot cannot be provided at the same time"):
+            mock_agent_with_tool_calls.run(
+                messages=messages,
+                break_point=agent_breakpoint,
+                snapshot=AgentSnapshot(component_inputs={}, component_visits={}, break_point=agent_breakpoint),
+            )
+
+    def test_breakpoint_with_invalid_component(self, mock_agent_with_tool_calls):
+        invalid_bp = Breakpoint(component_name="invalid_breakpoint")
+        with pytest.raises(ValueError):
+            AgentBreakpoint(break_point=invalid_bp, agent_name="test_agent")
+
+    def test_breakpoint_with_invalid_tool_name(self, mock_agent_with_tool_calls):
+        tool_breakpoint = ToolBreakpoint(component_name="tool_invoker", tool_name="invalid_tool")
+        with pytest.raises(ValueError, match="Tool 'invalid_tool' is not available in the agent's tools"):
+            agent_breakpoints = AgentBreakpoint(break_point=tool_breakpoint, agent_name="test_agent")
+            mock_agent_with_tool_calls.run(
+                messages=[ChatMessage.from_user("What's the weather in Berlin?")], break_point=agent_breakpoints
+            )
+
+
+class TestAsyncAgentBreakpoints:
+    @pytest.mark.asyncio
+    async def test_run_async_with_chat_generator_breakpoint(self, agent):
+        messages = [ChatMessage.from_user("What's the weather in Berlin?")]
+        chat_generator_bp = Breakpoint(component_name="chat_generator")
+        agent_breakpoint = AgentBreakpoint(break_point=chat_generator_bp, agent_name="test_agent")
+        with pytest.raises(BreakpointException) as exc_info:
+            await agent.run_async(messages=messages, break_point=agent_breakpoint)
+        assert exc_info.value.component == "chat_generator"
+        assert "messages" in exc_info.value.inputs["chat_generator"]["serialized_data"]
+
+    @pytest.mark.asyncio
+    async def test_run_async_with_tool_invoker_breakpoint(self, mock_agent_with_tool_calls):
+        messages = [ChatMessage.from_user("What's the weather in Berlin?")]
+        tool_bp = ToolBreakpoint(component_name="tool_invoker", tool_name="weather_tool")
+        agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test")
+        with pytest.raises(BreakpointException) as exc_info:
+            await mock_agent_with_tool_calls.run_async(messages=messages, break_point=agent_breakpoint)
+
+        assert exc_info.value.component == "tool_invoker"
+        assert "messages" in exc_info.value.inputs["tool_invoker"]["serialized_data"]
+
+    @pytest.mark.asyncio
+    async def test_resume_from_chat_generator_async(self, agent, tmp_path):
+        debug_path = str(tmp_path / "debug_snapshots")
+        messages = [ChatMessage.from_user("What's the weather in Berlin?")]
+        chat_generator_bp = Breakpoint(component_name="chat_generator", snapshot_file_path=debug_path)
+        agent_breakpoint = AgentBreakpoint(break_point=chat_generator_bp, agent_name="test_agent")
+
+        try:
+            await agent.run_async(messages=messages, break_point=agent_breakpoint)
+        except BreakpointException:
+            pass
+
+        snapshot_files = list(Path(debug_path).glob("test_agent_chat_generator_*.json"))
+
+        assert len(snapshot_files) > 0
+        latest_snapshot_file = str(max(snapshot_files, key=os.path.getctime))
+
+        result = await agent.run_async(
+            messages=[ChatMessage.from_user("Continue from where we left off.")],
+            snapshot=load_pipeline_snapshot(latest_snapshot_file).agent_snapshot,
+        )
+
+        assert "messages" in result
+        assert "last_message" in result
+        assert len(result["messages"]) > 0
+
+    @pytest.mark.asyncio
+    async def test_resume_from_tool_invoker_async(self, mock_agent_with_tool_calls, tmp_path):
+        debug_path = str(tmp_path / "debug_snapshots")
+        messages = [ChatMessage.from_user("What's the weather in Berlin?")]
+        tool_bp = ToolBreakpoint(component_name="tool_invoker", tool_name="weather_tool", snapshot_file_path=debug_path)
+        agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test_agent")
+
+        try:
+            await mock_agent_with_tool_calls.run_async(messages=messages, break_point=agent_breakpoint)
+        except BreakpointException:
+            pass
+
+        snapshot_files = list(Path(debug_path).glob("test_agent_tool_invoker_*.json"))
+
+        assert len(snapshot_files) > 0
+        latest_snapshot_file = str(max(snapshot_files, key=os.path.getctime))
+
+        result = await mock_agent_with_tool_calls.run_async(
+            messages=[ChatMessage.from_user("Continue from where we left off.")],
+            snapshot=load_pipeline_snapshot(latest_snapshot_file).agent_snapshot,
+        )
+
+        assert "messages" in result
+        assert "last_message" in result
+        assert len(result["messages"]) > 0
+
+    @pytest.mark.asyncio
+    async def test_invalid_combination_breakpoint_and_pipeline_snapshot_async(self, mock_agent_with_tool_calls):
+        tool_bp = ToolBreakpoint(component_name="tool_invoker", visit_count=0, tool_name="weather_tool")
+        agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test")
+        with pytest.raises(ValueError, match="break_point and snapshot cannot be provided at the same time"):
+            await mock_agent_with_tool_calls.run_async(
+                messages=[ChatMessage.from_user("What's the weather in Berlin?")],
+                break_point=agent_breakpoint,
+                snapshot=AgentSnapshot(component_inputs={}, component_visits={}, break_point=agent_breakpoint),
+            )
+
+    @pytest.mark.asyncio
+    async def test_breakpoint_with_invalid_component_async(self, mock_agent_with_tool_calls):
+        with pytest.raises(ValueError):
+            AgentBreakpoint(break_point=Breakpoint(component_name="invalid_breakpoint"), agent_name="test")
+
+    @pytest.mark.asyncio
+    async def test_breakpoint_with_invalid_tool_name_async(self, mock_agent_with_tool_calls):
+        tool_breakpoint = ToolBreakpoint(component_name="tool_invoker", tool_name="invalid_tool")
+        agent_breakpoint = AgentBreakpoint(break_point=tool_breakpoint, agent_name="test")
+        with pytest.raises(ValueError, match="Tool 'invalid_tool' is not available in the agent's tools"):
+            await mock_agent_with_tool_calls.run_async(
+                messages=[ChatMessage.from_user("What's the weather in Berlin?")], break_point=agent_breakpoint
+            )
diff --git a/test/components/agents/test_agent_breakpoints_inside_pipeline.py b/test/components/agents/test_agent_breakpoints_inside_pipeline.py
index 8237e0e746..12d609d14b 100644
--- a/test/components/agents/test_agent_breakpoints_inside_pipeline.py
+++ b/test/components/agents/test_agent_breakpoints_inside_pipeline.py
@@ -75,17 +75,14 @@ def run(self, sources: list[ByteStream]) -> dict[str, list[Document]]:
 
         documents = []
         for source in sources:
-            # Extract the HTML content from the ByteStream
             html_content = source.data.decode("utf-8")
 
             # Simple text extraction - remove HTML tags and extract meaningful content
-            # This is a simplified version that extracts the main content
             # Remove HTML tags
             text_content = re.sub(r"<[^>]+>", " ", html_content)
             # Remove extra whitespace
             text_content = re.sub(r"\s+", " ", text_content).strip()
 
-            # Create a Document with the extracted text
             document = Document(
                 content=text_content,
                 meta={"url": source.meta.get("url", "unknown"), "mime_type": source.mime_type, "source_type": "html"},
@@ -101,10 +98,6 @@ def add_database_tool_function(name: str, surname: str, job_title: Optional[str]
     )
 
 
-# We use this since the @tool decorator has issues with deserialization
-add_database_tool = create_tool_from_function(add_database_tool_function, name="add_database_tool")
-
-
 @pytest.fixture
 def pipeline_with_agent(monkeypatch):
     monkeypatch.setenv("OPENAI_API_KEY", "test_key")
@@ -163,6 +156,9 @@ def mock_run(messages, tools=None, **kwargs):
 
     generator.run = mock_run
 
+    # We use this since the @tool decorator has issues with deserialization
+    add_database_tool = create_tool_from_function(add_database_tool_function, name="add_database_tool")
+
     database_assistant = Agent(
         chat_generator=generator,
         tools=[add_database_tool],
diff --git a/test/components/agents/test_agent_breakpoints_isolation_async.py b/test/components/agents/test_agent_breakpoints_isolation_async.py
deleted file mode 100644
index a41da346f3..0000000000
--- a/test/components/agents/test_agent_breakpoints_isolation_async.py
+++ /dev/null
@@ -1,187 +0,0 @@
-# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
-#
-# SPDX-License-Identifier: Apache-2.0
-
-import os
-from pathlib import Path
-from unittest.mock import AsyncMock
-
-import pytest
-
-from haystack.components.agents import Agent
-from haystack.core.errors import BreakpointException
-from haystack.core.pipeline.breakpoint import load_pipeline_snapshot
-from haystack.dataclasses import ChatMessage, ToolCall
-from haystack.dataclasses.breakpoints import AgentBreakpoint, Breakpoint, ToolBreakpoint
-from haystack.tools import Tool
-from test.components.agents.test_agent import MockChatGeneratorWithRunAsync, weather_function
-
-AGENT_NAME = "isolated_agent"
-
-
-@pytest.fixture
-def weather_tool():
-    return Tool(
-        name="weather_tool",
-        description="Provides weather information for a given location.",
-        parameters={"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]},
-        function=weather_function,
-    )
-
-
-@pytest.fixture
-def mock_chat_generator():
-    generator = MockChatGeneratorWithRunAsync()
-    mock_run_async = AsyncMock()
-    mock_run_async.return_value = {
-        "replies": [
-            ChatMessage.from_assistant(
-                "I'll help you check the weather.",
-                tool_calls=[{"tool_name": "weather_tool", "tool_args": {"location": "Berlin"}}],
-            )
-        ]
-    }
-
-    async def mock_run_async_with_tools(messages, tools=None, **kwargs):
-        return mock_run_async.return_value
-
-    generator.run_async = mock_run_async_with_tools
-    return generator
-
-
-@pytest.fixture
-def agent(mock_chat_generator, weather_tool):
-    return Agent(
-        chat_generator=mock_chat_generator,
-        tools=[weather_tool],
-        system_prompt="You are a helpful assistant that can use tools to help users.",
-        max_agent_steps=10,  # Increase max steps to allow breakpoints to trigger
-    )
-
-
-@pytest.fixture
-def debug_path(tmp_path):
-    return str(tmp_path / "debug_snapshots")
-
-
-@pytest.fixture
-def mock_agent_with_tool_calls(monkeypatch, weather_tool):
-    monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
-    generator = MockChatGeneratorWithRunAsync()
-    mock_messages = [
-        ChatMessage.from_assistant("First response"),
-        ChatMessage.from_assistant(tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]),
-    ]
-    agent = Agent(chat_generator=generator, tools=[weather_tool], max_agent_steps=10)  # Increase max steps
-    agent.warm_up()
-    agent.chat_generator.run_async = AsyncMock(return_value={"replies": mock_messages})
-    return agent
-
-
-@pytest.mark.asyncio
-async def test_run_async_with_chat_generator_breakpoint(agent):
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    chat_generator_bp = Breakpoint(component_name="chat_generator", visit_count=0)
-    agent_breakpoint = AgentBreakpoint(break_point=chat_generator_bp, agent_name="test")
-    with pytest.raises(BreakpointException) as exc_info:
-        await agent.run_async(messages=messages, break_point=agent_breakpoint, agent_name=AGENT_NAME)
-    assert exc_info.value.component == "chat_generator"
-    assert "messages" in exc_info.value.inputs["chat_generator"]["serialized_data"]
-
-
-@pytest.mark.asyncio
-async def test_run_async_with_tool_invoker_breakpoint(mock_agent_with_tool_calls):
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    tool_bp = ToolBreakpoint(component_name="tool_invoker", visit_count=0, tool_name="weather_tool")
-    agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test")
-    with pytest.raises(BreakpointException) as exc_info:
-        await mock_agent_with_tool_calls.run_async(
-            messages=messages, break_point=agent_breakpoint, agent_name=AGENT_NAME
-        )
-
-    assert exc_info.value.component == "tool_invoker"
-    assert "messages" in exc_info.value.inputs["tool_invoker"]["serialized_data"]
-
-
-@pytest.mark.asyncio
-async def test_resume_from_chat_generator_async(agent, debug_path):
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    chat_generator_bp = Breakpoint(component_name="chat_generator", visit_count=0, snapshot_file_path=debug_path)
-    agent_breakpoint = AgentBreakpoint(break_point=chat_generator_bp, agent_name=AGENT_NAME)
-
-    try:
-        await agent.run_async(messages=messages, break_point=agent_breakpoint, agent_name=AGENT_NAME)
-    except BreakpointException:
-        pass
-
-    snapshot_files = list(Path(debug_path).glob(AGENT_NAME + "_chat_generator_*.json"))
-
-    assert len(snapshot_files) > 0
-    latest_snapshot_file = str(max(snapshot_files, key=os.path.getctime))
-
-    result = await agent.run_async(
-        messages=[ChatMessage.from_user("Continue from where we left off.")],
-        snapshot=load_pipeline_snapshot(latest_snapshot_file).agent_snapshot,
-    )
-
-    assert "messages" in result
-    assert "last_message" in result
-    assert len(result["messages"]) > 0
-
-
-@pytest.mark.asyncio
-async def test_resume_from_tool_invoker_async(mock_agent_with_tool_calls, debug_path):
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    tool_bp = ToolBreakpoint(
-        component_name="tool_invoker", visit_count=0, tool_name="weather_tool", snapshot_file_path=debug_path
-    )
-    agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name=AGENT_NAME)
-
-    try:
-        await mock_agent_with_tool_calls.run_async(
-            messages=messages, break_point=agent_breakpoint, agent_name=AGENT_NAME
-        )
-    except BreakpointException:
-        pass
-
-    snapshot_files = list(Path(debug_path).glob(AGENT_NAME + "_tool_invoker_*.json"))
-
-    assert len(snapshot_files) > 0
-    latest_snapshot_file = str(max(snapshot_files, key=os.path.getctime))
-
-    result = await mock_agent_with_tool_calls.run_async(
-        messages=[ChatMessage.from_user("Continue from where we left off.")],
-        snapshot=load_pipeline_snapshot(latest_snapshot_file).agent_snapshot,
-    )
-
-    assert "messages" in result
-    assert "last_message" in result
-    assert len(result["messages"]) > 0
-
-
-@pytest.mark.asyncio
-async def test_invalid_combination_breakpoint_and_pipeline_snapshot_async(mock_agent_with_tool_calls):
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    tool_bp = ToolBreakpoint(component_name="tool_invoker", visit_count=0, tool_name="weather_tool")
-    agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test")
-    with pytest.raises(ValueError, match="break_point and snapshot cannot be provided at the same time"):
-        await mock_agent_with_tool_calls.run_async(
-            messages=messages, break_point=agent_breakpoint, snapshot={"some": "snapshot"}
-        )
-
-
-@pytest.mark.asyncio
-async def test_breakpoint_with_invalid_component_async(mock_agent_with_tool_calls):
-    invalid_bp = Breakpoint(component_name="invalid_breakpoint", visit_count=0)
-    with pytest.raises(ValueError):
-        AgentBreakpoint(break_point=invalid_bp, agent_name="test")
-
-
-@pytest.mark.asyncio
-async def test_breakpoint_with_invalid_tool_name_async(mock_agent_with_tool_calls):
-    tool_breakpoint = ToolBreakpoint(component_name="tool_invoker", visit_count=0, tool_name="invalid_tool")
-    with pytest.raises(ValueError, match="Tool 'invalid_tool' is not available in the agent's tools"):
-        agent_breakpoint = AgentBreakpoint(break_point=tool_breakpoint, agent_name="test")
-        await mock_agent_with_tool_calls.run_async(
-            messages=[ChatMessage.from_user("What's the weather in Berlin?")], break_point=agent_breakpoint
-        )
diff --git a/test/components/agents/test_agent_breakpoints_isolation_sync.py b/test/components/agents/test_agent_breakpoints_isolation_sync.py
deleted file mode 100644
index 4afb4b92a9..0000000000
--- a/test/components/agents/test_agent_breakpoints_isolation_sync.py
+++ /dev/null
@@ -1,121 +0,0 @@
-# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
-#
-# SPDX-License-Identifier: Apache-2.0
-
-import os
-from pathlib import Path
-
-import pytest
-
-from haystack.core.errors import BreakpointException
-from haystack.core.pipeline.breakpoint import load_pipeline_snapshot
-from haystack.dataclasses import ChatMessage
-from haystack.dataclasses.breakpoints import AgentBreakpoint, Breakpoint, ToolBreakpoint
-from test.components.agents.test_agent_breakpoints_utils import (
-    agent_sync,
-    mock_agent_with_tool_calls_sync,
-    weather_tool,
-)
-
-AGENT_NAME = "isolated_agent"
-
-
-def test_run_with_chat_generator_breakpoint(agent_sync):  # noqa: F811
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    chat_generator_bp = Breakpoint(component_name="chat_generator", visit_count=0)
-    agent_breakpoint = AgentBreakpoint(break_point=chat_generator_bp, agent_name="test_agent")
-    with pytest.raises(BreakpointException) as exc_info:
-        agent_sync.run(messages=messages, break_point=agent_breakpoint, agent_name="test")
-    assert exc_info.value.component == "chat_generator"
-    assert "messages" in exc_info.value.inputs["chat_generator"]["serialized_data"]
-
-
-def test_run_with_tool_invoker_breakpoint(mock_agent_with_tool_calls_sync):  # noqa: F811
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    tool_bp = ToolBreakpoint(component_name="tool_invoker", visit_count=0, tool_name="weather_tool")
-    agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test_agent")
-    with pytest.raises(BreakpointException) as exc_info:
-        mock_agent_with_tool_calls_sync.run(messages=messages, break_point=agent_breakpoint, agent_name="test")
-
-    assert exc_info.value.component == "tool_invoker"
-    assert {"chat_generator", "tool_invoker"} == set(exc_info.value.inputs.keys())
-    assert "serialization_schema" in exc_info.value.inputs["chat_generator"]
-    assert "serialized_data" in exc_info.value.inputs["chat_generator"]
-    assert "messages" in exc_info.value.inputs["chat_generator"]["serialized_data"]
-
-
-def test_resume_from_chat_generator(agent_sync, tmp_path):  # noqa: F811
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    debug_path = str(tmp_path / "debug_snapshots")
-    chat_generator_bp = Breakpoint(component_name="chat_generator", visit_count=0, snapshot_file_path=debug_path)
-    agent_breakpoint = AgentBreakpoint(break_point=chat_generator_bp, agent_name=AGENT_NAME)
-
-    try:
-        agent_sync.run(messages=messages, break_point=agent_breakpoint, agent_name=AGENT_NAME)
-    except BreakpointException:
-        pass
-
-    snapshot_files = list(Path(debug_path).glob(AGENT_NAME + "_chat_generator_*.json"))
-    assert len(snapshot_files) > 0
-    latest_snapshot_file = str(max(snapshot_files, key=os.path.getctime))
-
-    result = agent_sync.run(
-        messages=[ChatMessage.from_user("Continue from where we left off.")],
-        snapshot=load_pipeline_snapshot(latest_snapshot_file).agent_snapshot,
-    )
-
-    assert "messages" in result
-    assert "last_message" in result
-    assert len(result["messages"]) > 0
-
-
-def test_resume_from_tool_invoker(mock_agent_with_tool_calls_sync, tmp_path):  # noqa: F811
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    debug_path = str(tmp_path / "debug_snapshots")
-    tool_bp = ToolBreakpoint(
-        component_name="tool_invoker", visit_count=0, tool_name=None, snapshot_file_path=debug_path
-    )
-    agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name=AGENT_NAME)
-
-    try:
-        mock_agent_with_tool_calls_sync.run(messages=messages, break_point=agent_breakpoint, agent_name=AGENT_NAME)
-    except BreakpointException:
-        pass
-
-    snapshot_files = list(Path(debug_path).glob(AGENT_NAME + "_tool_invoker_*.json"))
-    assert len(snapshot_files) > 0
-    latest_snapshot_file = str(max(snapshot_files, key=os.path.getctime))
-
-    result = mock_agent_with_tool_calls_sync.run(
-        messages=[ChatMessage.from_user("Continue from where we left off.")],
-        snapshot=load_pipeline_snapshot(latest_snapshot_file).agent_snapshot,
-    )
-
-    assert "messages" in result
-    assert "last_message" in result
-    assert len(result["messages"]) > 0
-
-
-def test_invalid_combination_breakpoint_and_pipeline_snapshot(mock_agent_with_tool_calls_sync):  # noqa: F811
-    messages = [ChatMessage.from_user("What's the weather in Berlin?")]
-    tool_bp = ToolBreakpoint(component_name="tool_invoker", visit_count=0, tool_name="weather_tool")
-    agent_breakpoint = AgentBreakpoint(break_point=tool_bp, agent_name="test_agent")
-    with pytest.raises(ValueError, match="break_point and snapshot cannot be provided at the same time"):
-        mock_agent_with_tool_calls_sync.run(
-            messages=messages, break_point=agent_breakpoint, snapshot={"some": "snapshot"}
-        )
-
-
-def test_breakpoint_with_invalid_component(mock_agent_with_tool_calls_sync):  # noqa: F811
-    invalid_bp = Breakpoint(component_name="invalid_breakpoint", visit_count=0)
-    with pytest.raises(ValueError):
-        AgentBreakpoint(break_point=invalid_bp, agent_name="test_agent")
-
-
-def test_breakpoint_with_invalid_tool_name(mock_agent_with_tool_calls_sync):  # noqa: F811
-    tool_breakpoint = ToolBreakpoint(component_name="tool_invoker", visit_count=0, tool_name="invalid_tool")
-    with pytest.raises(ValueError, match="Tool 'invalid_tool' is not available in the agent's tools"):
-        agent_breakpoints = AgentBreakpoint(break_point=tool_breakpoint, agent_name="test_agent")
-        mock_agent_with_tool_calls_sync.run(
-            messages=[ChatMessage.from_user("What's the weather in Berlin?")], break_point=agent_breakpoints
-        )
diff --git a/test/components/agents/test_agent_breakpoints_utils.py b/test/components/agents/test_agent_breakpoints_utils.py
deleted file mode 100644
index 3edd23100f..0000000000
--- a/test/components/agents/test_agent_breakpoints_utils.py
+++ /dev/null
@@ -1,110 +0,0 @@
-# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
-#
-# SPDX-License-Identifier: Apache-2.0
-
-from unittest.mock import AsyncMock, MagicMock
-
-import pytest
-
-from haystack.components.agents import Agent
-from haystack.components.generators.chat import OpenAIChatGenerator
-from haystack.dataclasses import ChatMessage, ToolCall
-from haystack.tools import Tool
-from test.components.agents.test_agent import (
-    MockChatGeneratorWithoutRunAsync,
-    MockChatGeneratorWithRunAsync,
-    weather_function,
-)
-
-
-# Common fixtures
-@pytest.fixture
-def weather_tool():
-    return Tool(
-        name="weather_tool",
-        description="Provides weather information for a given location.",
-        parameters={"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]},
-        function=weather_function,
-    )
-
-
-@pytest.fixture
-def debug_path(tmp_path):
-    return str(tmp_path / "debug_snapshots")
-
-
-@pytest.fixture
-def agent_sync(weather_tool):
-    generator = MockChatGeneratorWithoutRunAsync()
-    mock_run = MagicMock()
-    mock_run.return_value = {
-        "replies": [
-            ChatMessage.from_assistant(
-                "I'll help you check the weather.",
-                tool_calls=[{"tool_name": "weather_tool", "tool_args": {"location": "Berlin"}}],
-            )
-        ]
-    }
-
-    def mock_run_with_tools(messages, tools=None, **kwargs):
-        return mock_run.return_value
-
-    generator.run = mock_run_with_tools
-
-    return Agent(
-        chat_generator=generator,
-        tools=[weather_tool],
-        system_prompt="You are a helpful assistant that can use tools to help users.",
-    )
-
-
-@pytest.fixture
-def mock_agent_with_tool_calls_sync(monkeypatch, weather_tool):
-    monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
-    generator = OpenAIChatGenerator()
-    mock_messages = [
-        ChatMessage.from_assistant("First response"),
-        ChatMessage.from_assistant(tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]),
-    ]
-    agent = Agent(chat_generator=generator, tools=[weather_tool], max_agent_steps=1)
-    agent.warm_up()
-    agent.chat_generator.run = MagicMock(return_value={"replies": mock_messages})
-    return agent
-
-
-@pytest.fixture
-def agent_async(weather_tool):
-    generator = MockChatGeneratorWithRunAsync()
-    mock_run_async = AsyncMock()
-    mock_run_async.return_value = {
-        "replies": [
-            ChatMessage.from_assistant(
-                "I'll help you check the weather.",
-                tool_calls=[{"tool_name": "weather_tool", "tool_args": {"location": "Berlin"}}],
-            )
-        ]
-    }
-
-    async def mock_run_async_with_tools(messages, tools=None, **kwargs):
-        return mock_run_async.return_value
-
-    generator.run_async = mock_run_async_with_tools
-    return Agent(
-        chat_generator=generator,
-        tools=[weather_tool],
-        system_prompt="You are a helpful assistant that can use tools to help users.",
-    )
-
-
-@pytest.fixture
-def mock_agent_with_tool_calls_async(monkeypatch, weather_tool):
-    monkeypatch.setenv("OPENAI_API_KEY", "fake-key")
-    generator = MockChatGeneratorWithRunAsync()
-    mock_messages = [
-        ChatMessage.from_assistant("First response"),
-        ChatMessage.from_assistant(tool_calls=[ToolCall(tool_name="weather_tool", arguments={"location": "Berlin"})]),
-    ]
-    agent = Agent(chat_generator=generator, tools=[weather_tool], max_agent_steps=1)
-    agent.warm_up()
-    agent.chat_generator.run_async = AsyncMock(return_value={"replies": mock_messages})
-    return agent