From 024e8fc0649da654fe2c60826e80c2d9a2d2df70 Mon Sep 17 00:00:00 2001 From: Mackenzie Zastrow Date: Thu, 28 Aug 2025 10:25:59 -0400 Subject: [PATCH 1/3] Add invocation_state to ToolContext Addresses issue #579, #750 --- src/strands/tools/decorator.py | 4 +++- src/strands/types/tools.py | 3 +++ tests/strands/tools/test_decorator.py | 12 ++++++++++-- 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/src/strands/tools/decorator.py b/src/strands/tools/decorator.py index 75abac9ed..7ad3a92bd 100644 --- a/src/strands/tools/decorator.py +++ b/src/strands/tools/decorator.py @@ -268,7 +268,9 @@ def inject_special_parameters( invocation_state: Context for the tool invocation, including agent state. """ if self._context_param and self._context_param in self.signature.parameters: - tool_context = ToolContext(tool_use=tool_use, agent=invocation_state["agent"]) + tool_context = ToolContext( + tool_use=tool_use, agent=invocation_state["agent"], invocation_state=invocation_state + ) validated_input[self._context_param] = tool_context # Inject agent if requested (backward compatibility) diff --git a/src/strands/types/tools.py b/src/strands/types/tools.py index bb7c874f6..36e06d980 100644 --- a/src/strands/types/tools.py +++ b/src/strands/types/tools.py @@ -132,6 +132,8 @@ class ToolContext: tool_use: The complete ToolUse object containing tool invocation details. agent: The Agent instance executing this tool, providing access to conversation history, model configuration, and other agent state. + invocation_state: Keyword arguments passed to agent invocation methods (agent(), agent.invoke_async(), etc.). + Provides access to invocation-specific context and parameters. Note: This class is intended to be instantiated by the SDK. Direct construction by users @@ -140,6 +142,7 @@ class ToolContext: tool_use: ToolUse agent: "Agent" + invocation_state: dict[str, Any] ToolChoice = Union[ diff --git a/tests/strands/tools/test_decorator.py b/tests/strands/tools/test_decorator.py index e490c7bb0..6ff5cf757 100644 --- a/tests/strands/tools/test_decorator.py +++ b/tests/strands/tools/test_decorator.py @@ -1039,7 +1039,7 @@ def complex_schema_tool(union_param: Union[List[int], Dict[str, Any], str, None] assert "NoneType: None" in result["content"][0]["text"] -async def _run_context_injection_test(context_tool: AgentTool): +async def _run_context_injection_test(context_tool: AgentTool, additional_context=None): """Common test logic for context injection tests.""" tool: AgentTool = context_tool generator = tool.stream( @@ -1052,6 +1052,7 @@ async def _run_context_injection_test(context_tool: AgentTool): }, invocation_state={ "agent": Agent(name="test_agent"), + **(additional_context or {}), }, ) tool_results = [value async for value in generator] @@ -1081,6 +1082,8 @@ def context_tool(message: str, agent: Agent, tool_context: ToolContext) -> dict: tool_name = tool_context.tool_use["name"] agent_from_tool_context = tool_context.agent + assert tool_context.invocation_state["new_value"] == 13 + return { "status": "success", "content": [ @@ -1090,7 +1093,12 @@ def context_tool(message: str, agent: Agent, tool_context: ToolContext) -> dict: ], } - await _run_context_injection_test(context_tool) + await _run_context_injection_test( + context_tool, + { + "new_value": 13, + }, + ) @pytest.mark.asyncio From dd886704cf46c299cded6447401aece697834e9e Mon Sep 17 00:00:00 2001 From: Mackenzie Zastrow Date: Thu, 28 Aug 2025 13:18:54 -0400 Subject: [PATCH 2/3] Update doc strings for invocation_state --- src/strands/tools/decorator.py | 6 ++++-- src/strands/types/tools.py | 7 ++++--- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/strands/tools/decorator.py b/src/strands/tools/decorator.py index 7ad3a92bd..2ce6d946f 100644 --- a/src/strands/tools/decorator.py +++ b/src/strands/tools/decorator.py @@ -265,7 +265,8 @@ def inject_special_parameters( Args: validated_input: The validated input parameters (modified in place). tool_use: The tool use request containing tool invocation details. - invocation_state: Context for the tool invocation, including agent state. + invocation_state: Caller-provided kwargs that were passed to the agent when it was invoked (agent(), + agent.invoke_async(), etc.). """ if self._context_param and self._context_param in self.signature.parameters: tool_context = ToolContext( @@ -435,7 +436,8 @@ async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kw Args: tool_use: The tool use specification from the Agent. - invocation_state: Context for the tool invocation, including agent state. + invocation_state: Caller-provided kwargs that were passed to the agent when it was invoked (agent(), + agent.invoke_async(), etc.). **kwargs: Additional keyword arguments for future extensibility. Yields: diff --git a/src/strands/types/tools.py b/src/strands/types/tools.py index 36e06d980..1e0f4b841 100644 --- a/src/strands/types/tools.py +++ b/src/strands/types/tools.py @@ -132,8 +132,8 @@ class ToolContext: tool_use: The complete ToolUse object containing tool invocation details. agent: The Agent instance executing this tool, providing access to conversation history, model configuration, and other agent state. - invocation_state: Keyword arguments passed to agent invocation methods (agent(), agent.invoke_async(), etc.). - Provides access to invocation-specific context and parameters. + invocation_state: Caller-provided kwargs that were passed to the agent when it was invoked (agent(), + agent.invoke_async(), etc.). Note: This class is intended to be instantiated by the SDK. Direct construction by users @@ -249,7 +249,8 @@ def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kwargs: Args: tool_use: The tool use request containing tool ID and parameters. - invocation_state: Context for the tool invocation, including agent state. + invocation_state: Caller-provided kwargs that were passed to the agent when it was invoked (agent(), + agent.invoke_async(), etc.). **kwargs: Additional keyword arguments for future extensibility. Yields: From 873b82b6ae4d8d3bb4d64b9237beea14efd36867 Mon Sep 17 00:00:00 2001 From: Mackenzie Zastrow Date: Thu, 28 Aug 2025 15:19:00 -0400 Subject: [PATCH 3/3] Switch test to passing a reference --- tests/strands/tools/test_decorator.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/strands/tools/test_decorator.py b/tests/strands/tools/test_decorator.py index 6ff5cf757..02e7eb445 100644 --- a/tests/strands/tools/test_decorator.py +++ b/tests/strands/tools/test_decorator.py @@ -2,6 +2,7 @@ Tests for the function-based tool decorator pattern. """ +from asyncio import Queue from typing import Any, Dict, Optional, Union from unittest.mock import MagicMock @@ -1075,6 +1076,8 @@ async def _run_context_injection_test(context_tool: AgentTool, additional_contex async def test_tool_context_injection_default(): """Test that ToolContext is properly injected with default parameter name (tool_context).""" + value_to_pass = Queue() # a complex value that is not serializable + @strands.tool(context=True) def context_tool(message: str, agent: Agent, tool_context: ToolContext) -> dict: """Tool that uses ToolContext to access tool_use_id.""" @@ -1082,7 +1085,7 @@ def context_tool(message: str, agent: Agent, tool_context: ToolContext) -> dict: tool_name = tool_context.tool_use["name"] agent_from_tool_context = tool_context.agent - assert tool_context.invocation_state["new_value"] == 13 + assert tool_context.invocation_state["test_reference"] is value_to_pass return { "status": "success", @@ -1096,7 +1099,7 @@ def context_tool(message: str, agent: Agent, tool_context: ToolContext) -> dict: await _run_context_injection_test( context_tool, { - "new_value": 13, + "test_reference": value_to_pass, }, )