diff --git a/src/strands/agent/agent.py b/src/strands/agent/agent.py index 7c63c1e89d..8137f1887a 100644 --- a/src/strands/agent/agent.py +++ b/src/strands/agent/agent.py @@ -33,6 +33,7 @@ from .. import _identifier from .._async import run_async from ..event_loop.event_loop import event_loop_cycle +from ..tools._tool_helpers import generate_missing_tool_result_content if TYPE_CHECKING: from ..experimental.tools import ToolProvider @@ -57,7 +58,7 @@ from ..tools.watcher import ToolWatcher from ..types._events import AgentResultEvent, InitEventLoopEvent, ModelStreamChunkEvent, ToolInterruptEvent, TypedEvent from ..types.agent import AgentInput -from ..types.content import ContentBlock, Message, Messages +from ..types.content import ContentBlock, Message, Messages, SystemContentBlock from ..types.exceptions import ContextWindowOverflowException from ..types.interrupt import InterruptResponseContent from ..types.tools import ToolResult, ToolUse @@ -216,7 +217,7 @@ def __init__( model: Union[Model, str, None] = None, messages: Optional[Messages] = None, tools: Optional[list[Union[str, dict[str, str], "ToolProvider", Any]]] = None, - system_prompt: Optional[str] = None, + system_prompt: Optional[str | list[SystemContentBlock]] = None, structured_output_model: Optional[Type[BaseModel]] = None, callback_handler: Optional[ Union[Callable[..., Any], _DefaultCallbackHandlerSentinel] @@ -253,6 +254,7 @@ def __init__( If provided, only these tools will be available. If None, all tools will be available. system_prompt: System prompt to guide model behavior. + Can be a string or a list of SystemContentBlock objects for advanced features like caching. If None, the model will behave according to its default settings. structured_output_model: Pydantic model type(s) for structured output. When specified, all agent calls will attempt to return structured output of this type. @@ -280,14 +282,15 @@ def __init__( Defaults to None. session_manager: Manager for handling agent sessions including conversation history and state. If provided, enables session-based persistence and state management. - tool_executor: Definition of tool execution stragety (e.g., sequential, concurrent, etc.). + tool_executor: Definition of tool execution strategy (e.g., sequential, concurrent, etc.). Raises: ValueError: If agent id contains path separators. """ self.model = BedrockModel() if not model else BedrockModel(model_id=model) if isinstance(model, str) else model self.messages = messages if messages is not None else [] - self.system_prompt = system_prompt + # initializing self.system_prompt for backwards compatibility + self.system_prompt, self._system_prompt_content = self._initialize_system_prompt(system_prompt) self._default_structured_output_model = structured_output_model self.agent_id = _identifier.validate(agent_id or _DEFAULT_AGENT_ID, _identifier.Identifier.AGENT) self.name = name or _DEFAULT_AGENT_NAME @@ -816,6 +819,21 @@ def _convert_prompt_to_messages(self, prompt: AgentInput) -> Messages: messages: Messages | None = None if prompt is not None: + # Check if the latest message is toolUse + if len(self.messages) > 0 and any("toolUse" in content for content in self.messages[-1]["content"]): + # Add toolResult message after to have a valid conversation + logger.info( + "Agents latest message is toolUse, appending a toolResult message to have valid conversation." + ) + tool_use_ids = [ + content["toolUse"]["toolUseId"] for content in self.messages[-1]["content"] if "toolUse" in content + ] + self._append_message( + { + "role": "user", + "content": generate_missing_tool_result_content(tool_use_ids), + } + ) if isinstance(prompt, str): # String input - convert to user message messages = [{"role": "user", "content": [{"text": prompt}]}] @@ -965,6 +983,30 @@ def _filter_tool_parameters_for_recording(self, tool_name: str, input_params: di properties = tool_spec["inputSchema"]["json"]["properties"] return {k: v for k, v in input_params.items() if k in properties} + def _initialize_system_prompt( + self, system_prompt: str | list[SystemContentBlock] | None + ) -> tuple[str | None, list[SystemContentBlock] | None]: + """Initialize system prompt fields from constructor input. + + Maintains backwards compatibility by keeping system_prompt as str when string input + provided, avoiding breaking existing consumers. + + Maps system_prompt input to both string and content block representations: + - If string: system_prompt=string, _system_prompt_content=[{text: string}] + - If list with text elements: system_prompt=concatenated_text, _system_prompt_content=list + - If list without text elements: system_prompt=None, _system_prompt_content=list + - If None: system_prompt=None, _system_prompt_content=None + """ + if isinstance(system_prompt, str): + return system_prompt, [{"text": system_prompt}] + elif isinstance(system_prompt, list): + # Concatenate all text elements for backwards compatibility, None if no text found + text_parts = [block["text"] for block in system_prompt if "text" in block] + system_prompt_str = "\n".join(text_parts) if text_parts else None + return system_prompt_str, system_prompt + else: + return None, None + def _append_message(self, message: Message) -> None: """Appends a message to the agent's list of messages and invokes the callbacks for the MessageCreatedEvent.""" self.messages.append(message) diff --git a/src/strands/event_loop/event_loop.py b/src/strands/event_loop/event_loop.py index 3ea0097d84..66174c09fc 100644 --- a/src/strands/event_loop/event_loop.py +++ b/src/strands/event_loop/event_loop.py @@ -335,7 +335,12 @@ async def _handle_model_execution( tool_specs = agent.tool_registry.get_all_tool_specs() try: async for event in stream_messages( - agent.model, agent.system_prompt, agent.messages, tool_specs, structured_output_context.tool_choice + agent.model, + agent.system_prompt, + agent.messages, + tool_specs, + system_prompt_content=agent._system_prompt_content, + tool_choice=structured_output_context.tool_choice, ): yield event diff --git a/src/strands/event_loop/streaming.py b/src/strands/event_loop/streaming.py index 012a2d7620..c7b0b2caa3 100644 --- a/src/strands/event_loop/streaming.py +++ b/src/strands/event_loop/streaming.py @@ -22,7 +22,7 @@ TypedEvent, ) from ..types.citations import CitationsContentBlock -from ..types.content import ContentBlock, Message, Messages +from ..types.content import ContentBlock, Message, Messages, SystemContentBlock from ..types.streaming import ( ContentBlockDeltaEvent, ContentBlockStart, @@ -418,16 +418,22 @@ async def stream_messages( system_prompt: Optional[str], messages: Messages, tool_specs: list[ToolSpec], + *, tool_choice: Optional[Any] = None, + system_prompt_content: Optional[list[SystemContentBlock]] = None, + **kwargs: Any, ) -> AsyncGenerator[TypedEvent, None]: """Streams messages to the model and processes the response. Args: model: Model provider. - system_prompt: The system prompt to send. + system_prompt: The system prompt string, used for backwards compatibility with models that expect it. messages: List of messages to send. tool_specs: The list of tool specs. tool_choice: Optional tool choice constraint for forcing specific tool usage. + system_prompt_content: The authoritative system prompt content blocks that always contains the + system prompt data. + **kwargs: Additional keyword arguments for future extensibility. Yields: The reason for stopping, the final message, and the usage metrics @@ -436,7 +442,14 @@ async def stream_messages( messages = _normalize_messages(messages) start_time = time.time() - chunks = model.stream(messages, tool_specs if tool_specs else None, system_prompt, tool_choice=tool_choice) + + chunks = model.stream( + messages, + tool_specs if tool_specs else None, + system_prompt, + tool_choice=tool_choice, + system_prompt_content=system_prompt_content, + ) async for event in process_stream(chunks, start_time): yield event diff --git a/src/strands/models/bedrock.py b/src/strands/models/bedrock.py index c84cd0e3dc..4a7c816729 100644 --- a/src/strands/models/bedrock.py +++ b/src/strands/models/bedrock.py @@ -20,7 +20,7 @@ from ..event_loop import streaming from ..tools import convert_pydantic_to_tool_spec from ..tools._tool_helpers import noop_tool -from ..types.content import ContentBlock, Messages +from ..types.content import ContentBlock, Messages, SystemContentBlock from ..types.exceptions import ( ContextWindowOverflowException, ModelThrottledException, @@ -187,11 +187,11 @@ def get_config(self) -> BedrockConfig: """ return self.config - def format_request( + def _format_request( self, messages: Messages, tool_specs: Optional[list[ToolSpec]] = None, - system_prompt: Optional[str] = None, + system_prompt_content: Optional[list[SystemContentBlock]] = None, tool_choice: ToolChoice | None = None, ) -> dict[str, Any]: """Format a Bedrock converse stream request. @@ -201,6 +201,7 @@ def format_request( tool_specs: List of tool specifications to make available to the model. system_prompt: System prompt to provide context to the model. tool_choice: Selection strategy for tool invocation. + system_prompt_content: System prompt content blocks to provide context to the model. Returns: A Bedrock converse stream request. @@ -211,13 +212,20 @@ def format_request( ) if has_tool_content: tool_specs = [noop_tool.tool_spec] + + # Use system_prompt_content directly (copy for mutability) + system_blocks: list[SystemContentBlock] = system_prompt_content.copy() if system_prompt_content else [] + # Add cache point if configured (backwards compatibility) + if cache_prompt := self.config.get("cache_prompt"): + warnings.warn( + "cache_prompt is deprecated. Use SystemContentBlock with cachePoint instead.", UserWarning, stacklevel=3 + ) + system_blocks.append({"cachePoint": {"type": cache_prompt}}) + return { "modelId": self.config["model_id"], "messages": self._format_bedrock_messages(messages), - "system": [ - *([{"text": system_prompt}] if system_prompt else []), - *([{"cachePoint": {"type": self.config["cache_prompt"]}}] if self.config.get("cache_prompt") else []), - ], + "system": system_blocks, **( { "toolConfig": { @@ -590,6 +598,7 @@ async def stream( system_prompt: Optional[str] = None, *, tool_choice: ToolChoice | None = None, + system_prompt_content: Optional[list[SystemContentBlock]] = None, **kwargs: Any, ) -> AsyncGenerator[StreamEvent, None]: """Stream conversation with the Bedrock model. @@ -602,6 +611,7 @@ async def stream( tool_specs: List of tool specifications to make available to the model. system_prompt: System prompt to provide context to the model. tool_choice: Selection strategy for tool invocation. + system_prompt_content: System prompt content blocks to provide context to the model. **kwargs: Additional keyword arguments for future extensibility. Yields: @@ -620,7 +630,11 @@ def callback(event: Optional[StreamEvent] = None) -> None: loop = asyncio.get_event_loop() queue: asyncio.Queue[Optional[StreamEvent]] = asyncio.Queue() - thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt, tool_choice) + # Handle backward compatibility: if system_prompt is provided but system_prompt_content is None + if system_prompt and system_prompt_content is None: + system_prompt_content = [{"text": system_prompt}] + + thread = asyncio.to_thread(self._stream, callback, messages, tool_specs, system_prompt_content, tool_choice) task = asyncio.create_task(thread) while True: @@ -637,7 +651,7 @@ def _stream( callback: Callable[..., None], messages: Messages, tool_specs: Optional[list[ToolSpec]] = None, - system_prompt: Optional[str] = None, + system_prompt_content: Optional[list[SystemContentBlock]] = None, tool_choice: ToolChoice | None = None, ) -> None: """Stream conversation with the Bedrock model. @@ -649,7 +663,7 @@ def _stream( callback: Function to send events to the main thread. messages: List of message objects to be processed by the model. tool_specs: List of tool specifications to make available to the model. - system_prompt: System prompt to provide context to the model. + system_prompt_content: System prompt content blocks to provide context to the model. tool_choice: Selection strategy for tool invocation. Raises: @@ -658,7 +672,7 @@ def _stream( """ try: logger.debug("formatting request") - request = self.format_request(messages, tool_specs, system_prompt, tool_choice) + request = self._format_request(messages, tool_specs, system_prompt_content, tool_choice) logger.debug("request=<%s>", request) logger.debug("invoking model") diff --git a/src/strands/models/gemini.py b/src/strands/models/gemini.py index c288595e1c..c24d91a0d0 100644 --- a/src/strands/models/gemini.py +++ b/src/strands/models/gemini.py @@ -408,7 +408,13 @@ async def stream( if not error.message: raise - message = json.loads(error.message) + try: + message = json.loads(error.message) if error.message else {} + except json.JSONDecodeError as e: + logger.warning("error_message=<%s> | Gemini API returned non-JSON error", error.message) + # Re-raise the original ClientError (not JSONDecodeError) and make the JSON error the explicit cause + raise error from e + match message["error"]["status"]: case "RESOURCE_EXHAUSTED" | "UNAVAILABLE": raise ModelThrottledException(error.message) from error diff --git a/src/strands/models/model.py b/src/strands/models/model.py index 7f178660a0..b2fa738024 100644 --- a/src/strands/models/model.py +++ b/src/strands/models/model.py @@ -6,7 +6,7 @@ from pydantic import BaseModel -from ..types.content import Messages +from ..types.content import Messages, SystemContentBlock from ..types.streaming import StreamEvent from ..types.tools import ToolChoice, ToolSpec @@ -72,6 +72,7 @@ def stream( system_prompt: Optional[str] = None, *, tool_choice: ToolChoice | None = None, + system_prompt_content: list[SystemContentBlock] | None = None, **kwargs: Any, ) -> AsyncIterable[StreamEvent]: """Stream conversation with the model. @@ -87,6 +88,7 @@ def stream( tool_specs: List of tool specifications to make available to the model. system_prompt: System prompt to provide context to the model. tool_choice: Selection strategy for tool invocation. + system_prompt_content: System prompt content blocks for advanced features like caching. **kwargs: Additional keyword arguments for future extensibility. Yields: diff --git a/src/strands/multiagent/graph.py b/src/strands/multiagent/graph.py index 2d3d538fe9..b421b70c1a 100644 --- a/src/strands/multiagent/graph.py +++ b/src/strands/multiagent/graph.py @@ -26,6 +26,15 @@ from .._async import run_async from ..agent import Agent from ..agent.state import AgentState +from ..experimental.hooks.multiagent import ( + AfterMultiAgentInvocationEvent, + AfterNodeCallEvent, + BeforeMultiAgentInvocationEvent, + BeforeNodeCallEvent, + MultiAgentInitializedEvent, +) +from ..hooks import HookProvider, HookRegistry +from ..session import SessionManager from ..telemetry import get_tracer from ..types._events import ( MultiAgentHandoffEvent, @@ -40,6 +49,8 @@ logger = logging.getLogger(__name__) +_DEFAULT_GRAPH_ID = "default_graph" + @dataclass class GraphState: @@ -223,6 +234,9 @@ def __init__(self) -> None: self._execution_timeout: Optional[float] = None self._node_timeout: Optional[float] = None self._reset_on_revisit: bool = False + self._id: str = _DEFAULT_GRAPH_ID + self._session_manager: Optional[SessionManager] = None + self._hooks: Optional[list[HookProvider]] = None def add_node(self, executor: Agent | MultiAgentBase, node_id: str | None = None) -> GraphNode: """Add an Agent or MultiAgentBase instance as a node to the graph.""" @@ -313,6 +327,33 @@ def set_node_timeout(self, timeout: float) -> "GraphBuilder": self._node_timeout = timeout return self + def set_graph_id(self, graph_id: str) -> "GraphBuilder": + """Set graph id. + + Args: + graph_id: Unique graph id + """ + self._id = graph_id + return self + + def set_session_manager(self, session_manager: SessionManager) -> "GraphBuilder": + """Set session manager for the graph. + + Args: + session_manager: SessionManager instance + """ + self._session_manager = session_manager + return self + + def set_hook_providers(self, hooks: list[HookProvider]) -> "GraphBuilder": + """Set hook providers for the graph. + + Args: + hooks: Customer hooks user passes in + """ + self._hooks = hooks + return self + def build(self) -> "Graph": """Build and validate the graph with configured settings.""" if not self.nodes: @@ -338,6 +379,9 @@ def build(self) -> "Graph": execution_timeout=self._execution_timeout, node_timeout=self._node_timeout, reset_on_revisit=self._reset_on_revisit, + session_manager=self._session_manager, + hooks=self._hooks, + id=self._id, ) def _validate_graph(self) -> None: @@ -365,6 +409,9 @@ def __init__( execution_timeout: Optional[float] = None, node_timeout: Optional[float] = None, reset_on_revisit: bool = False, + session_manager: Optional[SessionManager] = None, + hooks: Optional[list[HookProvider]] = None, + id: str = _DEFAULT_GRAPH_ID, ) -> None: """Initialize Graph with execution limits and reset behavior. @@ -376,6 +423,9 @@ def __init__( execution_timeout: Total execution timeout in seconds (default: None - no limit) node_timeout: Individual node timeout in seconds (default: None - no limit) reset_on_revisit: Whether to reset node state when revisited (default: False) + session_manager: Session manager for persisting graph state and execution history (default: None) + hooks: List of hook providers for monitoring and extending graph execution behavior (default: None) + id: Unique graph id (default: None) """ super().__init__() @@ -391,6 +441,19 @@ def __init__( self.reset_on_revisit = reset_on_revisit self.state = GraphState() self.tracer = get_tracer() + self.session_manager = session_manager + self.hooks = HookRegistry() + if self.session_manager: + self.hooks.add_hook(self.session_manager) + if hooks: + for hook in hooks: + self.hooks.add_hook(hook) + + self._resume_next_nodes: list[GraphNode] = [] + self._resume_from_session = False + self.id = id + + self.hooks.invoke_callbacks(MultiAgentInitializedEvent(self)) def __call__( self, task: str | list[ContentBlock], invocation_state: dict[str, Any] | None = None, **kwargs: Any @@ -453,18 +516,25 @@ async def stream_async( if invocation_state is None: invocation_state = {} + self.hooks.invoke_callbacks(BeforeMultiAgentInvocationEvent(self, invocation_state)) + logger.debug("task=<%s> | starting graph execution", task) # Initialize state start_time = time.time() - self.state = GraphState( - status=Status.EXECUTING, - task=task, - total_nodes=len(self.nodes), - edges=[(edge.from_node, edge.to_node) for edge in self.edges], - entry_points=list(self.entry_points), - start_time=start_time, - ) + if not self._resume_from_session: + # Initialize state + self.state = GraphState( + status=Status.EXECUTING, + task=task, + total_nodes=len(self.nodes), + edges=[(edge.from_node, edge.to_node) for edge in self.edges], + entry_points=list(self.entry_points), + start_time=start_time, + ) + else: + self.state.status = Status.EXECUTING + self.state.start_time = start_time span = self.tracer.start_multiagent_span(task, "graph") with trace_api.use_span(span, end_on_exit=True): @@ -499,6 +569,9 @@ async def stream_async( raise finally: self.state.execution_time = round((time.time() - start_time) * 1000) + self.hooks.invoke_callbacks(AfterMultiAgentInvocationEvent(self)) + self._resume_from_session = False + self._resume_next_nodes.clear() def _validate_graph(self, nodes: dict[str, GraphNode]) -> None: """Validate graph nodes for duplicate instances.""" @@ -514,7 +587,7 @@ def _validate_graph(self, nodes: dict[str, GraphNode]) -> None: async def _execute_graph(self, invocation_state: dict[str, Any]) -> AsyncIterator[Any]: """Execute graph and yield TypedEvent objects.""" - ready_nodes = list(self.entry_points) + ready_nodes = self._resume_next_nodes if self._resume_from_session else list(self.entry_points) while ready_nodes: # Check execution limits before continuing @@ -703,7 +776,9 @@ def _is_node_ready_with_conditions(self, node: GraphNode, completed_batch: list[ async def _execute_node(self, node: GraphNode, invocation_state: dict[str, Any]) -> AsyncIterator[Any]: """Execute a single node and yield TypedEvent objects.""" - # Reset the node's state if reset_on_revisit is enabled and it's being revisited + self.hooks.invoke_callbacks(BeforeNodeCallEvent(self, node.node_id, invocation_state)) + + # Reset the node's state if reset_on_revisit is enabled, and it's being revisited if self.reset_on_revisit and node in self.state.completed_nodes: logger.debug("node_id=<%s> | resetting node state for revisit", node.node_id) node.reset_executor_state() @@ -844,6 +919,9 @@ async def _execute_node(self, node: GraphNode, invocation_state: dict[str, Any]) # Re-raise to stop graph execution (fail-fast behavior) raise + finally: + self.hooks.invoke_callbacks(AfterNodeCallEvent(self, node.node_id, invocation_state)) + def _accumulate_metrics(self, node_result: NodeResult) -> None: """Accumulate metrics from a node result.""" self.state.accumulated_usage["inputTokens"] += node_result.accumulated_usage.get("inputTokens", 0) @@ -928,3 +1006,94 @@ def _build_result(self) -> GraphResult: edges=self.state.edges, entry_points=self.state.entry_points, ) + + def serialize_state(self) -> dict[str, Any]: + """Serialize the current graph state to a dictionary.""" + compute_nodes = self._compute_ready_nodes_for_resume() + next_nodes = [n.node_id for n in compute_nodes] if compute_nodes else [] + return { + "type": "graph", + "id": self.id, + "status": self.state.status.value, + "completed_nodes": [n.node_id for n in self.state.completed_nodes], + "failed_nodes": [n.node_id for n in self.state.failed_nodes], + "node_results": {k: v.to_dict() for k, v in (self.state.results or {}).items()}, + "next_nodes_to_execute": next_nodes, + "current_task": self.state.task, + "execution_order": [n.node_id for n in self.state.execution_order], + } + + def deserialize_state(self, payload: dict[str, Any]) -> None: + """Restore graph state from a session dict and prepare for execution. + + This method handles two scenarios: + 1. If the graph execution ended (no next_nodes_to_execute, eg: Completed, or Failed with dead end nodes), + resets all nodes and graph state to allow re-execution from the beginning. + 2. If the graph execution was interrupted mid-execution (has next_nodes_to_execute), + restores the persisted state and prepares to resume execution from the next ready nodes. + + Args: + payload: Dictionary containing persisted state data including status, + completed nodes, results, and next nodes to execute. + """ + if not payload.get("next_nodes_to_execute"): + # Reset all nodes + for node in self.nodes.values(): + node.reset_executor_state() + # Reset graph state + self.state = GraphState() + self._resume_from_session = False + return + else: + self._from_dict(payload) + self._resume_from_session = True + + def _compute_ready_nodes_for_resume(self) -> list[GraphNode]: + if self.state.status == Status.PENDING: + return [] + ready_nodes: list[GraphNode] = [] + completed_nodes = set(self.state.completed_nodes) + for node in self.nodes.values(): + if node in completed_nodes: + continue + incoming = [e for e in self.edges if e.to_node is node] + if not incoming: + ready_nodes.append(node) + elif all(e.from_node in completed_nodes and e.should_traverse(self.state) for e in incoming): + ready_nodes.append(node) + + return ready_nodes + + def _from_dict(self, payload: dict[str, Any]) -> None: + self.state.status = Status(payload["status"]) + # Hydrate completed nodes & results + raw_results = payload.get("node_results") or {} + results: dict[str, NodeResult] = {} + for node_id, entry in raw_results.items(): + if node_id not in self.nodes: + continue + try: + results[node_id] = NodeResult.from_dict(entry) + except Exception: + logger.exception("Failed to hydrate NodeResult for node_id=%s; skipping.", node_id) + raise + self.state.results = results + + self.state.failed_nodes = set( + self.nodes[node_id] for node_id in (payload.get("failed_nodes") or []) if node_id in self.nodes + ) + + # Restore completed nodes from persisted data + completed_node_ids = payload.get("completed_nodes") or [] + self.state.completed_nodes = {self.nodes[node_id] for node_id in completed_node_ids if node_id in self.nodes} + + # Execution order (only nodes that still exist) + order_node_ids = payload.get("execution_order") or [] + self.state.execution_order = [self.nodes[node_id] for node_id in order_node_ids if node_id in self.nodes] + + # Task + self.state.task = payload.get("current_task", self.state.task) + + # next nodes to execute + next_nodes = [self.nodes[nid] for nid in (payload.get("next_nodes_to_execute") or []) if nid in self.nodes] + self._resume_next_nodes = next_nodes diff --git a/src/strands/multiagent/swarm.py b/src/strands/multiagent/swarm.py index cd0a2d74ce..accd564630 100644 --- a/src/strands/multiagent/swarm.py +++ b/src/strands/multiagent/swarm.py @@ -18,13 +18,22 @@ import logging import time from dataclasses import dataclass, field -from typing import Any, AsyncIterator, Callable, Tuple, cast +from typing import Any, AsyncIterator, Callable, Optional, Tuple, cast from opentelemetry import trace as trace_api from .._async import run_async from ..agent import Agent from ..agent.state import AgentState +from ..experimental.hooks.multiagent import ( + AfterMultiAgentInvocationEvent, + AfterNodeCallEvent, + BeforeMultiAgentInvocationEvent, + BeforeNodeCallEvent, + MultiAgentInitializedEvent, +) +from ..hooks import HookProvider, HookRegistry +from ..session import SessionManager from ..telemetry import get_tracer from ..tools.decorator import tool from ..types._events import ( @@ -40,6 +49,8 @@ logger = logging.getLogger(__name__) +_DEFAULT_SWARM_ID = "default_swarm" + @dataclass class SwarmNode: @@ -210,10 +221,14 @@ def __init__( node_timeout: float = 300.0, repetitive_handoff_detection_window: int = 0, repetitive_handoff_min_unique_agents: int = 0, + session_manager: Optional[SessionManager] = None, + hooks: Optional[list[HookProvider]] = None, + id: str = _DEFAULT_SWARM_ID, ) -> None: """Initialize Swarm with agents and configuration. Args: + id : Unique swarm id (default: None) nodes: List of nodes (e.g. Agent) to include in the swarm entry_point: Agent to start with. If None, uses the first agent (default: None) max_handoffs: Maximum handoffs to agents and users (default: 20) @@ -224,9 +239,11 @@ def __init__( Disabled by default (default: 0) repetitive_handoff_min_unique_agents: Minimum unique agents required in recent sequence Disabled by default (default: 0) + session_manager: Session manager for persisting graph state and execution history (default: None) + hooks: List of hook providers for monitoring and extending graph execution behavior (default: None) """ super().__init__() - + self.id = id self.entry_point = entry_point self.max_handoffs = max_handoffs self.max_iterations = max_iterations @@ -244,8 +261,19 @@ def __init__( ) self.tracer = get_tracer() + self.session_manager = session_manager + self.hooks = HookRegistry() + if hooks: + for hook in hooks: + self.hooks.add_hook(hook) + if self.session_manager: + self.hooks.add_hook(self.session_manager) + + self._resume_from_session = False + self._setup_swarm(nodes) self._inject_swarm_tools() + self.hooks.invoke_callbacks(MultiAgentInitializedEvent(self)) def __call__( self, task: str | list[ContentBlock], invocation_state: dict[str, Any] | None = None, **kwargs: Any @@ -260,7 +288,6 @@ def __call__( """ if invocation_state is None: invocation_state = {} - return run_async(lambda: self.invoke_async(task, invocation_state)) async def invoke_async( @@ -309,22 +336,24 @@ async def stream_async( if invocation_state is None: invocation_state = {} + self.hooks.invoke_callbacks(BeforeMultiAgentInvocationEvent(self, invocation_state)) + logger.debug("starting swarm execution") - # Initialize swarm state with configuration - if self.entry_point: - initial_node = self.nodes[str(self.entry_point.name)] - else: - initial_node = next(iter(self.nodes.values())) + if not self._resume_from_session: + # Initialize swarm state with configuration + initial_node = self._initial_node() - self.state = SwarmState( - current_node=initial_node, - task=task, - completion_status=Status.EXECUTING, - shared_context=self.shared_context, - ) + self.state = SwarmState( + current_node=initial_node, + task=task, + completion_status=Status.EXECUTING, + shared_context=self.shared_context, + ) + else: + self.state.completion_status = Status.EXECUTING + self.state.start_time = time.time() - start_time = time.time() span = self.tracer.start_multiagent_span(task, "swarm") with trace_api.use_span(span, end_on_exit=True): try: @@ -345,7 +374,9 @@ async def stream_async( self.state.completion_status = Status.FAILED raise finally: - self.state.execution_time = round((time.time() - start_time) * 1000) + self.state.execution_time = round((time.time() - self.state.start_time) * 1000) + self.hooks.invoke_callbacks(AfterMultiAgentInvocationEvent(self, invocation_state)) + self._resume_from_session = False # Yield final result after execution_time is set result = self._build_result() @@ -656,6 +687,7 @@ async def _execute_swarm(self, invocation_state: dict[str, Any]) -> AsyncIterato # TODO: Implement cancellation token to stop _execute_node from continuing try: # Execute with timeout wrapper for async generator streaming + self.hooks.invoke_callbacks(BeforeNodeCallEvent(self, current_node.node_id, invocation_state)) node_stream = self._stream_with_timeout( self._execute_node(current_node, self.state.task, invocation_state), self.node_timeout, @@ -666,6 +698,9 @@ async def _execute_swarm(self, invocation_state: dict[str, Any]) -> AsyncIterato self.state.node_history.append(current_node) + # After self.state add current node, swarm state finish updating, we persist here + self.hooks.invoke_callbacks(AfterNodeCallEvent(self, current_node.node_id, invocation_state)) + logger.debug("node=<%s> | node execution completed", current_node.node_id) # Check if handoff occurred during execution @@ -823,3 +858,84 @@ def _build_result(self) -> SwarmResult: execution_time=self.state.execution_time, node_history=self.state.node_history, ) + + def serialize_state(self) -> dict[str, Any]: + """Serialize the current swarm state to a dictionary.""" + status_str = self.state.completion_status.value + next_nodes = ( + [self.state.current_node.node_id] + if self.state.completion_status == Status.EXECUTING and self.state.current_node + else [] + ) + + return { + "type": "swarm", + "id": self.id, + "status": status_str, + "node_history": [n.node_id for n in self.state.node_history], + "node_results": {k: v.to_dict() for k, v in self.state.results.items()}, + "next_nodes_to_execute": next_nodes, + "current_task": self.state.task, + "context": { + "shared_context": getattr(self.state.shared_context, "context", {}) or {}, + "handoff_message": self.state.handoff_message, + }, + } + + def deserialize_state(self, payload: dict[str, Any]) -> None: + """Restore swarm state from a session dict and prepare for execution. + + This method handles two scenarios: + 1. If the persisted status is COMPLETED, FAILED resets all nodes and graph state + to allow re-execution from the beginning. + 2. Otherwise, restores the persisted state and prepares to resume execution + from the next ready nodes. + + Args: + payload: Dictionary containing persisted state data including status, + completed nodes, results, and next nodes to execute. + """ + if not payload.get("next_nodes_to_execute"): + for node in self.nodes.values(): + node.reset_executor_state() + self.state = SwarmState( + current_node=SwarmNode("", Agent()), + task="", + completion_status=Status.PENDING, + ) + self._resume_from_session = False + return + else: + self._from_dict(payload) + self._resume_from_session = True + + def _from_dict(self, payload: dict[str, Any]) -> None: + self.state.completion_status = Status(payload["status"]) + # Hydrate completed nodes & results + context = payload["context"] or {} + self.shared_context.context = context.get("shared_context") or {} + self.state.handoff_message = context.get("handoff_message") + + self.state.node_history = [self.nodes[nid] for nid in (payload.get("node_history") or []) if nid in self.nodes] + + raw_results = payload.get("node_results") or {} + results: dict[str, NodeResult] = {} + for node_id, entry in raw_results.items(): + if node_id not in self.nodes: + continue + try: + results[node_id] = NodeResult.from_dict(entry) + except Exception: + logger.exception("Failed to hydrate NodeResult for node_id=%s; skipping.", node_id) + raise + self.state.results = results + self.state.task = payload.get("current_task", self.state.task) + + next_node_ids = payload.get("next_nodes_to_execute") or [] + if next_node_ids: + self.state.current_node = self.nodes[next_node_ids[0]] if next_node_ids[0] else self._initial_node() + + def _initial_node(self) -> SwarmNode: + if self.entry_point: + return self.nodes[str(self.entry_point.name)] + return next(iter(self.nodes.values())) # First SwarmNode diff --git a/src/strands/session/repository_session_manager.py b/src/strands/session/repository_session_manager.py index 86c6044a65..a042452d39 100644 --- a/src/strands/session/repository_session_manager.py +++ b/src/strands/session/repository_session_manager.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Optional from ..agent.state import AgentState +from ..tools._tool_helpers import generate_missing_tool_result_content from ..types.content import Message from ..types.exceptions import SessionException from ..types.session import ( @@ -159,6 +160,50 @@ def initialize(self, agent: "Agent", **kwargs: Any) -> None: # Restore the agents messages array including the optional prepend messages agent.messages = prepend_messages + [session_message.to_message() for session_message in session_messages] + # Fix broken session histories: https://github.com/strands-agents/sdk-python/issues/859 + agent.messages = self._fix_broken_tool_use(agent.messages) + + def _fix_broken_tool_use(self, messages: list[Message]) -> list[Message]: + """Add tool_result after orphaned tool_use messages. + + Before 1.15.0, strands had a bug where they persisted sessions with a potentially broken messages array. + This method retroactively fixes that issue by adding a tool_result outside of session management. After 1.15.0, + this bug is no longer present. + """ + for index, message in enumerate(messages): + # Check all but the latest message in the messages array + # The latest message being orphaned is handled in the agent class + if index + 1 < len(messages): + if any("toolUse" in content for content in message["content"]): + tool_use_ids = [ + content["toolUse"]["toolUseId"] for content in message["content"] if "toolUse" in content + ] + + # Check if there are more messages after the current toolUse message + tool_result_ids = [ + content["toolResult"]["toolUseId"] + for content in messages[index + 1]["content"] + if "toolResult" in content + ] + + missing_tool_use_ids = list(set(tool_use_ids) - set(tool_result_ids)) + # If there area missing tool use ids, that means the messages history is broken + if missing_tool_use_ids: + logger.warning( + "Session message history has an orphaned toolUse with no toolResult. " + "Adding toolResult content blocks to create valid conversation." + ) + # Create the missing toolResult content blocks + missing_content_blocks = generate_missing_tool_result_content(missing_tool_use_ids) + + if tool_result_ids: + # If there were any toolResult ids, that means only some of the content blocks are missing + messages[index + 1]["content"].extend(missing_content_blocks) + else: + # The message following the toolUse was not a toolResult, so lets insert it + messages.insert(index + 1, {"role": "user", "content": missing_content_blocks}) + return messages + def sync_multi_agent(self, source: "MultiAgentBase", **kwargs: Any) -> None: """Serialize and update the multi-agent state into the session repository. diff --git a/src/strands/tools/_tool_helpers.py b/src/strands/tools/_tool_helpers.py index d640f23b8f..d023caeec1 100644 --- a/src/strands/tools/_tool_helpers.py +++ b/src/strands/tools/_tool_helpers.py @@ -1,6 +1,7 @@ """Helpers for tools.""" -from strands.tools.decorator import tool +from ..tools.decorator import tool +from ..types.content import ContentBlock # https://github.com/strands-agents/sdk-python/issues/998 @@ -13,3 +14,17 @@ def noop_tool() -> None: summarization will fail. As a workaround, we register the no-op tool. """ pass + + +def generate_missing_tool_result_content(tool_use_ids: list[str]) -> list[ContentBlock]: + """Generate ToolResult content blocks for orphaned ToolUse message.""" + return [ + { + "toolResult": { + "toolUseId": tool_use_id, + "status": "error", + "content": [{"text": "Tool was interrupted."}], + } + } + for tool_use_id in tool_use_ids + ] diff --git a/src/strands/types/content.py b/src/strands/types/content.py index c3eddca4db..4d0bbe4124 100644 --- a/src/strands/types/content.py +++ b/src/strands/types/content.py @@ -103,11 +103,11 @@ class SystemContentBlock(TypedDict, total=False): """Contains configurations for instructions to provide the model for how to handle input. Attributes: - guardContent: A content block to assess with the guardrail. + cachePoint: A cache point configuration to optimize conversation history. text: A system prompt for the model. """ - guardContent: GuardContent + cachePoint: CachePoint text: str diff --git a/tests/fixtures/mocked_model_provider.py b/tests/fixtures/mocked_model_provider.py index 56817a6e4b..24de958bcd 100644 --- a/tests/fixtures/mocked_model_provider.py +++ b/tests/fixtures/mocked_model_provider.py @@ -58,6 +58,9 @@ async def stream( tool_specs: Optional[list[ToolSpec]] = None, system_prompt: Optional[str] = None, tool_choice: Optional[Any] = None, + *, + system_prompt_content=None, + **kwargs: Any, ) -> AsyncGenerator[Any, None]: events = self.map_agent_message_to_events(self.agent_responses[self.index]) for event in events: diff --git a/tests/strands/agent/test_agent.py b/tests/strands/agent/test_agent.py index 52840f1a25..3a0bc2dfbf 100644 --- a/tests/strands/agent/test_agent.py +++ b/tests/strands/agent/test_agent.py @@ -330,6 +330,7 @@ def test_agent__call__( [tool.tool_spec], system_prompt, tool_choice=None, + system_prompt_content=[{"text": system_prompt}], ), unittest.mock.call( [ @@ -367,6 +368,7 @@ def test_agent__call__( [tool.tool_spec], system_prompt, tool_choice=None, + system_prompt_content=[{"text": system_prompt}], ), ], ) @@ -487,6 +489,7 @@ def test_agent__call__retry_with_reduced_context(mock_model, agent, tool, agener unittest.mock.ANY, unittest.mock.ANY, tool_choice=None, + system_prompt_content=unittest.mock.ANY, ) conversation_manager_spy.reduce_context.assert_called_once() @@ -631,6 +634,7 @@ def test_agent__call__retry_with_overwritten_tool(mock_model, agent, tool, agene unittest.mock.ANY, unittest.mock.ANY, tool_choice=None, + system_prompt_content=unittest.mock.ANY, ) assert conversation_manager_spy.reduce_context.call_count == 2 @@ -2162,6 +2166,82 @@ def shell(command: str): assert agent.messages[-1] == {"content": [{"text": "I invoked a tool!"}], "role": "assistant"} +def test_agent_string_system_prompt(): + """Test initialization with string system prompt.""" + system_prompt = "You are a helpful assistant." + agent = Agent(system_prompt=system_prompt) + + assert agent.system_prompt == system_prompt + assert agent._system_prompt_content == [{"text": system_prompt}] + + +def test_agent_single_text_block_system_prompt(): + """Test initialization with single text SystemContentBlock.""" + text = "You are a helpful assistant." + system_prompt_content = [{"text": text}] + agent = Agent(system_prompt=system_prompt_content) + + assert agent.system_prompt == text + assert agent._system_prompt_content == system_prompt_content + + +def test_agent_multiple_blocks_system_prompt(): + """Test initialization with multiple SystemContentBlocks.""" + system_prompt_content = [ + {"text": "You are a helpful assistant."}, + {"cachePoint": {"type": "default"}}, + {"text": "Additional instructions."}, + ] + agent = Agent(system_prompt=system_prompt_content) + + assert agent.system_prompt == "You are a helpful assistant.\nAdditional instructions." + assert agent._system_prompt_content == system_prompt_content + + +def test_agent_single_non_text_block_system_prompt(): + """Test initialization with single non-text SystemContentBlock.""" + system_prompt_content = [{"cachePoint": {"type": "default"}}] + agent = Agent(system_prompt=system_prompt_content) + + assert agent.system_prompt is None + assert agent._system_prompt_content == system_prompt_content + + +def test_agent_none_system_prompt(): + """Test initialization with None system prompt.""" + agent = Agent(system_prompt=None) + + assert agent.system_prompt is None + assert agent._system_prompt_content is None + + +def test_agent_empty_list_system_prompt(): + """Test initialization with empty list system prompt.""" + agent = Agent(system_prompt=[]) + + assert agent.system_prompt is None + assert agent._system_prompt_content == [] + + +def test_agent_backwards_compatibility_string_access(): + """Test that string system prompts maintain backwards compatibility.""" + system_prompt = "You are a helpful assistant." + agent = Agent(system_prompt=system_prompt) + + # Should be able to access as string for backwards compatibility + assert agent.system_prompt == system_prompt + + +def test_agent_backwards_compatibility_single_text_block(): + """Test that single text blocks maintain backwards compatibility.""" + text = "You are a helpful assistant." + system_prompt_content = [{"text": text}] + agent = Agent(system_prompt=system_prompt_content) + + # Should extract text for backwards compatibility + assert agent.system_prompt == text + + @pytest.mark.parametrize( "content, expected", [ @@ -2215,3 +2295,143 @@ def test_redact_user_content(content, expected): agent = Agent() result = agent._redact_user_content(content, "REDACTED") assert result == expected + + +def test_agent_fixes_orphaned_tool_use_on_new_prompt(mock_model, agenerator): + """Test that agent adds toolResult for orphaned toolUse when called with new prompt.""" + mock_model.mock_stream.return_value = agenerator( + [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"start": {"text": ""}}}, + {"contentBlockDelta": {"delta": {"text": "Fixed!"}}}, + {"contentBlockStop": {}}, + {"messageStop": {"stopReason": "end_turn"}}, + ] + ) + + # Start with orphaned toolUse message + messages = [ + { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "orphaned-123", "name": "tool_decorated", "input": {"random_string": "test"}}} + ], + } + ] + + agent = Agent(model=mock_model, messages=messages) + + # Call with new prompt should fix orphaned toolUse + agent("Continue conversation") + + # Should have added toolResult message + assert len(agent.messages) >= 3 + assert agent.messages[1] == { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "orphaned-123", + "status": "error", + "content": [{"text": "Tool was interrupted."}], + } + } + ], + } + + +def test_agent_fixes_multiple_orphaned_tool_uses(mock_model, agenerator): + """Test that agent handles multiple orphaned toolUse messages.""" + mock_model.mock_stream.return_value = agenerator( + [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"start": {"text": ""}}}, + {"contentBlockDelta": {"delta": {"text": "Fixed multiple!"}}}, + {"contentBlockStop": {}}, + {"messageStop": {"stopReason": "end_turn"}}, + ] + ) + + messages = [ + { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "orphaned-123", + "name": "tool_decorated", + "input": {"random_string": "test1"}, + } + }, + { + "toolUse": { + "toolUseId": "orphaned-456", + "name": "tool_decorated", + "input": {"random_string": "test2"}, + } + }, + ], + } + ] + + agent = Agent(model=mock_model, messages=messages) + agent("Continue") + + # Should have toolResult for both toolUse IDs + assert agent.messages[1] == { + "role": "user", + "content": [ + { + "toolResult": { + "toolUseId": "orphaned-123", + "status": "error", + "content": [{"text": "Tool was interrupted."}], + } + }, + { + "toolResult": { + "toolUseId": "orphaned-456", + "status": "error", + "content": [{"text": "Tool was interrupted."}], + } + }, + ], + } + + +def test_agent_skips_fix_for_valid_conversation(mock_model, agenerator): + """Test that agent doesn't modify valid toolUse/toolResult pairs.""" + mock_model.mock_stream.return_value = agenerator( + [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"start": {"text": ""}}}, + {"contentBlockDelta": {"delta": {"text": "No fix needed!"}}}, + {"contentBlockStop": {}}, + {"messageStop": {"stopReason": "end_turn"}}, + ] + ) + + # Valid conversation with toolUse followed by toolResult + messages = [ + { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "valid-123", "name": "tool_decorated", "input": {"random_string": "test"}}} + ], + }, + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "valid-123", "status": "success", "content": [{"text": "result"}]}} + ], + }, + ] + + agent = Agent(model=mock_model, messages=messages) + original_length = len(agent.messages) + + agent("Continue") + + # Should not have added any toolResult messages + # Only the new user message and assistant response should be added + assert len(agent.messages) == original_length + 2 diff --git a/tests/strands/event_loop/test_event_loop.py b/tests/strands/event_loop/test_event_loop.py index 72c63e8970..72fe1b4bda 100644 --- a/tests/strands/event_loop/test_event_loop.py +++ b/tests/strands/event_loop/test_event_loop.py @@ -379,6 +379,7 @@ async def test_event_loop_cycle_tool_result( tool_registry.get_all_tool_specs(), "p1", tool_choice=None, + system_prompt_content=unittest.mock.ANY, ) diff --git a/tests/strands/event_loop/test_streaming.py b/tests/strands/event_loop/test_streaming.py index e75af40035..714fbac271 100644 --- a/tests/strands/event_loop/test_streaming.py +++ b/tests/strands/event_loop/test_streaming.py @@ -802,9 +802,10 @@ async def test_stream_messages(agenerator, alist): stream = strands.event_loop.streaming.stream_messages( mock_model, - system_prompt="test prompt", + system_prompt_content=[{"text": "test prompt"}], messages=[{"role": "assistant", "content": [{"text": "a"}, {"text": " \n"}]}], tool_specs=None, + system_prompt="test prompt", ) tru_events = await alist(stream) @@ -845,6 +846,135 @@ async def test_stream_messages(agenerator, alist): None, "test prompt", tool_choice=None, + system_prompt_content=[{"text": "test prompt"}], + ) + + +@pytest.mark.asyncio +async def test_stream_messages_with_system_prompt_content(agenerator, alist): + """Test stream_messages with SystemContentBlock input.""" + mock_model = unittest.mock.MagicMock() + mock_model.stream.return_value = agenerator( + [ + {"contentBlockDelta": {"delta": {"text": "test"}}}, + {"contentBlockStop": {}}, + ] + ) + + system_prompt_content = [{"text": "You are a helpful assistant."}, {"cachePoint": {"type": "default"}}] + + stream = strands.event_loop.streaming.stream_messages( + mock_model, + system_prompt_content=system_prompt_content, + messages=[{"role": "user", "content": [{"text": "Hello"}]}], + tool_specs=[], + system_prompt=None, + ) + + await alist(stream) + + # Verify model.stream was called with both parameters + mock_model.stream.assert_called_with( + [{"role": "user", "content": [{"text": "Hello"}]}], + None, + None, + tool_choice=None, + system_prompt_content=system_prompt_content, + ) + + +@pytest.mark.asyncio +async def test_stream_messages_single_text_block_backwards_compatibility(agenerator, alist): + """Test that single text block extracts system_prompt for backwards compatibility.""" + mock_model = unittest.mock.MagicMock() + mock_model.stream.return_value = agenerator( + [ + {"contentBlockDelta": {"delta": {"text": "test"}}}, + {"contentBlockStop": {}}, + ] + ) + + system_prompt_content = [{"text": "You are a helpful assistant."}] + + stream = strands.event_loop.streaming.stream_messages( + mock_model, + system_prompt_content=system_prompt_content, + messages=[{"role": "user", "content": [{"text": "Hello"}]}], + tool_specs=[], + system_prompt="You are a helpful assistant.", + ) + + await alist(stream) + + # Verify model.stream was called with extracted system_prompt for backwards compatibility + mock_model.stream.assert_called_with( + [{"role": "user", "content": [{"text": "Hello"}]}], + None, + "You are a helpful assistant.", + tool_choice=None, + system_prompt_content=system_prompt_content, + ) + + +@pytest.mark.asyncio +async def test_stream_messages_empty_system_prompt_content(agenerator, alist): + """Test stream_messages with empty system_prompt_content.""" + mock_model = unittest.mock.MagicMock() + mock_model.stream.return_value = agenerator( + [ + {"contentBlockDelta": {"delta": {"text": "test"}}}, + {"contentBlockStop": {}}, + ] + ) + + stream = strands.event_loop.streaming.stream_messages( + mock_model, + messages=[{"role": "user", "content": [{"text": "Hello"}]}], + tool_specs=[], + system_prompt=None, + system_prompt_content=[], + ) + + await alist(stream) + + # Verify model.stream was called with None system_prompt + mock_model.stream.assert_called_with( + [{"role": "user", "content": [{"text": "Hello"}]}], + None, + None, + tool_choice=None, + system_prompt_content=[], + ) + + +@pytest.mark.asyncio +async def test_stream_messages_none_system_prompt_content(agenerator, alist): + """Test stream_messages with None system_prompt_content.""" + mock_model = unittest.mock.MagicMock() + mock_model.stream.return_value = agenerator( + [ + {"contentBlockDelta": {"delta": {"text": "test"}}}, + {"contentBlockStop": {}}, + ] + ) + + stream = strands.event_loop.streaming.stream_messages( + mock_model, + system_prompt_content=None, + messages=[{"role": "user", "content": [{"text": "Hello"}]}], + tool_specs=None, + system_prompt=None, + ) + + tru_events = await alist(stream) + + # Verify model.stream was called with None system_prompt and empty lists + mock_model.stream.assert_called_with( + [{"role": "user", "content": [{"text": "Hello"}]}], + None, + None, + tool_choice=None, + system_prompt_content=None, ) # Ensure that we're getting typed events coming out of process_stream @@ -875,9 +1005,10 @@ async def test_stream_messages_normalizes_messages(agenerator, alist): await alist( strands.event_loop.streaming.stream_messages( mock_model, - system_prompt="test prompt", + system_prompt_content=[{"text": "test prompt"}], messages=messages, tool_specs=None, + system_prompt="test prompt", ) ) diff --git a/tests/strands/event_loop/test_streaming_structured_output.py b/tests/strands/event_loop/test_streaming_structured_output.py index e170445276..4645e1724a 100644 --- a/tests/strands/event_loop/test_streaming_structured_output.py +++ b/tests/strands/event_loop/test_streaming_structured_output.py @@ -50,9 +50,10 @@ async def test_stream_messages_with_tool_choice(agenerator, alist): stream = strands.event_loop.streaming.stream_messages( mock_model, - system_prompt="test prompt", + system_prompt_content=[{"text": "test prompt"}], messages=[{"role": "user", "content": [{"text": "Generate a test model"}]}], tool_specs=[tool_spec], + system_prompt="test prompt", tool_choice=tool_choice, ) @@ -64,6 +65,7 @@ async def test_stream_messages_with_tool_choice(agenerator, alist): [tool_spec], "test prompt", tool_choice=tool_choice, + system_prompt_content=[{"text": "test prompt"}], ) # Verify we get the expected events @@ -113,9 +115,10 @@ async def test_stream_messages_with_forced_structured_output(agenerator, alist): stream = strands.event_loop.streaming.stream_messages( mock_model, - system_prompt="Extract user information", + system_prompt_content=[{"text": "Extract user information"}], messages=[{"role": "user", "content": [{"text": "Alice is 30 years old"}]}], tool_specs=[tool_spec], + system_prompt="Extract user information", tool_choice=tool_choice, ) @@ -127,6 +130,7 @@ async def test_stream_messages_with_forced_structured_output(agenerator, alist): [tool_spec], "Extract user information", tool_choice=tool_choice, + system_prompt_content=[{"text": "Extract user information"}], ) assert len(tru_events) > 0 diff --git a/tests/strands/experimental/hooks/multiagent/test_multi_agent_hooks.py b/tests/strands/experimental/hooks/multiagent/test_multi_agent_hooks.py new file mode 100644 index 0000000000..4e97a92177 --- /dev/null +++ b/tests/strands/experimental/hooks/multiagent/test_multi_agent_hooks.py @@ -0,0 +1,130 @@ +import pytest + +from strands import Agent +from strands.experimental.hooks.multiagent.events import ( + AfterMultiAgentInvocationEvent, + AfterNodeCallEvent, + BeforeMultiAgentInvocationEvent, + BeforeNodeCallEvent, + MultiAgentInitializedEvent, +) +from strands.multiagent.graph import Graph, GraphBuilder +from strands.multiagent.swarm import Swarm +from tests.fixtures.mock_multiagent_hook_provider import MockMultiAgentHookProvider +from tests.fixtures.mocked_model_provider import MockedModelProvider + + +@pytest.fixture +def hook_provider(): + return MockMultiAgentHookProvider( + [ + BeforeMultiAgentInvocationEvent, + AfterMultiAgentInvocationEvent, + AfterNodeCallEvent, + BeforeNodeCallEvent, + MultiAgentInitializedEvent, + ] + ) + + +@pytest.fixture +def mock_model(): + agent_messages = [ + {"role": "assistant", "content": [{"text": "Task completed"}]}, + {"role": "assistant", "content": [{"text": "Task completed by agent 2"}]}, + {"role": "assistant", "content": [{"text": "Additional response"}]}, + ] + return MockedModelProvider(agent_messages) + + +@pytest.fixture +def agent1(mock_model): + return Agent(model=mock_model, system_prompt="You are agent 1.", name="agent1") + + +@pytest.fixture +def agent2(mock_model): + return Agent(model=mock_model, system_prompt="You are agent 2.", name="agent2") + + +@pytest.fixture +def swarm(agent1, agent2, hook_provider): + swarm = Swarm(nodes=[agent1, agent2], hooks=[hook_provider]) + return swarm + + +@pytest.fixture +def graph(agent1, agent2, hook_provider): + builder = GraphBuilder() + builder.add_node(agent1, "agent1") + builder.add_node(agent2, "agent2") + builder.add_edge("agent1", "agent2") + builder.set_entry_point("agent1") + graph = Graph(nodes=builder.nodes, edges=builder.edges, entry_points=builder.entry_points, hooks=[hook_provider]) + return graph + + +def test_swarm_complete_hook_lifecycle(swarm, hook_provider): + """E2E test verifying complete hook lifecycle for Swarm.""" + result = swarm("test task") + + length, events = hook_provider.get_events() + assert length == 5 + assert result.status.value == "completed" + + events_list = list(events) + + # Check event types and basic properties, ignoring invocation_state + assert isinstance(events_list[0], MultiAgentInitializedEvent) + assert events_list[0].source == swarm + + assert isinstance(events_list[1], BeforeMultiAgentInvocationEvent) + assert events_list[1].source == swarm + + assert isinstance(events_list[2], BeforeNodeCallEvent) + assert events_list[2].source == swarm + assert events_list[2].node_id == "agent1" + + assert isinstance(events_list[3], AfterNodeCallEvent) + assert events_list[3].source == swarm + assert events_list[3].node_id == "agent1" + + assert isinstance(events_list[4], AfterMultiAgentInvocationEvent) + assert events_list[4].source == swarm + + +def test_graph_complete_hook_lifecycle(graph, hook_provider): + """E2E test verifying complete hook lifecycle for Graph.""" + result = graph("test task") + + length, events = hook_provider.get_events() + assert length == 7 + assert result.status.value == "completed" + + events_list = list(events) + + # Check event types and basic properties, ignoring invocation_state + assert isinstance(events_list[0], MultiAgentInitializedEvent) + assert events_list[0].source == graph + + assert isinstance(events_list[1], BeforeMultiAgentInvocationEvent) + assert events_list[1].source == graph + + assert isinstance(events_list[2], BeforeNodeCallEvent) + assert events_list[2].source == graph + assert events_list[2].node_id == "agent1" + + assert isinstance(events_list[3], AfterNodeCallEvent) + assert events_list[3].source == graph + assert events_list[3].node_id == "agent1" + + assert isinstance(events_list[4], BeforeNodeCallEvent) + assert events_list[4].source == graph + assert events_list[4].node_id == "agent2" + + assert isinstance(events_list[5], AfterNodeCallEvent) + assert events_list[5].source == graph + assert events_list[5].node_id == "agent2" + + assert isinstance(events_list[6], AfterMultiAgentInvocationEvent) + assert events_list[6].source == graph diff --git a/tests/strands/models/test_bedrock.py b/tests/strands/models/test_bedrock.py index 0f68c8f179..2809e8a725 100644 --- a/tests/strands/models/test_bedrock.py +++ b/tests/strands/models/test_bedrock.py @@ -296,7 +296,7 @@ def test_update_config(model, model_id): def test_format_request_default(model, messages, model_id): - tru_request = model.format_request(messages) + tru_request = model._format_request(messages) exp_request = { "inferenceConfig": {}, "modelId": model_id, @@ -309,7 +309,7 @@ def test_format_request_default(model, messages, model_id): def test_format_request_additional_request_fields(model, messages, model_id, additional_request_fields): model.update_config(additional_request_fields=additional_request_fields) - tru_request = model.format_request(messages) + tru_request = model._format_request(messages) exp_request = { "additionalModelRequestFields": additional_request_fields, "inferenceConfig": {}, @@ -323,7 +323,7 @@ def test_format_request_additional_request_fields(model, messages, model_id, add def test_format_request_additional_response_field_paths(model, messages, model_id, additional_response_field_paths): model.update_config(additional_response_field_paths=additional_response_field_paths) - tru_request = model.format_request(messages) + tru_request = model._format_request(messages) exp_request = { "additionalModelResponseFieldPaths": additional_response_field_paths, "inferenceConfig": {}, @@ -337,7 +337,7 @@ def test_format_request_additional_response_field_paths(model, messages, model_i def test_format_request_guardrail_config(model, messages, model_id, guardrail_config): model.update_config(**guardrail_config) - tru_request = model.format_request(messages) + tru_request = model._format_request(messages) exp_request = { "guardrailConfig": { "guardrailIdentifier": guardrail_config["guardrail_id"], @@ -361,7 +361,7 @@ def test_format_request_guardrail_config_without_trace_or_stream_processing_mode "guardrail_version": "v1", } ) - tru_request = model.format_request(messages) + tru_request = model._format_request(messages) exp_request = { "guardrailConfig": { "guardrailIdentifier": "g1", @@ -379,7 +379,7 @@ def test_format_request_guardrail_config_without_trace_or_stream_processing_mode def test_format_request_inference_config(model, messages, model_id, inference_config): model.update_config(**inference_config) - tru_request = model.format_request(messages) + tru_request = model._format_request(messages) exp_request = { "inferenceConfig": { "maxTokens": inference_config["max_tokens"], @@ -396,7 +396,7 @@ def test_format_request_inference_config(model, messages, model_id, inference_co def test_format_request_system_prompt(model, messages, model_id, system_prompt): - tru_request = model.format_request(messages, system_prompt=system_prompt) + tru_request = model._format_request(messages, system_prompt_content=[{"text": system_prompt}]) exp_request = { "inferenceConfig": {}, "modelId": model_id, @@ -407,8 +407,54 @@ def test_format_request_system_prompt(model, messages, model_id, system_prompt): assert tru_request == exp_request +def test_format_request_system_prompt_content(model, messages, model_id): + """Test _format_request with SystemContentBlock input.""" + system_prompt_content = [{"text": "You are a helpful assistant."}, {"cachePoint": {"type": "default"}}] + + tru_request = model._format_request(messages, system_prompt_content=system_prompt_content) + exp_request = { + "inferenceConfig": {}, + "modelId": model_id, + "messages": messages, + "system": system_prompt_content, + } + + assert tru_request == exp_request + + +def test_format_request_system_prompt_content_with_cache_prompt_config(model, messages, model_id): + """Test _format_request with SystemContentBlock and cache_prompt config (backwards compatibility).""" + system_prompt_content = [{"text": "You are a helpful assistant."}] + model.update_config(cache_prompt="default") + + with pytest.warns(UserWarning, match="cache_prompt is deprecated"): + tru_request = model._format_request(messages, system_prompt_content=system_prompt_content) + + exp_request = { + "inferenceConfig": {}, + "modelId": model_id, + "messages": messages, + "system": [{"text": "You are a helpful assistant."}, {"cachePoint": {"type": "default"}}], + } + + assert tru_request == exp_request + + +def test_format_request_empty_system_prompt_content(model, messages, model_id): + """Test _format_request with empty SystemContentBlock list.""" + tru_request = model._format_request(messages, system_prompt_content=[]) + exp_request = { + "inferenceConfig": {}, + "modelId": model_id, + "messages": messages, + "system": [], + } + + assert tru_request == exp_request + + def test_format_request_tool_specs(model, messages, model_id, tool_spec): - tru_request = model.format_request(messages, [tool_spec]) + tru_request = model._format_request(messages, tool_specs=[tool_spec]) exp_request = { "inferenceConfig": {}, "modelId": model_id, @@ -425,7 +471,7 @@ def test_format_request_tool_specs(model, messages, model_id, tool_spec): def test_format_request_tool_choice_auto(model, messages, model_id, tool_spec): tool_choice = {"auto": {}} - tru_request = model.format_request(messages, [tool_spec], tool_choice=tool_choice) + tru_request = model._format_request(messages, [tool_spec], tool_choice=tool_choice) exp_request = { "inferenceConfig": {}, "modelId": model_id, @@ -442,7 +488,7 @@ def test_format_request_tool_choice_auto(model, messages, model_id, tool_spec): def test_format_request_tool_choice_any(model, messages, model_id, tool_spec): tool_choice = {"any": {}} - tru_request = model.format_request(messages, [tool_spec], tool_choice=tool_choice) + tru_request = model._format_request(messages, [tool_spec], tool_choice=tool_choice) exp_request = { "inferenceConfig": {}, "modelId": model_id, @@ -459,7 +505,7 @@ def test_format_request_tool_choice_any(model, messages, model_id, tool_spec): def test_format_request_tool_choice_tool(model, messages, model_id, tool_spec): tool_choice = {"tool": {"name": "test_tool"}} - tru_request = model.format_request(messages, [tool_spec], tool_choice=tool_choice) + tru_request = model._format_request(messages, [tool_spec], tool_choice=tool_choice) exp_request = { "inferenceConfig": {}, "modelId": model_id, @@ -476,7 +522,10 @@ def test_format_request_tool_choice_tool(model, messages, model_id, tool_spec): def test_format_request_cache(model, messages, model_id, tool_spec, cache_type): model.update_config(cache_prompt=cache_type, cache_tools=cache_type) - tru_request = model.format_request(messages, [tool_spec]) + + with pytest.warns(UserWarning, match="cache_prompt is deprecated"): + tru_request = model._format_request(messages, tool_specs=[tool_spec]) + exp_request = { "inferenceConfig": {}, "modelId": model_id, @@ -609,6 +658,51 @@ async def test_stream(bedrock_client, model, messages, tool_spec, model_id, addi bedrock_client.converse_stream.assert_called_once_with(**request) +@pytest.mark.asyncio +async def test_stream_with_system_prompt_content(bedrock_client, model, messages, alist): + """Test stream method with system_prompt_content parameter.""" + bedrock_client.converse_stream.return_value = {"stream": ["e1", "e2"]} + + system_prompt_content = [{"text": "You are a helpful assistant."}, {"cachePoint": {"type": "default"}}] + + response = model.stream(messages, system_prompt_content=system_prompt_content) + tru_chunks = await alist(response) + exp_chunks = ["e1", "e2"] + + assert tru_chunks == exp_chunks + + # Verify the request was formatted with system_prompt_content + expected_request = { + "inferenceConfig": {}, + "modelId": "m1", + "messages": messages, + "system": system_prompt_content, + } + bedrock_client.converse_stream.assert_called_once_with(**expected_request) + + +@pytest.mark.asyncio +async def test_stream_backwards_compatibility_single_text_block(bedrock_client, model, messages, alist): + """Test that single text block in system_prompt_content works with legacy system_prompt.""" + bedrock_client.converse_stream.return_value = {"stream": ["e1", "e2"]} + + system_prompt_content = [{"text": "You are a helpful assistant."}] + + response = model.stream( + messages, system_prompt="You are a helpful assistant.", system_prompt_content=system_prompt_content + ) + await alist(response) + + # Verify the request was formatted with system_prompt_content + expected_request = { + "inferenceConfig": {}, + "modelId": "m1", + "messages": messages, + "system": system_prompt_content, + } + bedrock_client.converse_stream.assert_called_once_with(**expected_request) + + @pytest.mark.asyncio async def test_stream_stream_input_guardrails( bedrock_client, model, messages, tool_spec, model_id, additional_request_fields, alist @@ -1510,7 +1604,7 @@ def test_format_request_cleans_tool_result_content_blocks(model, model_id): } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) tool_result = formatted_request["messages"][0]["content"][0]["toolResult"] expected = {"toolUseId": "tool123", "content": [{"text": "Tool output"}]} @@ -1538,7 +1632,7 @@ def test_format_request_removes_status_field_when_configured(model, model_id): } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) tool_result = formatted_request["messages"][0]["content"][0]["toolResult"] expected = {"toolUseId": "tool123", "content": [{"text": "Tool output"}]} @@ -1579,7 +1673,7 @@ def test_explicit_boolean_values_preserved(bedrock_client): } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) # Verify toolResult contains status field by default tool_result = formatted_request["messages"][0]["content"][0]["toolResult"] @@ -1601,7 +1695,7 @@ def test_format_request_filters_sdk_unknown_member_content_blocks(model, model_i } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) content = formatted_request["messages"][0]["content"] assert len(content) == 2 @@ -1683,7 +1777,7 @@ def test_format_request_filters_image_content_blocks(model, model_id): } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) image_block = formatted_request["messages"][0]["content"][0]["image"] expected = {"format": "png", "source": {"bytes": b"image_data"}} @@ -1711,7 +1805,7 @@ def test_format_request_filters_nested_image_s3_fields(model, model_id): } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) image_source = formatted_request["messages"][0]["content"][0]["image"]["source"] assert image_source == {"bytes": b"image_data"} @@ -1737,7 +1831,7 @@ def test_format_request_filters_document_content_blocks(model, model_id): } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) document_block = formatted_request["messages"][0]["content"][0]["document"] expected = {"name": "test.pdf", "source": {"bytes": b"pdf_data"}, "format": "pdf"} @@ -1761,7 +1855,7 @@ def test_format_request_filters_nested_reasoning_content(model, model_id): } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) reasoning_text = formatted_request["messages"][0]["content"][0]["reasoningContent"]["reasoningText"] assert reasoning_text == {"text": "thinking...", "signature": "abc123"} @@ -1785,7 +1879,7 @@ def test_format_request_filters_video_content_blocks(model, model_id): } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) video_block = formatted_request["messages"][0]["content"][0]["video"] expected = {"format": "mp4", "source": {"bytes": b"video_data"}} @@ -1810,7 +1904,7 @@ def test_format_request_filters_cache_point_content_blocks(model, model_id): } ] - formatted_request = model.format_request(messages) + formatted_request = model._format_request(messages) cache_point_block = formatted_request["messages"][0]["content"][0]["cachePoint"] expected = {"type": "default"} @@ -1839,14 +1933,14 @@ def test_update_config_validation_warns_on_unknown_keys(model, captured_warnings def test_tool_choice_supported_no_warning(model, messages, tool_spec, captured_warnings): """Test that toolChoice doesn't emit warning for supported providers.""" tool_choice = {"auto": {}} - model.format_request(messages, [tool_spec], tool_choice=tool_choice) + model._format_request(messages, [tool_spec], tool_choice=tool_choice) assert len(captured_warnings) == 0 def test_tool_choice_none_no_warning(model, messages, captured_warnings): """Test that None toolChoice doesn't emit warning.""" - model.format_request(messages, tool_choice=None) + model._format_request(messages, tool_choice=None) assert len(captured_warnings) == 0 @@ -1945,7 +2039,7 @@ def test_format_request_filters_output_schema(model, messages, model_id): "outputSchema": {"type": "object", "properties": {"result": {"type": "string"}}}, } - request = model.format_request(messages, [tool_spec_with_output_schema]) + request = model._format_request(messages, [tool_spec_with_output_schema]) tool_spec = request["toolConfig"]["tools"][0]["toolSpec"] @@ -1956,3 +2050,23 @@ def test_format_request_filters_output_schema(model, messages, model_id): assert tool_spec["name"] == "test_tool" assert tool_spec["description"] == "Test tool with output schema" assert tool_spec["inputSchema"] == {"type": "object", "properties": {}} + + +@pytest.mark.asyncio +async def test_stream_backward_compatibility_system_prompt(bedrock_client, model, messages, alist): + """Test that system_prompt is converted to system_prompt_content when system_prompt_content is None.""" + bedrock_client.converse_stream.return_value = {"stream": ["e1", "e2"]} + + system_prompt = "You are a helpful assistant." + + response = model.stream(messages, system_prompt=system_prompt) + await alist(response) + + # Verify the request was formatted with system_prompt converted to system_prompt_content + expected_request = { + "inferenceConfig": {}, + "modelId": "m1", + "messages": messages, + "system": [{"text": system_prompt}], + } + bedrock_client.converse_stream.assert_called_once_with(**expected_request) diff --git a/tests/strands/models/test_gemini.py b/tests/strands/models/test_gemini.py index 9eb5a9a7f3..a8f5351cc8 100644 --- a/tests/strands/models/test_gemini.py +++ b/tests/strands/models/test_gemini.py @@ -1,4 +1,5 @@ import json +import logging import unittest.mock import pydantic @@ -621,3 +622,18 @@ async def test_structured_output(gemini_client, model, messages, model_id, weath "model": model_id, } gemini_client.aio.models.generate_content.assert_called_with(**exp_request) + + +@pytest.mark.asyncio +async def test_stream_handles_non_json_error(gemini_client, model, messages, caplog, alist): + error_message = "Invalid API key" + gemini_client.aio.models.generate_content_stream.side_effect = genai.errors.ClientError( + error_message, {"message": error_message} + ) + + with caplog.at_level(logging.WARNING): + with pytest.raises(genai.errors.ClientError, match=error_message): + await alist(model.stream(messages)) + + assert "Gemini API returned non-JSON error" in caplog.text + assert f"error_message=<{error_message}>" in caplog.text diff --git a/tests/strands/multiagent/test_graph.py b/tests/strands/multiagent/test_graph.py index 07037a447a..b32356cb4d 100644 --- a/tests/strands/multiagent/test_graph.py +++ b/tests/strands/multiagent/test_graph.py @@ -10,6 +10,7 @@ from strands.hooks.registry import HookProvider, HookRegistry from strands.multiagent.base import MultiAgentBase, MultiAgentResult, NodeResult from strands.multiagent.graph import Graph, GraphBuilder, GraphEdge, GraphNode, GraphResult, GraphState, Status +from strands.session.file_session_manager import FileSessionManager from strands.session.session_manager import SessionManager @@ -1979,3 +1980,56 @@ async def stream_without_result(*args, **kwargs): mock_strands_tracer.start_multiagent_span.assert_called() mock_use_span.assert_called_once() + + +@pytest.mark.asyncio +async def test_graph_persisted(mock_strands_tracer, mock_use_span): + """Test graph persistence functionality.""" + # Create mock session manager + session_manager = Mock(spec=FileSessionManager) + session_manager.read_multi_agent().return_value = None + + # Create simple graph with session manager + builder = GraphBuilder() + agent = create_mock_agent("test_agent") + builder.add_node(agent, "test_node") + builder.set_entry_point("test_node") + builder.set_session_manager(session_manager) + + graph = builder.build() + + # Test get_state_from_orchestrator + state = graph.serialize_state() + assert state["type"] == "graph" + assert state["id"] == "default_graph" + assert "status" in state + assert "completed_nodes" in state + assert "node_results" in state + + # Test apply_state_from_dict with persisted state + persisted_state = { + "status": "executing", + "completed_nodes": [], + "failed_nodes": [], + "node_results": {}, + "current_task": "persisted task", + "execution_order": [], + "next_nodes_to_execute": ["test_node"], + } + + graph.deserialize_state(persisted_state) + assert graph.state.task == "persisted task" + + # Execute graph to test persistence integration + result = await graph.invoke_async("Test persistence") + + # Verify execution completed + assert result.status == Status.COMPLETED + assert len(result.results) == 1 + assert "test_node" in result.results + + # Test state serialization after execution + final_state = graph.serialize_state() + assert final_state["status"] == "completed" + assert len(final_state["completed_nodes"]) == 1 + assert "test_node" in final_state["node_results"] diff --git a/tests/strands/multiagent/test_swarm.py b/tests/strands/multiagent/test_swarm.py index 14a0ac1d69..e8a6a5f79a 100644 --- a/tests/strands/multiagent/test_swarm.py +++ b/tests/strands/multiagent/test_swarm.py @@ -9,6 +9,7 @@ from strands.hooks.registry import HookRegistry from strands.multiagent.base import Status from strands.multiagent.swarm import SharedContext, Swarm, SwarmNode, SwarmResult, SwarmState +from strands.session.file_session_manager import FileSessionManager from strands.session.session_manager import SessionManager from strands.types._events import MultiAgentNodeStartEvent from strands.types.content import ContentBlock @@ -1098,3 +1099,53 @@ async def failing_execute_swarm(*args, **kwargs): # Verify the swarm status is FAILED assert swarm.state.completion_status == Status.FAILED + + +@pytest.mark.asyncio +async def test_swarm_persistence(mock_strands_tracer, mock_use_span): + """Test swarm persistence functionality.""" + # Create mock session manager + session_manager = Mock(spec=FileSessionManager) + session_manager.read_multi_agent.return_value = None + + # Create simple swarm with session manager + agent = create_mock_agent("test_agent") + swarm = Swarm([agent], session_manager=session_manager) + + # Test get_state_from_orchestrator + state = swarm.serialize_state() + assert state["type"] == "swarm" + assert state["id"] == "default_swarm" + assert "status" in state + assert "node_history" in state + assert "node_results" in state + assert "context" in state + + # Test apply_state_from_dict with persisted state + persisted_state = { + "status": "executing", + "node_history": [], + "node_results": {}, + "current_task": "persisted task", + "next_nodes_to_execute": ["test_agent"], + "context": {"shared_context": {"test_agent": {"key": "value"}}, "handoff_message": "test handoff"}, + } + + swarm._from_dict(persisted_state) + assert swarm.state.task == "persisted task" + assert swarm.state.handoff_message == "test handoff" + assert swarm.shared_context.context["test_agent"]["key"] == "value" + + # Execute swarm to test persistence integration + result = await swarm.invoke_async("Test persistence") + + # Verify execution completed + assert result.status == Status.COMPLETED + assert len(result.results) == 1 + assert "test_agent" in result.results + + # Test state serialization after execution + final_state = swarm.serialize_state() + assert final_state["status"] == "completed" + assert len(final_state["node_history"]) == 1 + assert "test_agent" in final_state["node_results"] diff --git a/tests/strands/session/test_repository_session_manager.py b/tests/strands/session/test_repository_session_manager.py index e346f01e08..ed0ec9072e 100644 --- a/tests/strands/session/test_repository_session_manager.py +++ b/tests/strands/session/test_repository_session_manager.py @@ -233,3 +233,183 @@ def test_initialize_multi_agent_existing(session_manager, mock_multi_agent): # Verify deserialize_state was called with existing state mock_multi_agent.deserialize_state.assert_called_once_with(existing_state) + + +def test_fix_broken_tool_use_adds_missing_tool_results(session_manager): + """Test that _fix_broken_tool_use adds missing toolResult messages.""" + conversation_manager = SlidingWindowConversationManager() + + # Create agent in repository first + session_agent = SessionAgent( + agent_id="existing-agent", + state={"key": "value"}, + conversation_manager_state=conversation_manager.get_state(), + ) + session_manager.session_repository.create_agent("test-session", session_agent) + + broken_messages = [ + { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": "orphaned-123", "name": "test_tool", "input": {"input": "test"}}}], + }, + {"role": "user", "content": [{"text": "Some other message"}]}, + ] + # Create some session messages + for index, broken_message in enumerate(broken_messages): + broken_session_message = SessionMessage( + message=broken_message, + message_id=index, + ) + session_manager.session_repository.create_message("test-session", "existing-agent", broken_session_message) + + # Initialize agent + agent = Agent(agent_id="existing-agent") + session_manager.initialize(agent) + + fixed_messages = agent.messages + + # Should insert toolResult message between toolUse and other message + assert len(fixed_messages) == 3 + assert "toolResult" in fixed_messages[1]["content"][0] + assert fixed_messages[1]["content"][0]["toolResult"]["toolUseId"] == "orphaned-123" + assert fixed_messages[1]["content"][0]["toolResult"]["status"] == "error" + assert fixed_messages[1]["content"][0]["toolResult"]["content"][0]["text"] == "Tool was interrupted." + + +def test_fix_broken_tool_use_extends_partial_tool_results(session_manager): + """Test fixing messages where some toolResults are missing.""" + conversation_manager = SlidingWindowConversationManager() + # Create agent in repository first + session_agent = SessionAgent( + agent_id="existing-agent", + state={"key": "value"}, + conversation_manager_state=conversation_manager.get_state(), + ) + session_manager.session_repository.create_agent("test-session", session_agent) + + broken_messages = [ + { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "complete-123", "name": "test_tool", "input": {"input": "test1"}}}, + {"toolUse": {"toolUseId": "missing-456", "name": "test_tool", "input": {"input": "test2"}}}, + ], + }, + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "complete-123", "status": "success", "content": [{"text": "result"}]}} + ], + }, + ] + # Create some session messages + for index, broken_message in enumerate(broken_messages): + broken_session_message = SessionMessage( + message=broken_message, + message_id=index, + ) + session_manager.session_repository.create_message("test-session", "existing-agent", broken_session_message) + + # Initialize agent + agent = Agent(agent_id="existing-agent") + session_manager.initialize(agent) + + fixed_messages = agent.messages + + # Should add missing toolResult to existing message + assert len(fixed_messages) == 2 + assert len(fixed_messages[1]["content"]) == 2 + + tool_use_ids = {tr["toolResult"]["toolUseId"] for tr in fixed_messages[1]["content"]} + assert tool_use_ids == {"complete-123", "missing-456"} + + # Check the added toolResult has correct properties + missing_result = next(tr for tr in fixed_messages[1]["content"] if tr["toolResult"]["toolUseId"] == "missing-456") + assert missing_result["toolResult"]["status"] == "error" + assert missing_result["toolResult"]["content"][0]["text"] == "Tool was interrupted." + + +def test_fix_broken_tool_use_handles_multiple_orphaned_tools(session_manager): + """Test fixing multiple orphaned toolUse messages.""" + + conversation_manager = SlidingWindowConversationManager() + # Create agent in repository first + session_agent = SessionAgent( + agent_id="existing-agent", + state={"key": "value"}, + conversation_manager_state=conversation_manager.get_state(), + ) + session_manager.session_repository.create_agent("test-session", session_agent) + + broken_messages = [ + { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "orphaned-123", "name": "test_tool", "input": {"input": "test1"}}}, + {"toolUse": {"toolUseId": "orphaned-456", "name": "test_tool", "input": {"input": "test2"}}}, + ], + }, + {"role": "user", "content": [{"text": "Next message"}]}, + ] + # Create some session messages + for index, broken_message in enumerate(broken_messages): + broken_session_message = SessionMessage( + message=broken_message, + message_id=index, + ) + session_manager.session_repository.create_message("test-session", "existing-agent", broken_session_message) + + # Initialize agent + agent = Agent(agent_id="existing-agent") + session_manager.initialize(agent) + + fixed_messages = agent.messages + + # Should insert message with both toolResults + assert len(fixed_messages) == 3 + assert len(fixed_messages[1]["content"]) == 2 + + tool_use_ids = {tr["toolResult"]["toolUseId"] for tr in fixed_messages[1]["content"]} + assert tool_use_ids == {"orphaned-123", "orphaned-456"} + + +def test_fix_broken_tool_use_ignores_last_message(session_manager): + """Test that orphaned toolUse in the last message is not fixed.""" + messages = [ + {"role": "user", "content": [{"text": "Hello"}]}, + { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "last-message-123", "name": "test_tool", "input": {"input": "test"}}} + ], + }, + ] + + fixed_messages = session_manager._fix_broken_tool_use(messages) + + # Should remain unchanged since toolUse is in last message + assert fixed_messages == messages + + +def test_fix_broken_tool_use_does_not_change_valid_message(session_manager): + """Test that orphaned toolUse in the last message is not fixed.""" + messages = [ + {"role": "user", "content": [{"text": "Hello"}]}, + { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "last-message-123", "name": "test_tool", "input": {"input": "test"}}} + ], + }, + { + "role": "user", + "content": [ + {"toolResult": {"toolUseId": "last-message-123", "input": {"input": "test"}, "status": "success"}} + ], + }, + ] + + fixed_messages = session_manager._fix_broken_tool_use(messages) + + # Should remain unchanged since toolUse is in last message + assert fixed_messages == messages diff --git a/tests_integ/models/test_model_bedrock.py b/tests_integ/models/test_model_bedrock.py index 9dff66fde4..2c2e125ad9 100644 --- a/tests_integ/models/test_model_bedrock.py +++ b/tests_integ/models/test_model_bedrock.py @@ -267,3 +267,16 @@ def test_redacted_content_handling(): assert "reasoningContent" in result.message["content"][0] assert "redactedContent" in result.message["content"][0]["reasoningContent"] assert isinstance(result.message["content"][0]["reasoningContent"]["redactedContent"], bytes) + + +def test_multi_prompt_system_content(): + """Test multi-prompt system content blocks.""" + system_prompt_content = [ + {"text": "You are a helpful assistant."}, + {"text": "Always be concise."}, + {"text": "End responses with 'Done.'"}, + ] + + agent = Agent(system_prompt=system_prompt_content, load_tools_from_directory=False) + # just verifying there is no failure + agent("Hello!") diff --git a/tests_integ/models/test_model_openai.py b/tests_integ/models/test_model_openai.py index 115a0819d5..7beb3013cd 100644 --- a/tests_integ/models/test_model_openai.py +++ b/tests_integ/models/test_model_openai.py @@ -221,3 +221,13 @@ def test_rate_limit_throttling_integration_no_retries(model): # Verify it's a rate limit error error_message = str(exc_info.value).lower() assert "rate limit" in error_message or "tokens per min" in error_message + + +def test_content_blocks_handling(model): + """Test that content blocks are handled properly without failures.""" + content = [{"text": "What is 2+2?"}, {"text": "Please be brief."}] + + agent = Agent(model=model, load_tools_from_directory=False) + result = agent(content) + + assert "4" in result.message["content"][0]["text"] diff --git a/tests_integ/test_bedrock_cache_point.py b/tests_integ/test_bedrock_cache_point.py index 82bca22a26..5299146bb1 100644 --- a/tests_integ/test_bedrock_cache_point.py +++ b/tests_integ/test_bedrock_cache_point.py @@ -1,4 +1,5 @@ from strands import Agent +from strands.models import BedrockModel from strands.types.content import Messages @@ -29,3 +30,31 @@ def cache_point_callback_handler(**kwargs): agent = Agent(messages=messages, callback_handler=cache_point_callback_handler, load_tools_from_directory=False) agent("What is favorite color?") assert cache_point_usage > 0 + + +def test_bedrock_multi_prompt_and_duplicate_cache_point(): + """Test multi-prompt system with cache point.""" + system_prompt_content = [ + {"text": "You are a helpful assistant." * 500}, # Long text for cache + {"cachePoint": {"type": "default"}}, + {"text": "Always respond with enthusiasm!"}, + ] + + cache_point_usage = 0 + + def cache_point_callback_handler(**kwargs): + nonlocal cache_point_usage + if "event" in kwargs and kwargs["event"] and "metadata" in kwargs["event"] and kwargs["event"]["metadata"]: + metadata = kwargs["event"]["metadata"] + if "usage" in metadata and metadata["usage"]: + if "cacheReadInputTokens" in metadata["usage"] or "cacheWriteInputTokens" in metadata["usage"]: + cache_point_usage += 1 + + agent = Agent( + model=BedrockModel(model_id="us.anthropic.claude-sonnet-4-20250514-v1:0", cache_prompt="default"), + system_prompt=system_prompt_content, + callback_handler=cache_point_callback_handler, + load_tools_from_directory=False, + ) + agent("Hello!") + assert cache_point_usage > 0 diff --git a/tests_integ/test_multiagent_graph.py b/tests_integ/test_multiagent_graph.py index a7335feb79..08343a5546 100644 --- a/tests_integ/test_multiagent_graph.py +++ b/tests_integ/test_multiagent_graph.py @@ -1,4 +1,6 @@ from typing import Any, AsyncIterator +from unittest.mock import patch +from uuid import uuid4 import pytest @@ -13,6 +15,7 @@ ) from strands.multiagent.base import MultiAgentBase, MultiAgentResult, NodeResult, Status from strands.multiagent.graph import GraphBuilder +from strands.session.file_session_manager import FileSessionManager from strands.types.content import ContentBlock from tests.fixtures.mock_hook_provider import MockHookProvider @@ -458,3 +461,127 @@ async def test_graph_metrics_accumulation(): # Verify accumulated metrics are sum of node metrics total_tokens = sum(node_result.accumulated_usage["totalTokens"] for node_result in result.results.values()) assert result.accumulated_usage["totalTokens"] == total_tokens, "Accumulated tokens don't match sum of node tokens" + + +@pytest.mark.asyncio +async def test_graph_interrupt_and_resume(): + """Test graph interruption and resume functionality with FileSessionManager.""" + + session_id = str(uuid4()) + + # Create real agents + agent1 = Agent(model="us.amazon.nova-pro-v1:0", system_prompt="You are agent 1", name="agent1") + agent2 = Agent(model="us.amazon.nova-pro-v1:0", system_prompt="You are agent 2", name="agent2") + agent3 = Agent(model="us.amazon.nova-pro-v1:0", system_prompt="You are agent 3", name="agent3") + + session_manager = FileSessionManager(session_id=session_id) + + builder = GraphBuilder() + builder.add_node(agent1, "node1") + builder.add_node(agent2, "node2") + builder.add_node(agent3, "node3") + builder.add_edge("node1", "node2") + builder.add_edge("node2", "node3") + builder.set_entry_point("node1") + builder.set_session_manager(session_manager) + + graph = builder.build() + + # Mock agent2 to fail on first execution + async def failing_stream_async(*args, **kwargs): + raise Exception("Simulated failure in agent2") + yield # This line is never reached, but makes it an async generator + + with patch.object(agent2, "stream_async", side_effect=failing_stream_async): + try: + await graph.invoke_async("This is a test task, just do it shortly") + raise AssertionError("Expected exception was not raised") + except Exception as e: + assert "Simulated failure in agent2" in str(e) + + # Verify partial execution was persisted + persisted_state = session_manager.read_multi_agent(session_id, graph.id) + assert persisted_state is not None + assert persisted_state["type"] == "graph" + assert persisted_state["status"] == "failed" + assert len(persisted_state["completed_nodes"]) == 1 # Only node1 completed + assert "node1" in persisted_state["completed_nodes"] + assert "node2" in persisted_state["next_nodes_to_execute"] + assert "node2" in persisted_state["failed_nodes"] + + # Track execution count before resume + initial_execution_count = graph.state.execution_count + + # Execute graph again + result = await graph.invoke_async("Test task") + + # Verify successful completion + assert result.status == Status.COMPLETED + assert len(result.results) == 3 + + execution_order_ids = [node.node_id for node in result.execution_order] + assert execution_order_ids == ["node1", "node2", "node3"] + + # Verify only 2 additional nodes were executed + assert result.execution_count == initial_execution_count + 2 + + final_state = session_manager.read_multi_agent(session_id, graph.id) + assert final_state["status"] == "completed" + assert len(final_state["completed_nodes"]) == 3 + + # Clean up + session_manager.delete_session(session_id) + + +@pytest.mark.asyncio +async def test_self_loop_resume_from_persisted_state(tmp_path): + """Test resuming self-loop from persisted state where next node is itself.""" + + session_id = f"self_loop_resume_{uuid4()}" + session_manager = FileSessionManager(session_id=session_id, storage_dir=str(tmp_path)) + + counter_agent = Agent( + model="us.amazon.nova-pro-v1:0", + system_prompt="You are a counter. Just respond with 'Count: 1', 'Count: 2', Stop at 5.", + ) + + def should_continue_loop(state): + loop_executions = len([node for node in state.execution_order if node.node_id == "loop_node"]) + return loop_executions < 5 + + builder = GraphBuilder() + builder.add_node(counter_agent, "loop_node") + builder.add_edge("loop_node", "loop_node", condition=should_continue_loop) + builder.set_entry_point("loop_node") + builder.set_session_manager(session_manager) + builder.reset_on_revisit(True) + + graph = builder.build() + + call_count = 0 + original_stream = counter_agent.stream_async + + async def failing_after_two(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count <= 2: + async for event in original_stream(*args, **kwargs): + yield event + else: + raise Exception("Simulated failure after two executions") + + with patch.object(counter_agent, "stream_async", side_effect=failing_after_two): + try: + await graph.invoke_async("Count till 5") + except Exception as e: + assert "Simulated failure after two executions" in str(e) + + persisted_state = session_manager.read_multi_agent(session_id, graph.id) + assert persisted_state["status"] == "failed" + assert "loop_node" in persisted_state.get("failed_nodes") + assert len(persisted_state.get("execution_order")) == 2 + + result = await graph.invoke_async("Continue counting to 5") + assert result.status == Status.COMPLETED + assert len(result.execution_order) == 5 + assert all(node.node_id == "loop_node" for node in result.execution_order) diff --git a/tests_integ/test_multiagent_swarm.py b/tests_integ/test_multiagent_swarm.py index ae9129fbb1..7710306196 100644 --- a/tests_integ/test_multiagent_swarm.py +++ b/tests_integ/test_multiagent_swarm.py @@ -1,3 +1,6 @@ +from unittest.mock import patch +from uuid import uuid4 + import pytest from strands import Agent, tool @@ -10,7 +13,9 @@ BeforeToolCallEvent, MessageAddedEvent, ) +from strands.multiagent.base import Status from strands.multiagent.swarm import Swarm +from strands.session.file_session_manager import FileSessionManager from strands.types.content import ContentBlock from tests.fixtures.mock_hook_provider import MockHookProvider @@ -319,3 +324,55 @@ async def test_swarm_get_agent_results_flattening(): assert len(agent_results) == 1 assert isinstance(agent_results[0], AgentResult) assert agent_results[0].message is not None + + +@pytest.mark.asyncio +async def test_swarm_interrupt_and_resume(researcher_agent, analyst_agent, writer_agent): + """Test swarm interruption after analyst_agent and resume functionality.""" + session_id = str(uuid4()) + + # Create session manager + session_manager = FileSessionManager(session_id=session_id) + + # Create swarm with session manager + swarm = Swarm([researcher_agent, analyst_agent, writer_agent], session_manager=session_manager) + + # Mock analyst_agent's _invoke method to fail + async def failing_invoke(*args, **kwargs): + raise Exception("Simulated failure in analyst") + yield # This line is never reached, but makes it an async generator + + with patch.object(analyst_agent, "stream_async", side_effect=failing_invoke): + # First execution - should fail at analyst + result = await swarm.invoke_async("Research AI trends and create a brief report") + try: + assert result.status == Status.FAILED + except Exception as e: + assert "Simulated failure in analyst" in str(e) + + # Verify partial execution was persisted + persisted_state = session_manager.read_multi_agent(session_id, swarm.id) + assert persisted_state is not None + assert persisted_state["type"] == "swarm" + assert persisted_state["status"] == "failed" + assert len(persisted_state["node_history"]) == 1 # At least researcher executed + + # Track execution count before resume + initial_execution_count = len(persisted_state["node_history"]) + + # Execute swarm again - should automatically resume from saved state + result = await swarm.invoke_async("Research AI trends and create a brief report") + + # Verify successful completion + assert result.status == Status.COMPLETED + assert len(result.results) > 0 + + assert len(result.node_history) >= initial_execution_count + 1 + + node_names = [node.node_id for node in result.node_history] + assert "researcher" in node_names + # Either analyst or writer (or both) should have executed to complete the task + assert "analyst" in node_names or "writer" in node_names + + # Clean up + session_manager.delete_session(session_id)