feat: add InvocationState TypedDict

Author: dbschmigelski

src/strands/agent/agent.py

@@ -38,6 +38,7 @@
 from ..tools.watcher import ToolWatcher
 from ..types.content import ContentBlock, Message, Messages
 from ..types.exceptions import ContextWindowOverflowException
+from ..types.invocation import InvocationState
 from ..types.tools import ToolResult, ToolUse
 from ..types.traces import AttributeValue
 from .agent_result import AgentResult
@@ -523,7 +524,7 @@ async def stream_async(self, prompt: Union[str, list[ContentBlock]], **kwargs: A
                 raise
 
     async def _run_loop(
-        self, message: Message, invocation_state: dict[str, Any]
+        self, message: Message, invocation_state: InvocationState
     ) -> AsyncGenerator[dict[str, Any], None]:
         """Execute the agent's event loop with the given message and parameters.
 
@@ -563,7 +564,7 @@ async def _run_loop(
             self.conversation_manager.apply_management(self)
             self.hooks.invoke_callbacks(AfterInvocationEvent(agent=self))
 
-    async def _execute_event_loop_cycle(self, invocation_state: dict[str, Any]) -> AsyncGenerator[dict[str, Any], None]:
+    async def _execute_event_loop_cycle(self, invocation_state: InvocationState) -> AsyncGenerator[dict[str, Any], None]:
         """Execute the event loop cycle with retry logic for context window limits.
 
         This internal method handles the execution of the event loop cycle and implements

src/strands/event_loop/event_loop.py

@@ -34,6 +34,7 @@
     MaxTokensReachedException,
     ModelThrottledException,
 )
+from ..types.invocation import InvocationState
 from ..types.streaming import Metrics, StopReason
 from ..types.tools import ToolChoice, ToolChoiceAuto, ToolConfig, ToolGenerator, ToolResult, ToolUse
 from ._recover_message_on_max_tokens_reached import recover_message_on_max_tokens_reached
@@ -49,7 +50,7 @@
 MAX_DELAY = 240  # 4 minutes
 
 
-async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) -> AsyncGenerator[dict[str, Any], None]:
+async def event_loop_cycle(agent: "Agent", invocation_state: InvocationState) -> AsyncGenerator[dict[str, Any], None]:
     """Execute a single cycle of the event loop.
 
     This core function processes a single conversation turn, handling model inference, tool execution, and error
@@ -273,7 +274,7 @@ async def event_loop_cycle(agent: "Agent", invocation_state: dict[str, Any]) ->
     yield {"stop": (stop_reason, message, agent.event_loop_metrics, invocation_state["request_state"])}
 
 
-async def recurse_event_loop(agent: "Agent", invocation_state: dict[str, Any]) -> AsyncGenerator[dict[str, Any], None]:
+async def recurse_event_loop(agent: "Agent", invocation_state: InvocationState) -> AsyncGenerator[dict[str, Any], None]:
     """Make a recursive call to event_loop_cycle with the current state.
 
     This function is used when the event loop needs to continue processing after tool execution.
@@ -306,7 +307,7 @@ async def recurse_event_loop(agent: "Agent", invocation_state: dict[str, Any]) -
     recursive_trace.end()
 
 
