from typing import Any, AsyncIterator

from unittest.mock import patch

from uuid import uuid4

import pytest

from strands import Agent, tool

from strands.hooks import (

AfterInvocationEvent,

AfterModelCallEvent,

AgentInitializedEvent,

BeforeInvocationEvent,

BeforeModelCallEvent,

MessageAddedEvent,

)

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

@tool

def calculate_sum(a: int, b: int) -> int:

"""Calculate the sum of two numbers."""

return a + b

@tool

def multiply_numbers(x: int, y: int) -> int:

"""Multiply two numbers together."""

return x * y

@pytest.fixture

def hook_provider():

return MockHookProvider("all")

@pytest.fixture

def math_agent(hook_provider):

"""Create an agent specialized in mathematical operations."""

return Agent(

model="us.amazon.nova-pro-v1:0",

system_prompt="You are a mathematical assistant. Always provide clear, step-by-step calculations.",

hooks=[hook_provider],

tools=[calculate_sum, multiply_numbers],

)

@pytest.fixture

def analysis_agent(hook_provider):

"""Create an agent specialized in data analysis."""

return Agent(

model="us.amazon.nova-pro-v1:0",

hooks=[hook_provider],

system_prompt="You are a data analysis expert. Provide insights and interpretations of numerical results.",

)

@pytest.fixture

def summary_agent(hook_provider):

"""Create an agent specialized in summarization."""

return Agent(

model="us.amazon.nova-lite-v1:0",

hooks=[hook_provider],

system_prompt="You are a summarization expert. Create concise, clear summaries of complex information.",

)

@pytest.fixture

def validation_agent(hook_provider):

"""Create an agent specialized in validation."""

return Agent(

model="us.amazon.nova-pro-v1:0",

hooks=[hook_provider],

system_prompt="You are a validation expert. Check results for accuracy and completeness.",

)

@pytest.fixture

def image_analysis_agent(hook_provider):

"""Create an agent specialized in image analysis."""

return Agent(

hooks=[hook_provider],

system_prompt=(

"You are an image analysis expert. Describe what you see in images and provide detailed analysis."

),

)

@pytest.fixture

def nested_computation_graph(math_agent, analysis_agent):

"""Create a nested graph for mathematical computation and analysis."""

builder = GraphBuilder()

# Add agents to nested graph

builder.add_node(math_agent, "calculator")

builder.add_node(analysis_agent, "analyzer")

# Connect them sequentially

builder.add_edge("calculator", "analyzer")

builder.set_entry_point("calculator")

return builder.build()

@pytest.mark.asyncio

async def test_graph_execution_with_string(math_agent, summary_agent, validation_agent, nested_computation_graph):

# Define conditional functions

def should_validate(state):

"""Condition to determine if validation should run."""

return any(node.node_id == "computation_subgraph" for node in state.completed_nodes)

def proceed_to_second_summary(state):

"""Condition to skip additional summary."""

return False # Skip for this test

builder = GraphBuilder()

summary_agent_duplicate = Agent(

model="us.amazon.nova-lite-v1:0",

system_prompt="You are a summarization expert. Create concise, clear summaries of complex information.",

)

# Add various node types

builder.add_node(nested_computation_graph, "computation_subgraph") # Nested Graph node

builder.add_node(math_agent, "secondary_math") # Agent node

builder.add_node(validation_agent, "validator") # Agent node with condition

builder.add_node(summary_agent, "primary_summary") # Agent node

builder.add_node(summary_agent_duplicate, "secondary_summary") # Another Agent node

# Add edges with various configurations

builder.add_edge("computation_subgraph", "secondary_math") # Graph -> Agent

builder.add_edge("computation_subgraph", "validator", condition=should_validate) # Conditional edge

builder.add_edge("secondary_math", "primary_summary") # Agent -> Agent

builder.add_edge("validator", "primary_summary") # Agent -> Agent

builder.add_edge("primary_summary", "secondary_summary", condition=proceed_to_second_summary) # Conditional (false)

builder.set_entry_point("computation_subgraph")

graph = builder.build()

task = (

"Calculate 15 + 27 and 8 * 6, analyze both results, perform additional calculations, validate everything, "

"and provide a comprehensive summary"

)

