toolset refactor POC

docs-website/docs/tools/toolset.mdx

@@ -73,21 +73,28 @@ def multiply_numbers(
 
 
 math_toolset.add(multiply_numbers)
+```
+
+### Combining Toolsets
+
+To use multiple Toolsets together, pass them as a list wherever tools are accepted:
 
-# or, you can merge toolsets together
-math_toolset.add(another_toolset)
+```python
+agent = Agent(
+    chat_generator=OpenAIChatGenerator(), tools=[math_toolset, another_toolset]
+)
 ```
 
 ### Run-Scoped Copies and Tool Selection
 
-A `Toolset` is never mutated in place during an [`Agent`](../pipeline-components/agents-1/agent.mdx) run. Each run operates on an isolated, run-scoped copy of the configured `Toolset`, created with the `spawn()` method. This makes concurrent runs that share the same `Toolset` instance safe: per-run state, such as an active tool-name selection or a [`SearchableToolset`](searchabletoolset.mdx)'s discovered tools, cannot leak or collide across runs.
+An [`Agent`](../pipeline-components/agents-1/agent.mdx) run never modifies your configured `Toolset`. A `Toolset` with per-run state, such as a [`SearchableToolset`](searchabletoolset.mdx), is copied for each run through its `spawn()` method, so concurrent runs cannot leak state (like discovered tools) into each other. A plain `Toolset` has no per-run state and is shared as is; just avoid adding or removing tools while runs are in progress.
 
-You can also restrict an `Agent` to a subset of tools at runtime by passing tool names, for example `agent.run(tools=["tool_a", "tool_b"])`. When a `Toolset` is configured, the selection is applied to the live (run-scoped) `Toolset` rather than flattening it into a static list, so dynamic behavior like a `SearchableToolset`'s search and lazy loading keeps working over the selected subset.
+You can also restrict an `Agent` to a subset of tools at runtime by passing tool names, for example `agent.run(tools=["tool_a", "tool_b"])`. The selection applies only to that run, and dynamic behavior like a `SearchableToolset`'s search keeps working over the selected subset.
 
 Two methods support this and can be overridden when subclassing:
 
-- `get_selectable_tools()`: Returns every tool available for name-based selection, ignoring any active selection restriction. Override it if your subclass's iteration does not surface every selectable tool.
-- `spawn()`: Returns an isolated, run-scoped copy of the `Toolset`. Override it if your subclass holds additional run-scoped state.
+- `get_selectable_tools()`: Returns every tool available for name-based selection. Override it if your subclass's iteration does not surface every selectable tool.
+- `spawn()`: Returns the `Toolset` itself, which has no run-scoped state to isolate. Override it to return an isolated, run-scoped copy if your subclass holds run-scoped state.
 
 ## Usage
 

haystack/components/agents/agent.py

@@ -51,7 +51,6 @@
     warm_up_hooks_async,
 )
 from haystack.tools import (
-    Tool,
     Toolset,
     ToolsType,
     _check_duplicate_tool_names,
@@ -478,7 +477,6 @@ def __init__(  # noqa: PLR0913
         self.tool_concurrency_limit = tool_concurrency_limit
         self.tool_streaming_callback_passthrough = tool_streaming_callback_passthrough
         self.hooks = hooks
-        self._tools_warmed_up = False
         self._hooks_warmed_up = False
 
         # --- State schema ---
@@ -569,13 +567,6 @@ def _register_prompt_variables(self) -> None:
             else:
                 component.set_input_type(self, name=var_name, type=Any, default=None)
 
-    def _warm_up_tools(self) -> None:
-        """Warm up the configured tools once."""
-        if not self._tools_warmed_up:
-            if self.tools:
-                warm_up_tools(tools=self.tools)
-            self._tools_warmed_up = True
-
     def _warm_up_hooks(self) -> None:
         """Warm up the configured hooks once."""
         if not self._hooks_warmed_up:
@@ -590,14 +581,14 @@ async def _warm_up_hooks_async(self) -> None:
 
     def warm_up(self) -> None:
         """Warm up the tools, hooks, and the underlying chat generator."""
-        self._warm_up_tools()
+        warm_up_tools(tools=self.tools)
         self._warm_up_hooks()
         if hasattr(self.chat_generator, "warm_up"):
             self.chat_generator.warm_up()
 
     async def warm_up_async(self) -> None:
         """Warm up the tools, hooks, and the underlying chat generator on the serving event loop."""
-        self._warm_up_tools()
+        warm_up_tools(tools=self.tools)
         await self._warm_up_hooks_async()
         if hasattr(self.chat_generator, "warm_up_async"):
             await self.chat_generator.warm_up_async()
@@ -791,22 +782,16 @@ def _select_tools(self, tools: ToolsType | list[str] | None = None) -> ToolsType
             or if any provided tool name is not valid.
         :raises TypeError: If tools is not a list of Tool objects, a Toolset, or a list of tool names (strings).
         """
-        # Toolsets are spawned into per-run copies (see _spawn_tools / _select_tools_by_name) so concurrent runs
+        # Toolsets are spawned per run (see _spawn_tools / _select_tools_by_name) so concurrent runs
         # sharing the same configured Toolset don't corrupt each other's run-scoped state.
         if tools is None:
             return _spawn_tools(tools=self.tools)
 
         if isinstance(tools, list) and all(isinstance(t, str) for t in tools):
             return _select_tools_by_name(self.tools, cast(list[str], tools))
 
-        if isinstance(tools, Toolset):
-            # Per-run tools are not covered by the Agent's own warm_up(), so warm them up here.
-            # warm_up() is expected to be idempotent, so re-warming on every run is cheap.
-            warm_up_tools(tools=tools)
-            return _spawn_tools(tools=tools)
-
-        if isinstance(tools, list):
-            selected = cast(list[Tool | Toolset], tools)  # mypy can't narrow the Union type from isinstance check
+        if isinstance(tools, (Toolset, list)):
+            selected = cast(ToolsType, tools)  # mypy can't narrow the Union type from the isinstance checks
             # Per-run tools are not covered by the Agent's own warm_up(), so warm them up here.
             # warm_up() is expected to be idempotent, so re-warming on every run is cheap.
             warm_up_tools(tools=selected)

haystack/components/agents/utils.py

@@ -90,36 +90,56 @@ def _record_tool_calls(state: State, tool_messages: list[ChatMessage]) -> None:
 # ---------------------------
 
 
+def _spawn_selection_copy(item: Tool | Toolset, selected_tool_names: set[str]) -> Toolset | None:
+    """
+    Return the per-run copy carrying the selection, or None if the item does not provide one.
+
+    A Toolset with run-scoped state (e.g. SearchableToolset) overrides `spawn()` to return a copy that
+    applies `selected_tool_names` itself. A plain Toolset returns itself from `spawn()` (it has nothing
+    to isolate), and a standalone Tool has no `spawn()`: in both cases the caller applies the selection.
+
+    :param item: A configured Tool or Toolset.
+    :param selected_tool_names: The tool names selected for this run.
+    :returns: The selection-carrying per-run copy, or None.
+    """
+    if not isinstance(item, Toolset):
+        return None
+    spawned = item.spawn(selected_tool_names=selected_tool_names)
+    return spawned if spawned is not item else None
+
+
 def _select_tools_by_name(configured_tools: ToolsType, names: list[str]) -> list[Tool | Toolset]:
     """
     Select configured tools by name for a single run.
 
-    Standalone Tools are kept when their name is requested. A Toolset that exposes a requested name is replaced by a
-    per-run `spawn()` (an isolated copy) with the requested names registered as its `_selected_tool_names`, so
-    dynamic toolsets such as SearchableToolset preserve their behavior (search/lazy-loading) over the selected subset
-    without mutating the shared, configured Toolset.
+    Standalone Tools are kept when their name is requested. A Toolset with run-scoped state (one overriding
+    `spawn()`, such as SearchableToolset) is replaced by a per-run copy carrying the requested names, so its
+    dynamic behavior (search/lazy-loading) is preserved without mutating the shared, configured Toolset. Any
+    other Toolset is warmed up and reduced to the matching Tools.
 
     :param configured_tools: The tools configured on the Agent.
     :param names: The requested tool names.
-    :returns: The selected standalone Tools and/or spawned, selection-scoped Toolsets.
+    :returns: The selected Tools and/or selection-scoped Toolset copies.
     :raises ValueError: If no tools were configured, or if any requested name is not a valid tool name.
     """
-    if not configured_tools:
+    if configured_tools is None:
         raise ValueError("No tools were configured for the Agent at initialization.")
 
     requested_names = set(names)
     items: list[Tool | Toolset] = (
         [configured_tools] if isinstance(configured_tools, Toolset) else list(configured_tools)
     )
 
-    # Resolve selectable names per item. For Toolsets we use get_selectable_tools() so dynamic toolsets
-    # (e.g. SearchableToolset) offer their full catalog by name, not just the tools exposed by iteration.
-    selectable_per_item: list[tuple[Tool | Toolset, set[str]]] = []
-    valid_tool_names: set[str] = set()
+    # Resolve the tools each item offers for selection
+    selectable_per_item: list[tuple[Tool | Toolset, list[Tool]]] = []
     for item in items:
-        item_names = {tool.name for tool in item.get_selectable_tools()} if isinstance(item, Toolset) else {item.name}
-        selectable_per_item.append((item, item_names))
-        valid_tool_names |= item_names
+        selectable = item.get_selectable_tools() if isinstance(item, Toolset) else [item]
+        selectable_per_item.append((item, selectable))
+
+    valid_tool_names = {tool.name for _, selectable in selectable_per_item for tool in selectable}
+    # A dynamic Toolset may look empty before its catalog is resolved, so emptiness is checked here.
+    if not valid_tool_names:
+        raise ValueError("No tools were configured for the Agent at initialization.")
 
     invalid_tool_names = requested_names - valid_tool_names
     if invalid_tool_names:
@@ -128,27 +148,28 @@ def _select_tools_by_name(configured_tools: ToolsType, names: list[str]) -> list
         )
 
     selected: list[Tool | Toolset] = []
-    for item, item_names in selectable_per_item:
-        matched = requested_names & item_names
+    for item, selectable in selectable_per_item:
+        matched = requested_names & {tool.name for tool in selectable}
         if not matched:
             continue
-        if isinstance(item, Toolset):
-            # Apply the selection to a per-run copy so the shared, configured Toolset is never mutated.
-            spawned = item.spawn()
-            spawned._selected_tool_names = matched
-            selected.append(spawned)
+        run_copy = _spawn_selection_copy(item, matched)
+        if run_copy is not None:
+            selected.append(run_copy)
         else:
-            selected.append(item)
+            # Select from `selectable`, the list the names were validated against: iterating a dynamic
+            # Toolset could silently miss tools.
+            selected.extend(tool for tool in selectable if tool.name in matched)
     return selected
 
 
 def _spawn_tools(tools: ToolsType) -> ToolsType:
     """
-    Return per-run copies of `tools`, replacing each Toolset with an isolated `spawn()` (Tools are passed through).
+    Return per-run copies of `tools`, replacing each Toolset with its `spawn()` (Tools are passed through).
 
     This isolates run-scoped Toolset state (e.g. a SearchableToolset's discovered tools and any active name
     selection) so that concurrent runs sharing the same configured Toolset — such as parallel sub-agent tool calls
-    or concurrent requests against one Agent — don't corrupt each other.
+    or concurrent requests against one Agent — don't corrupt each other. A plain Toolset has no run-scoped state
+    and its `spawn()` returns itself unchanged.
     """
     if isinstance(tools, Toolset):
         return tools.spawn()

haystack/tools/searchable_toolset.py

@@ -129,16 +129,15 @@ def __init__(
         self._bootstrap_tool: Tool | None = None
         self._document_store: InMemoryDocumentStore | None = None
         self._passthrough: bool | None = None
-        self._is_warmed_up = False
+
+        # Optional per-run name filter, set on the copies returned by spawn(). When set, iteration only
+        # yields tools whose name is in this set, and search is scoped to it. None means no filtering.
+        self._selected_tool_names: set[str] | None = None
 
         # Initialize parent with empty tools list - we manage tools dynamically
         super().__init__(tools=[])
 
-    def __add__(self, other: Tool | Toolset | list[Tool]) -> "Toolset":
-        """Concatenation is not supported for SearchableToolset."""
-        raise NotImplementedError("SearchableToolset does not support concatenation.")
-
-    def add(self, tool: Tool | Toolset) -> None:
+    def add(self, tool: Tool) -> None:
         """Adding new tools after initialization is not supported for SearchableToolset."""
         raise NotImplementedError("SearchableToolset does not support adding new tools after initialization.")
 
@@ -153,7 +152,7 @@ def warm_up(self) -> None:
 
         :raises ValueError: If the flattened catalog contains tools with duplicate names.
         """
-        if self._is_warmed_up:
+        if self._passthrough is not None:
             return
 
         # Warm up the catalog first (triggers lazy connections like MCPToolset), then flatten — lazy toolsets will
@@ -176,8 +175,6 @@ def warm_up(self) -> None:
             self._document_store.write_documents(documents, policy=DuplicatePolicy.OVERWRITE)
             self._bootstrap_tool = self._create_search_tool()
 
-        self._is_warmed_up = True
-
     def get_selectable_tools(self) -> list[Tool]:
         """
         Return the full catalog of tools that can be selected by name.
@@ -187,8 +184,7 @@ def get_selectable_tools(self) -> list[Tool]:
 
         :returns: The flattened catalog of tools.
         """
-        if not self._is_warmed_up:
-            self.warm_up()
+        self.warm_up()
         return list(self._catalog)
 
     def clear(self) -> None:
@@ -200,21 +196,23 @@ def clear(self) -> None:
         """
         self._discovered_tools.clear()
 
-    def spawn(self) -> "SearchableToolset":
+    def spawn(self, selected_tool_names: set[str] | None = None) -> "SearchableToolset":
         """
-        Return an isolated copy for a single run.
+        Return an isolated copy for a single run, carrying the given name selection.
 
         The copy shares the read-only catalog and BM25 index but gets fresh discovered tools and name selection,
-        plus a bootstrap search tool bound to the copy. This way concurrent runs sharing the same configured
-        SearchableToolset don't share discovered tools or collide on the active selection.
+        plus a bootstrap search tool bound to the copy; the selection scopes both iteration and search. This way
+        concurrent runs sharing the same configured SearchableToolset don't share discovered tools or collide on
+        the active selection.
 
+        :param selected_tool_names: Optional catalog tool names this run is restricted to. None means no
+            restriction.
         :returns: A run-scoped copy of this SearchableToolset.
         """
-        if not self._is_warmed_up:
-            self.warm_up()
+        self.warm_up()
         new = copy.copy(self)
         new._discovered_tools = {}
-        new._selected_tool_names = None
+        new._selected_tool_names = set(selected_tool_names) if selected_tool_names is not None else None
         # Rebuild the bootstrap tool so its closure is bound to the copy's discovered tools / selection
         # rather than the original's. The document store and catalog are read-only and stay shared.
         if not self._passthrough:
@@ -307,24 +305,17 @@ def __iter__(self) -> Iterator[Tool]:
         set, but the bootstrap search tool is always exposed so search keeps working over the selected subset.
         Automatically calls warm_up() if needed to ensure the bootstrap tool is available.
         """
-        # Unlike base Toolset/MCPToolset, which expose a placeholder tool before warm_up, this toolset materializes
-        # everything (flattened catalog, bootstrap tool, passthrough decision) in warm_up.
+        # This toolset materializes everything (flattened catalog, bootstrap tool, passthrough decision) in warm_up.
         # Without warming here, iterating before warm_up would yield nothing, so we warm up to make the toolset usable
         # at all.
-        if not self._is_warmed_up:
-            self.warm_up()
+        self.warm_up()
         if self._passthrough:
             yield from (tool for tool in self._catalog if self._is_selected(tool.name))
         else:
             if self._bootstrap_tool is not None:
                 yield self._bootstrap_tool
             yield from (tool for tool in self._discovered_tools.values() if self._is_selected(tool.name))
 
-    def __len__(self) -> int:
-        """Return the number of currently available tools."""
-        # the number of tools is computed by invoking __iter__ on the toolset
-        return sum(1 for _ in self)
-
     def __contains__(self, item: str | Tool) -> bool:
         """
         Check if a tool is available by Tool instance or tool name string.
@@ -338,16 +329,6 @@ def __contains__(self, item: str | Tool) -> bool:
             return any(tool == item for tool in self)
         raise TypeError(f"Invalid item type: {type(item)}. Must be Tool or str.")
 
-    def __getitem__(self, index: int) -> Tool:
-        """
-        Get a tool by index.
-
-        :param index: Index of the tool to retrieve.
-        :returns: The tool at the given index.
-        :raises IndexError: If the index is out of range.
-        """
-        return list(self)[index]
-
     def to_dict(self) -> dict[str, Any]:
         """
         Serialize the toolset to a dictionary.

haystack/tools/skills/skill_toolset.py

@@ -103,20 +103,13 @@ def warm_up(self) -> None:
         self._load_skill_tool.description = self._load_skill_description()
         self._is_warmed_up = True
 
-    def add(self, tool: Tool | Toolset) -> None:
+    def add(self, tool: Tool) -> None:
         """Adding tools is not supported: a SkillToolset's tools are fixed and defined by its store."""
         raise NotImplementedError(
             "SkillToolset does not support adding tools. To combine it with other tools, pass it to the Agent "
             "alongside them, e.g. tools=[skill_toolset, other_tool]."
         )
 
-    def __add__(self, other: Tool | Toolset | list[Tool]) -> "Toolset":
-        """Concatenation is not supported for SearchableToolset."""
-        raise NotImplementedError(
-            "SkillToolset does not support concatenation. To combine it with other tools, pass it to the Agent "
-            "alongside them, e.g. tools=[skill_toolset, other_tool]."
-        )
-
     def _load_skill_description(self) -> str:
         """
         Build the `load_skill` tool description, including the catalog of discovered skills.