-async def run_tool(agent: "Agent", tool_use: ToolUse, invocation_state: dict[str, Any]) -> ToolGenerator:
+async def run_tool(agent: "Agent", tool_use: ToolUse, invocation_state: InvocationState) -> ToolGenerator:
     """Process a tool invocation.
 
     Looks up the tool in the registry and streams it with the provided parameters.
@@ -429,7 +430,7 @@ async def _handle_tool_execution(
     cycle_trace: Trace,
     cycle_span: Any,
     cycle_start_time: float,
-    invocation_state: dict[str, Any],
+    invocation_state: InvocationState,
 ) -> AsyncGenerator[dict[str, Any], None]:
     tool_uses: list[ToolUse] = []
     tool_results: list[ToolResult] = []

src/strands/experimental/hooks/events.py

@@ -8,6 +8,7 @@
 
 from ...hooks import HookEvent
 from ...types.content import Message
+from ...types.invocation import InvocationState
 from ...types.streaming import StopReason
 from ...types.tools import AgentTool, ToolResult, ToolUse
 
@@ -30,7 +31,7 @@ class BeforeToolInvocationEvent(HookEvent):
 
     selected_tool: Optional[AgentTool]
     tool_use: ToolUse
-    invocation_state: dict[str, Any]
+    invocation_state: InvocationState
 
     def _can_write(self, name: str) -> bool:
         return name in ["selected_tool", "tool_use"]
@@ -57,7 +58,7 @@ class AfterToolInvocationEvent(HookEvent):
 
     selected_tool: Optional[AgentTool]
     tool_use: ToolUse
-    invocation_state: dict[str, Any]
+    invocation_state: InvocationState
     result: ToolResult
     exception: Optional[Exception] = None
 

src/strands/telemetry/tracer.py

@@ -13,6 +13,8 @@
 from opentelemetry.instrumentation.threading import ThreadingInstrumentor
 from opentelemetry.trace import Span, StatusCode
 
+from ..types.invocation import InvocationState
+
 from ..agent.agent_result import AgentResult
 from ..types.content import ContentBlock, Message, Messages
 from ..types.streaming import StopReason, Usage
@@ -343,7 +345,7 @@ def end_tool_call_span(
 
     def start_event_loop_cycle_span(
         self,
-        invocation_state: Any,
+        invocation_state: InvocationState,
         messages: Messages,
         parent_span: Optional[Span] = None,
         **kwargs: Any,

src/strands/tools/decorator.py

@@ -61,6 +61,7 @@ def my_tool(param1: str, param2: int = 42) -> dict:
 from pydantic import BaseModel, Field, create_model
 from typing_extensions import override
 
+from ..types.invocation import InvocationState
 from ..types.tools import AgentTool, JSONSchema, StrandsContext, ToolGenerator, ToolSpec, ToolUse
 
 logger = logging.getLogger(__name__)
@@ -253,7 +254,7 @@ def validate_input(self, input_data: dict[str, Any]) -> dict[str, Any]:
             raise ValueError(f"Validation failed for input parameters: {error_msg}") from e
 
     def inject_special_parameters(
-        self, validated_input: dict[str, Any], tool_use: ToolUse, invocation_state: dict[str, Any]
+        self, validated_input: dict[str, Any], tool_use: ToolUse, invocation_state: InvocationState
     ) -> None:
         """Inject special framework-provided parameters into the validated input.
 
@@ -413,7 +414,7 @@ def tool_type(self) -> str:
         return "function"
 
     @override
-    async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kwargs: Any) -> ToolGenerator:
+    async def stream(self, tool_use: ToolUse, invocation_state: InvocationState, **kwargs: Any) -> ToolGenerator:
         """Stream the tool with a tool use specification.
 
         This method handles tool use streams from a Strands Agent. It validates the input,

src/strands/tools/mcp/mcp_agent_tool.py

@@ -11,6 +11,7 @@
 from mcp.types import Tool as MCPTool
 from typing_extensions import override
 
+from ...types.invocation import InvocationState
 from ...types.tools import AgentTool, ToolGenerator, ToolSpec, ToolUse
 
 if TYPE_CHECKING:
@@ -75,7 +76,7 @@ def tool_type(self) -> str:
         return "python"
 
     @override
-    async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kwargs: Any) -> ToolGenerator:
+    async def stream(self, tool_use: ToolUse, invocation_state: InvocationState, **kwargs: Any) -> ToolGenerator:
         """Stream the MCP tool.
 
         This method delegates the tool stream to the MCP server connection, passing the tool use ID, tool name, and

src/strands/tools/tools.py

@@ -12,6 +12,7 @@
 
 from typing_extensions import override
 
+from ..types.invocation import InvocationState
 from ..types.tools import AgentTool, ToolFunc, ToolGenerator, ToolSpec, ToolUse
 
 logger = logging.getLogger(__name__)