result = await graph.invoke_async(task)

# Verify results

assert result.status.value == "completed"

assert result.total_nodes == 5

assert result.completed_nodes == 4 # All except secondary_summary (blocked by false condition)

assert result.failed_nodes == 0

assert len(result.results) == 4

# Verify execution order - extract node_ids from GraphNode objects

execution_order_ids = [node.node_id for node in result.execution_order]

# With parallel execution, secondary_math and validator can complete in any order

assert execution_order_ids[0] == "computation_subgraph" # First

assert execution_order_ids[3] == "primary_summary" # Last

assert set(execution_order_ids[1:3]) == {"secondary_math", "validator"} # Middle two in any order

# Verify specific nodes completed

assert "computation_subgraph" in result.results

assert "secondary_math" in result.results

assert "validator" in result.results

assert "primary_summary" in result.results

assert "secondary_summary" not in result.results # Should be blocked by condition

# Verify nested graph execution

nested_result = result.results["computation_subgraph"].result

assert nested_result.status.value == "completed"

@pytest.mark.asyncio

async def test_graph_execution_with_image(image_analysis_agent, summary_agent, yellow_img, hook_provider):

"""Test graph execution with multi-modal image input."""

builder = GraphBuilder()

# Add agents to graph

builder.add_node(image_analysis_agent, "image_analyzer")

builder.add_node(summary_agent, "summarizer")

# Connect them sequentially

builder.add_edge("image_analyzer", "summarizer")

builder.set_entry_point("image_analyzer")

graph = builder.build()

# Create content blocks with text and image

content_blocks: list[ContentBlock] = [

{"text": "Analyze this image and describe what you see:"},

{"image": {"format": "png", "source": {"bytes": yellow_img}}},

]

# Execute the graph with multi-modal input

result = await graph.invoke_async(content_blocks)

# Verify results

assert result.status.value == "completed"

assert result.total_nodes == 2

assert result.completed_nodes == 2

assert result.failed_nodes == 0

assert len(result.results) == 2

# Verify execution order

execution_order_ids = [node.node_id for node in result.execution_order]

assert execution_order_ids == ["image_analyzer", "summarizer"]

# Verify both nodes completed

assert "image_analyzer" in result.results

assert "summarizer" in result.results

expected_hook_events = [

AgentInitializedEvent,

BeforeInvocationEvent,

MessageAddedEvent,

BeforeModelCallEvent,

AfterModelCallEvent,

MessageAddedEvent,

AfterInvocationEvent,

]

assert hook_provider.extract_for(image_analysis_agent).event_types_received == expected_hook_events

assert hook_provider.extract_for(summary_agent).event_types_received == expected_hook_events

class CustomStreamingNode(MultiAgentBase):

"""Custom node that wraps an agent and adds custom streaming events."""

def __init__(self, agent: Agent, name: str):

self.agent = agent

self.name = name

async def invoke_async(

self, task: str | list[ContentBlock], invocation_state: dict[str, Any] | None = None, **kwargs: Any

) -> MultiAgentResult:

result = await self.agent.invoke_async(task, **kwargs)

node_result = NodeResult(result=result, status=Status.COMPLETED)

return MultiAgentResult(status=Status.COMPLETED, results={self.name: node_result})

async def stream_async(

self, task: str | list[ContentBlock], invocation_state: dict[str, Any] | None = None, **kwargs: Any

) -> AsyncIterator[dict[str, Any]]:

yield {"custom_event": "start", "node": self.name}

result = await self.agent.invoke_async(task, **kwargs)

yield {"custom_event": "agent_complete", "node": self.name}

node_result = NodeResult(result=result, status=Status.COMPLETED)

yield {"result": MultiAgentResult(status=Status.COMPLETED, results={self.name: node_result})}

@pytest.mark.asyncio

async def test_graph_streaming_with_agents(alist):

"""Test that Graph properly streams events from agent nodes."""

math_agent = Agent(

name="math",

model="us.amazon.nova-pro-v1:0",

system_prompt="You are a math assistant.",

tools=[calculate_sum],

)

summary_agent = Agent(

name="summary",

model="us.amazon.nova-lite-v1:0",

system_prompt="You are a summary assistant.",

)

builder = GraphBuilder()

builder.add_node(math_agent, "math")

builder.add_node(summary_agent, "summary")

builder.add_edge("math", "summary")

builder.set_entry_point("math")

builder.set_node_timeout(900.0) # Verify timeout doesn't interfere with streaming

graph = builder.build()

# Collect events

events = await alist(graph.stream_async("Calculate 5 + 3 and summarize the result"))

# Count event categories

node_start_events = [e for e in events if e.get("type") == "multiagent_node_start"]

node_stream_events = [e for e in events if e.get("type") == "multiagent_node_stream"]

node_stop_events = [e for e in events if e.get("type") == "multiagent_node_stop"]

handoff_events = [e for e in events if e.get("type") == "multiagent_handoff"]

result_events = [e for e in events if "result" in e and e.get("type") != "multiagent_node_stream"]

# Verify we got multiple events of each type

assert len(node_start_events) >= 2, f"Expected at least 2 node_start events, got {len(node_start_events)}"

assert len(node_stream_events) > 10, f"Expected many node_stream events, got {len(node_stream_events)}"

assert len(node_stop_events) >= 2, f"Expected at least 2 node_stop events, got {len(node_stop_events)}"

assert len(handoff_events) >= 1, f"Expected at least 1 handoff event, got {len(handoff_events)}"

assert len(result_events) >= 1, f"Expected at least 1 result event, got {len(result_events)}"

# Verify handoff event structure

handoff = handoff_events[0]

assert "from_node_ids" in handoff, "Handoff event missing from_node_ids"

assert "to_node_ids" in handoff, "Handoff event missing to_node_ids"

assert isinstance(handoff["from_node_ids"], list), "from_node_ids should be a list"

assert isinstance(handoff["to_node_ids"], list), "to_node_ids should be a list"

assert "math" in handoff["from_node_ids"], "Expected math in from_node_ids"

assert "summary" in handoff["to_node_ids"], "Expected summary in to_node_ids"

# Verify we have events for both nodes

math_events = [e for e in events if e.get("node_id") == "math"]

summary_events = [e for e in events if e.get("node_id") == "summary"]

assert len(math_events) > 0, "Expected events from math node"

assert len(summary_events) > 0, "Expected events from summary node"

@pytest.mark.asyncio

async def test_graph_streaming_with_custom_node(alist):

"""Test that Graph properly streams events from custom MultiAgentBase nodes."""

math_agent = Agent(

name="math",

model="us.amazon.nova-pro-v1:0",

system_prompt="You are a math assistant.",

tools=[calculate_sum],

)

summary_agent = Agent(

name="summary",

model="us.amazon.nova-lite-v1:0",

system_prompt="You are a summary assistant.",

)

# Create a custom node

custom_node = CustomStreamingNode(summary_agent, "custom_summary")

builder = GraphBuilder()

builder.add_node(math_agent, "math")

builder.add_node(custom_node, "custom_summary")

builder.add_edge("math", "custom_summary")

builder.set_entry_point("math")

graph = builder.build()

# Collect events

events = await alist(graph.stream_async("Calculate 5 + 3 and summarize the result"))

# Count event categories

node_start_events = [e for e in events if e.get("type") == "multiagent_node_start"]

node_stream_events = [e for e in events if e.get("type") == "multiagent_node_stream"]

result_events = [e for e in events if "result" in e and e.get("type") != "multiagent_node_stream"]

# Extract custom events from wrapped node_stream events

# Structure: {"type": "multiagent_node_stream", "node_id": "...", "event": {...}}

custom_events = []

for e in node_stream_events:

if e.get("type") == "multiagent_node_stream" and "event" in e:

inner_event = e["event"]

if isinstance(inner_event, dict) and "custom_event" in inner_event:

custom_events.append(inner_event)

# Verify we got multiple events of each type

assert len(node_start_events) >= 2, f"Expected at least 2 node_start events, got {len(node_start_events)}"

assert len(node_stream_events) > 5, f"Expected many node_stream events, got {len(node_stream_events)}"

assert len(custom_events) >= 2, f"Expected at least 2 custom events (start, complete), got {len(custom_events)}"

assert len(result_events) >= 1, f"Expected at least 1 result event, got {len(result_events)}"

# Verify custom events are properly structured

custom_start = [e for e in custom_events if e.get("custom_event") == "start"]

custom_complete = [e for e in custom_events if e.get("custom_event") == "agent_complete"]

assert len(custom_start) >= 1, "Expected at least 1 custom start event"

assert len(custom_complete) >= 1, "Expected at least 1 custom complete event"

@pytest.mark.asyncio

async def test_nested_graph_streaming(alist):

"""Test that nested graphs properly propagate streaming events."""

math_agent = Agent(

name="math",

model="us.amazon.nova-pro-v1:0",

system_prompt="You are a math assistant.",

tools=[calculate_sum],

)

analysis_agent = Agent(

name="analysis",

model="us.amazon.nova-lite-v1:0",

system_prompt="You are an analysis assistant.",

)

# Create nested graph

nested_builder = GraphBuilder()

nested_builder.add_node(math_agent, "calculator")

nested_builder.add_node(analysis_agent, "analyzer")

nested_builder.add_edge("calculator", "analyzer")

nested_builder.set_entry_point("calculator")

nested_graph = nested_builder.build()

# Create outer graph with nested graph

summary_agent = Agent(

name="summary",

model="us.amazon.nova-lite-v1:0",

system_prompt="You are a summary assistant.",

)

outer_builder = GraphBuilder()

outer_builder.add_node(nested_graph, "computation")

outer_builder.add_node(summary_agent, "summary")

outer_builder.add_edge("computation", "summary")

outer_builder.set_entry_point("computation")

outer_graph = outer_builder.build()

# Collect events

events = await alist(outer_graph.stream_async("Calculate 7 + 8 and provide a summary"))

# Count event categories

node_start_events = [e for e in events if e.get("type") == "multiagent_node_start"]

node_stream_events = [e for e in events if e.get("type") == "multiagent_node_stream"]

result_events = [e for e in events if "result" in e and e.get("type") != "multiagent_node_stream"]

# Verify we got multiple events

assert len(node_start_events) >= 2, f"Expected at least 2 node_start events, got {len(node_start_events)}"

assert len(node_stream_events) > 10, f"Expected many node_stream events, got {len(node_stream_events)}"

assert len(result_events) >= 1, f"Expected at least 1 result event, got {len(result_events)}"

# Verify we have events from nested nodes

computation_events = [e for e in events if e.get("node_id") == "computation"]

summary_events = [e for e in events if e.get("node_id") == "summary"]

assert len(computation_events) > 0, "Expected events from computation (nested graph) node"

assert len(summary_events) > 0, "Expected events from summary node"

@pytest.mark.asyncio

async def test_graph_metrics_accumulation():

"""Test that graph properly accumulates metrics from agent nodes."""

math_agent = Agent(

name="math",

model="us.amazon.nova-pro-v1:0",

system_prompt="You are a math assistant.",

tools=[calculate_sum],

)

summary_agent = Agent(

name="summary",

model="us.amazon.nova-lite-v1:0",

system_prompt="You are a summary assistant.",

)

builder = GraphBuilder()

builder.add_node(math_agent, "math")

builder.add_node(summary_agent, "summary")

builder.add_edge("math", "summary")

builder.set_entry_point("math")

graph = builder.build()

result = await graph.invoke_async("Calculate 5 + 3 and summarize the result")

# Verify result has accumulated metrics

assert result.accumulated_usage is not None

assert result.accumulated_usage["totalTokens"] > 0, "Expected non-zero total tokens"

assert result.accumulated_usage["inputTokens"] > 0, "Expected non-zero input tokens"

assert result.accumulated_usage["outputTokens"] > 0, "Expected non-zero output tokens"

assert result.accumulated_metrics is not None

assert result.accumulated_metrics["latencyMs"] > 0, "Expected non-zero latency"

# Verify individual node results have metrics

for node_id, node_result in result.results.items():

assert node_result.accumulated_usage is not None, f"Node {node_id} missing usage metrics"

assert node_result.accumulated_usage["totalTokens"] > 0, f"Node {node_id} has zero total tokens"

assert node_result.accumulated_metrics is not None, f"Node {node_id} missing metrics"

# 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)