fix: unwrap Annotated in _unwrap_optional for State parameter detection (#13169)

Co-authored-by: anakin87 <[email protected]>

haystack/components/agents/tool_calling.py

@@ -17,7 +17,7 @@ from haystack.dataclasses import ChatMessage, ToolCall
 from haystack.dataclasses.streaming_chunk import StreamingCallbackT, StreamingChunk, _invoke_streaming_callback
 from haystack.tools import ComponentTool, Tool, ToolsType, _check_duplicate_tool_names, flatten_tools_or_toolsets
 from haystack.tools.errors import ToolInvocationError
-from haystack.tools.parameters_schema_utils import _unwrap_optional
+from haystack.tools.parameters_schema_utils import _is_state_annotation
 from haystack.tracing.utils import _serializable_value
 
 logger = logging.getLogger(__name__)
@@ -307,7 +307,7 @@ def _inject_state_args(tool: Tool, llm_args: dict[str, Any], state: State) -> di
 
     # We also inject the full State object for any parameter annotated as State
     for param_name, param_type in func_params.items():
-        if _unwrap_optional(param_type) is State:
+        if _is_state_annotation(param_type):
             final_args[param_name] = state
 
     return final_args
@@ -424,7 +424,7 @@ def _state_io_for_call(tool: Tool, llm_args: dict[str, Any]) -> tuple[_StateKeys
     """
     func_params = _get_func_params(tool)
     # Check if State is in func_params
-    if any(_unwrap_optional(param_type) is State for param_type in func_params.values()):
+    if any(_is_state_annotation(param_type) for param_type in func_params.values()):
         return _ALL_STATE_KEYS, _ALL_STATE_KEYS
 
     # Calculate reads

haystack/tools/component_tool.py

@@ -8,7 +8,6 @@ from typing import Any, get_args, get_origin
 from pydantic import Field, TypeAdapter, create_model
 
 from haystack import logging
-from haystack.components.agents.state.state import State
 from haystack.core.component import Component
 from haystack.core.serialization import component_to_dict, generate_qualified_class_name
 from haystack.tools import Tool
@@ -17,6 +16,7 @@ from haystack.tools.from_function import _remove_title_from_schema
 from haystack.tools.parameters_schema_utils import (
     _contains_callable_type,
     _get_component_param_descriptions,
+    _is_state_annotation,
     _resolve_type,
     _unwrap_optional,
 )
@@ -361,7 +361,7 @@ class ComponentTool(Tool):
                 continue
 
             # Skip State-typed parameters - Agent tool execution injects them at runtime
-            if _unwrap_optional(input_type) is State:
+            if _is_state_annotation(input_type):
                 continue
 
             description = param_descriptions.get(input_name, f"Input '{input_name}' for the component.")

haystack/tools/from_function.py

@@ -8,11 +8,10 @@ from typing import Any, overload
 
 from pydantic import create_model
 
-from haystack.components.agents.state.state import State
 from haystack.core.type_utils import _resolve_parameter_types
 
 from .errors import SchemaGenerationError
-from .parameters_schema_utils import _contains_callable_type, _unwrap_optional
+from .parameters_schema_utils import _contains_callable_type, _is_state_annotation
 from .tool import Tool
 
 
@@ -149,7 +148,7 @@ def create_tool_from_function(
             continue
 
         # Skip State-typed parameters (including Optional[State]) - Agent tool execution injects them at runtime
-        if _unwrap_optional(annotation) is State:
+        if _is_state_annotation(annotation):
             continue
 
         if annotation is param.empty:

haystack/tools/parameters_schema_utils.py

@@ -9,12 +9,13 @@ from collections.abc import Callable as ABCCallable
 from dataclasses import MISSING, fields, is_dataclass
 from inspect import getdoc
 from types import NoneType
-from typing import Any, Union, get_args, get_origin
+from typing import Annotated, Any, Union, get_args, get_origin
 
 from docstring_parser import parse
 from pydantic import BaseModel, Field, create_model
 
 from haystack import logging
+from haystack.components.agents.state.state import State
 from haystack.dataclasses import ChatMessage
 from haystack.utils.type_serialization import _is_union_type
 
@@ -36,6 +37,18 @@ def _unwrap_optional(type_hint: Any) -> Any:
     return type_hint
 
 
+def _is_state_annotation(annotation: Any) -> bool:
+    """Recognize State through Annotated and optional wrappers."""
+    if get_origin(annotation) is Annotated:
+        return _is_state_annotation(get_args(annotation)[0])
+
+    unwrapped = _unwrap_optional(annotation)
+    if unwrapped is not annotation:
+        return _is_state_annotation(unwrapped)
+
+    return annotation is State
+
+
 def _contains_callable_type(type_hint: Any) -> bool:
     """
     Check if a type hint contains a Callable type, including within Union types.

releasenotes/notes/fix-annotated-state-tool-parameters-13166.yaml

@@ -0,0 +1,7 @@
+---
+fixes:
+  - |
+    Fixed a ``SchemaGenerationError`` when using ``@tool``, ``create_tool_from_function``, or ``ComponentTool``
+    with a parameter typed as ``Annotated[State, ...]``, including optional variants.
+    These parameters are now excluded from the schema sent to the language model and receive
+    the Agent's live ``State`` object when the tool runs, just like parameters typed as ``State``.

test/components/agents/test_tool_calling.py

@@ -7,7 +7,7 @@ import datetime
 import json
 import time
 from collections.abc import Callable
-from typing import Any
+from typing import Annotated, Any
 from unittest.mock import patch
 
 import pytest
@@ -582,13 +582,15 @@ class TestRunTool:
         )
         assert len(tool_messages_2) == 3
 
-    def test_run_injects_state_object_into_tool(self):
+    @pytest.mark.parametrize("state_annotation", [State, Annotated[State, "Live agent state"]])
+    def test_run_injects_state_object_into_tool(self, state_annotation: Any) -> None:
         received_state = {}
 
         def function_with_state(city: str, state: State) -> str:
             received_state["state"] = state
             return f"Weather in {city}: sunny"
 
+        function_with_state.__annotations__["state"] = state_annotation
         state_tool = Tool(
             name="state_tool",
             description="A tool that receives the live State object.",
@@ -606,13 +608,15 @@ class TestRunTool:
         assert received_state["state"] is state
 
     @pytest.mark.asyncio
-    async def test_run_async_injects_state_object_into_tool(self):
+    @pytest.mark.parametrize("state_annotation", [State, Annotated[State, "Live agent state"]])
+    async def test_run_async_injects_state_object_into_tool(self, state_annotation: Any) -> None:
         received_state = {}
 
         def function_with_state(city: str, state: State) -> str:
             received_state["state"] = state
             return f"Weather in {city}: sunny"
 
+        function_with_state.__annotations__["state"] = state_annotation
         state_tool = Tool(
             name="state_tool",
             description="A tool that receives the live State object.",
@@ -996,12 +1000,13 @@ def _writer_tool(state_key: str = "documents", value: str = "written") -> Tool:
     )
 
 
-def _state_param_tool():
+def _state_param_tool(state_annotation: Any = State) -> Tool:
     """A tool that receives the live State object, so it may read/write any key."""
 
     def fn(state: State) -> str:
         return "done"
 
+    fn.__annotations__["state"] = state_annotation
     return Tool(
         name="state_tool",
         description="Receives the live State.",
@@ -1053,7 +1058,8 @@ class TestScheduleToolCalls:
         tools = [_make_retrieval_tool(), _make_retrieval_tool()]
         assert _schedule_tool_calls(calls, tools) == [[0], [1]]
 
-    def test_state_param_tool_is_a_barrier(self):
+    @pytest.mark.parametrize("state_annotation", [State, Annotated[State, "Live agent state"]])
+    def test_state_param_tool_is_a_barrier(self, state_annotation: Any) -> None:
         # A State-typed tool reads/writes every key, so it serializes everything around it even when the other
         # calls touch disjoint keys ("a" and "b"): writer("b") must run before the state tool (which reads "b"),
         # and reader("a") must run after it (the state tool may write "a").
@@ -1062,7 +1068,7 @@ class TestScheduleToolCalls:
             ToolCall(tool_name="state_tool", arguments={}),
             ToolCall(tool_name="writer_tool", arguments={}),
         ]
-        tools = [_reader_tool("a"), _state_param_tool(), _writer_tool("b")]
+        tools = [_reader_tool("a"), _state_param_tool(state_annotation=state_annotation), _writer_tool("b")]
         assert _schedule_tool_calls(calls, tools) == [[2], [1], [0]]
 
     def test_llm_supplied_arg_is_not_a_state_read(self):

test/tools/test_component_tool.py

@@ -5,7 +5,7 @@
 import json
 import os
 from dataclasses import dataclass
-from typing import Any
+from typing import Annotated, Any
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
@@ -504,7 +504,8 @@ class TestComponentTool:
         assert "streaming_callback" not in param_names
         assert "messages" in param_names
 
-    def test_from_component_with_state_param_excluded_from_schema(self):
+    @pytest.mark.parametrize("state_annotation", [State, Annotated[State, "Live agent state"]])
+    def test_from_component_with_state_param_excluded_from_schema(self, state_annotation: Any) -> None:
         @component
         class ComponentWithState:
             """A component that takes State as a direct input."""
@@ -513,6 +514,8 @@ class TestComponentTool:
             def run(self, query: str, state: State) -> dict:
                 return {"result": query}
 
+            run.__annotations__["state"] = state_annotation
+
         tool = ComponentTool(component=ComponentWithState(), name="state_comp", description="test")
 
         param_names = list(tool.parameters.get("properties", {}).keys())

test/tools/test_from_function.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections.abc import Callable
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
 import jsonschema
 import pytest
@@ -157,11 +157,13 @@ def test_from_function_with_callable_params_skipped():
     assert "query" in param_names
 
 
-def test_from_function_state_param_excluded_from_schema():
+@pytest.mark.parametrize("state_annotation", [State, Annotated[State, "Live agent state"]])
+def test_from_function_state_param_excluded_from_schema(state_annotation: Any) -> None:
     def function_with_state(city: str, state: State) -> str:
         """Get weather for a city, with access to agent state."""
         return f"Weather in {city}: sunny"
 
+    function_with_state.__annotations__["state"] = state_annotation
     tool = create_tool_from_function(function=function_with_state)
 
     assert tool.name == "function_with_state"

test/tools/test_parameters_schema_utils.py

@@ -3,14 +3,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-from typing import Union
+from typing import Annotated, Any, Optional, Union
 
 import pytest
 from pydantic import Field, create_model
 
+from haystack.components.agents import State
 from haystack.dataclasses import ByteStream, ChatMessage, Document, TextContent, ToolCall, ToolCallResult
 from haystack.tools.from_function import _remove_title_from_schema
-from haystack.tools.parameters_schema_utils import _resolve_type
+from haystack.tools.parameters_schema_utils import _is_state_annotation, _resolve_type
 
 BYTE_STREAM_SCHEMA = {
     "type": "object",
@@ -397,3 +398,19 @@ def test_resolve_type_pep_604():
 
     resolved = _resolve_type(dict[str, int] | list[str])
     assert resolved == Union[dict[str, int], list[str]]
+
+
+@pytest.mark.parametrize(
+    "annotation, expected",
+    [
+        (State, True),
+        (Annotated[State, "state"], True),
+        (Annotated[State | None, "state"], True),
+        (Optional[Annotated[State, "state"]], True),  # noqa: UP045 - exercise typing.Optional compatibility
+        (Annotated[str, "description"], False),
+        (list[State], False),
+        (State | str, False),
+    ],
+)
+def test_is_state_annotation(annotation: Any, expected: bool) -> None:
+    assert _is_state_annotation(annotation) is expected