fix(tool/decorator): simplify validation logic

Author: Ratish1

src/strands/tools/decorator.py

@@ -54,8 +54,6 @@ def my_tool(param1: str, param2: int = 42) -> dict:
     TypeVar,
     Union,
     cast,
-    get_args,
-    get_origin,
     get_type_hints,
     overload,
 )
@@ -101,39 +99,12 @@ def __init__(self, func: Callable[..., Any], context_param: str | None = None) -
         self.type_hints = get_type_hints(func)
         self._context_param = context_param
 
+        self._validate_signature()
+
         # Parse the docstring with docstring_parser
         doc_str = inspect.getdoc(func) or ""
         self.doc = docstring_parser.parse(doc_str)
 
-        def _contains_tool_context(tp: Any) -> bool:
-            """Return True if the annotation `tp` (possibly Union/Optional) includes ToolContext."""
-            if tp is None:
-                return False
-            origin = get_origin(tp)
-            if origin is Union:
-                return any(_contains_tool_context(a) for a in get_args(tp))
-            # Handle direct ToolContext type
-            return tp is ToolContext
-
-        for param in self.signature.parameters.values():
-            # Prefer resolved type hints (handles forward refs); fall back to annotation
-            ann = self.type_hints.get(param.name, param.annotation)
-            if ann is inspect._empty:
-                continue
-
-            if _contains_tool_context(ann):
-                # If decorator didn't opt-in to context injection, complain
-                if self._context_param is None:
-                    raise TypeError(
-                        f"Parameter '{param.name}' is of type 'ToolContext' but '@tool(context=True)' is missing."
-                    )
-                # If decorator specified a different param name, complain
-                if param.name != self._context_param:
-                    raise TypeError(
-                        f"Parameter '{param.name}' is of type 'ToolContext' but has the wrong name. "
-                        f"It should be named '{self._context_param}'."
-                    )
-
         # Get parameter descriptions from parsed docstring
         self.param_descriptions = {
             param.arg_name: param.description or f"Parameter {param.arg_name}" for param in self.doc.params
@@ -142,6 +113,21 @@ def _contains_tool_context(tp: Any) -> bool:
         # Create a Pydantic model for validation
         self.input_model = self._create_input_model()
 
+    def _validate_signature(self) -> None:
+        """Verify that ToolContext is used correctly in the function signature."""
+        # Find and validate the ToolContext parameter
+        for param in self.signature.parameters.values():
+            if param.annotation is ToolContext:
+                if self._context_param is None:
+                    raise ValueError("@tool(context=True) must be set if passing in ToolContext param")
+
+                if param.name != self._context_param:
+                    raise ValueError(
+                        f"param_name=<{param.name}> | ToolContext param must be named '{self._context_param}'"
+                    )
+                # Found the parameter, no need to check further
+                break
+
     def _create_input_model(self) -> Type[BaseModel]:
         """Create a Pydantic model from function signature for input validation.
 

tests/strands/tools/test_decorator.py

@@ -1366,41 +1366,35 @@ async def async_generator() -> AsyncGenerator:
 
 
 def test_tool_with_mismatched_tool_context_param_name_raises_error():
-    """Verify that a TypeError is raised for a mismatched tool_context parameter name."""
-    with pytest.raises(TypeError) as excinfo:
+    """Verify that a ValueError is raised for a mismatched tool_context parameter name."""
+    with pytest.raises(ValueError) as excinfo:
 
         @strands.tool(context=True)
         def my_tool(context: ToolContext):
             pass
 
-    assert (
-        "Parameter 'context' is of type 'ToolContext' but has the wrong name. It should be named 'tool_context'."
-        in str(excinfo.value)
-    )
+    assert "ToolContext param must be named 'tool_context'" in str(excinfo.value)
+    assert "param_name=<context>" in str(excinfo.value)
 
 
 def test_tool_with_tool_context_but_no_context_flag_raises_error():
-    """Verify that a TypeError is raised if ToolContext is used without context=True."""
-    with pytest.raises(TypeError) as excinfo:
+    """Verify that a ValueError is raised if ToolContext is used without context=True."""
+    with pytest.raises(ValueError) as excinfo:
 
         @strands.tool
         def my_tool(tool_context: ToolContext):
             pass
 
-    assert "Parameter 'tool_context' is of type 'ToolContext' but '@tool(context=True)' is missing." in str(
-        excinfo.value
-    )
+    assert "@tool(context=True) must be set" in str(excinfo.value)
 
 
 def test_tool_with_tool_context_named_custom_context_raises_error_if_mismatched():
-    """Verify that a TypeError is raised when context param name doesn't match the decorator value."""
-    with pytest.raises(TypeError) as excinfo:
+    """Verify that a ValueError is raised when context param name doesn't match the decorator value."""
+    with pytest.raises(ValueError) as excinfo:
 
         @strands.tool(context="my_context")
         def my_tool(tool_context: ToolContext):
             pass
 
-    assert (
-        "Parameter 'tool_context' is of type 'ToolContext' but has the wrong name. It should be named 'my_context'."
-        in str(excinfo.value)
-    )
+    assert "ToolContext param must be named 'my_context'" in str(excinfo.value)
+    assert "param_name=<tool_context>" in str(excinfo.value)