@@ -198,7 +199,7 @@ def tool_type(self) -> str:
         return "python"
 
     @override
-    async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kwargs: Any) -> ToolGenerator:
+    async def stream(self, tool_use: ToolUse, invocation_state: InvocationState, **kwargs: Any) -> ToolGenerator:
         """Stream the Python function with the given tool use request.
 
         Args:
@@ -210,7 +211,7 @@ async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kw
             Tool events with the last being the tool result.
         """
         if inspect.iscoroutinefunction(self._tool_func):
-            result = await self._tool_func(tool_use, **invocation_state)
+            result = await self._tool_func(tool_use, **invocation_state) # this will fail if invocation state and kwargs overlap
         else:
             result = await asyncio.to_thread(self._tool_func, tool_use, **invocation_state)
 

src/strands/types/__init__.py

@@ -1,5 +1,6 @@
 """SDK type definitions."""
 
 from .collections import PaginatedList
+from .invocation import InvocationState
 
 __all__ = ["PaginatedList"]

src/strands/types/invocation.py

@@ -0,0 +1,45 @@
+"""Types for agent invocation state and context."""
+
+from typing import Any, Dict, TypedDict, TYPE_CHECKING
+
+from opentelemetry.trace import Span
+
+if TYPE_CHECKING:
+    from ..agent import Agent
+    from ..telemetry import Trace
+
+class InvocationState(TypedDict, total=False):
+    """Type definition for invocation_state used throughout the agent framework.
+
+    This TypedDict defines the structure of the invocation_state dictionary that is
+    passed through the agent's event loop and tool execution pipeline. All fields
+    are optional since invocation_state is built incrementally during execution.
+
+    Core Framework Fields:
+        agent: The Agent instance executing the invocation (added for backward compatibility).
+        event_loop_cycle_id: Unique identifier for the current event loop cycle.
+        request_state: State dictionary maintained across event loop cycles.
+        event_loop_cycle_trace: Trace object for monitoring the current cycle.
+        event_loop_cycle_span: Span object for distributed tracing.
+        event_loop_parent_cycle_id: UUID of the parent cycle for recursive calls. # always uuid? or just string?
+        event_loop_parent_span: Parent span for tracing hierarchy.
+
+    Additional Fields:
+        Any additional keyword arguments passed during agent invocation or added
+        by hooks and tools during execution are also included in this state.
+    """
+
+    # Core agent reference
+    agent: "Agent"  # Forward reference to avoid circular imports
+
+    # Event loop cycle management
+    event_loop_cycle_id: str
+    event_loop_parent_cycle_id: str
+
+    # State management
+    request_state: Dict[str, Any]
+
+    # Tracing and monitoring
+    event_loop_cycle_trace: "Trace"  # "Trace"  # Trace object type varies by implementation
+    event_loop_cycle_span: Span   # Span object type varies by implementation
+    event_loop_parent_span: Span  # Parent span for tracing hierarchy

src/strands/types/tools.py

@@ -8,6 +8,8 @@
 from abc import ABC, abstractmethod
 from typing import Any, AsyncGenerator, Awaitable, Callable, Literal, Protocol, Union
 
+from .invocation import InvocationState
+
 from typing_extensions import TypedDict
 
 from .media import DocumentContent, ImageContent
@@ -129,7 +131,7 @@ class StrandsContext(TypedDict, total=False):
     """
 
     tool_use: ToolUse
-    invocation_state: dict[str, Any]
+    invocation_state: InvocationState
 
 
 ToolChoice = Union[
@@ -231,7 +233,7 @@ def supports_hot_reload(self) -> bool:
 
     @abstractmethod
     # pragma: no cover
-    def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kwargs: Any) -> ToolGenerator:
+    def stream(self, tool_use: ToolUse, invocation_state: InvocationState, **kwargs: Any) -> ToolGenerator:
         """Stream tool events and return the final result.
 
         Args: