from uuid import uuid4

import pytest

from strands import Agent, tool

from strands.experimental.hooks.multiagent import BeforeNodeCallEvent

from strands.hooks import (

AfterInvocationEvent,

AfterModelCallEvent,

AfterToolCallEvent,

BeforeInvocationEvent,

BeforeModelCallEvent,

BeforeToolCallEvent,

MessageAddedEvent,

)

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

@tool

def web_search(query: str) -> str:

"""Search the web for information."""

# Mock implementation

return f"Results for '{query}': 25% yearly growth assumption, reaching $1.81 trillion by 2030"

@tool

def calculate(expression: str) -> str:

"""Calculate the result of a mathematical expression."""

try:

return f"The result of {expression} is {eval(expression)}"

except Exception as e:

return f"Error calculating {expression}: {str(e)}"

@pytest.fixture

def hook_provider():

return MockHookProvider("all")

@pytest.fixture

def researcher_agent(hook_provider):

"""Create an agent specialized in research."""

return Agent(

name="researcher",

system_prompt=(

"You are a research specialist who excels at finding information. When you need to perform calculations or"

" format documents, hand off to the appropriate specialist."

),

hooks=[hook_provider],

tools=[web_search],

)

@pytest.fixture

def analyst_agent(hook_provider):

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

return Agent(

name="analyst",

system_prompt=(

"You are a data analyst who excels at calculations and numerical analysis. When you need"

" research or document formatting, hand off to the appropriate specialist."

),

hooks=[hook_provider],

tools=[calculate],

)

@pytest.fixture

def writer_agent(hook_provider):

"""Create an agent specialized in writing and formatting."""

return Agent(

name="writer",

hooks=[hook_provider],

system_prompt=(

"You are a professional writer who excels at formatting and presenting information. When you need research"

" or calculations, hand off to the appropriate specialist."

),

)

@pytest.fixture

def exit_hook():

class ExitHook:

def __init__(self):

self.should_exit = True

def register_hooks(self, registry):

registry.add_callback(BeforeNodeCallEvent, self.exit_before_analyst)

def exit_before_analyst(self, event):

if event.node_id == "analyst" and self.should_exit:

raise SystemExit("Controlled exit before analyst")

return ExitHook()

@pytest.fixture

def verify_hook():

class VerifyHook:

def __init__(self):

self.first_node = None

def register_hooks(self, registry):

registry.add_callback(BeforeNodeCallEvent, self.capture_first_node)

def capture_first_node(self, event):

if self.first_node is None:

self.first_node = event.node_id

return VerifyHook()

def test_swarm_execution_with_string(researcher_agent, analyst_agent, writer_agent, hook_provider):

"""Test swarm execution with string input."""

# Create the swarm

swarm = Swarm([researcher_agent, analyst_agent, writer_agent])

# Define a task that requires collaboration

task = (

"Research the current AI agent market trends, calculate the growth rate assuming 25% yearly growth, "

"and create a basic report"

)

# Execute the swarm

result = swarm(task)

# Verify results

assert result.status.value == "completed"

assert len(result.results) > 0

assert result.execution_time > 0

assert result.execution_count > 0

# Verify agent history - at least one agent should have been used

assert len(result.node_history) > 0

# Just ensure that hooks are emitted; actual content is not verified

researcher_hooks = hook_provider.extract_for(researcher_agent).event_types_received

assert BeforeInvocationEvent in researcher_hooks

assert MessageAddedEvent in researcher_hooks

assert BeforeModelCallEvent in researcher_hooks

assert BeforeToolCallEvent in researcher_hooks

assert AfterToolCallEvent in researcher_hooks

assert AfterModelCallEvent in researcher_hooks

assert AfterInvocationEvent in researcher_hooks

@pytest.mark.asyncio

async def test_swarm_execution_with_image(researcher_agent, analyst_agent, writer_agent, yellow_img):

"""Test swarm execution with image input."""

# Create the swarm

swarm = Swarm([researcher_agent, analyst_agent, writer_agent])

# Create content blocks with text and image

content_blocks: list[ContentBlock] = [

{"text": "Analyze this image and create a report about what you see:"},

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

]

# Execute the swarm with multi-modal input

result = await swarm.invoke_async(content_blocks)

# Verify results

assert result.status.value == "completed"

assert len(result.results) > 0

assert result.execution_time > 0

assert result.execution_count > 0

# Verify agent history - at least one agent should have been used

assert len(result.node_history) > 0

@pytest.mark.asyncio

async def test_swarm_streaming(alist):

"""Test that Swarm properly streams all event types during execution."""

researcher = Agent(

name="researcher",

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

system_prompt="You are a researcher. When you need calculations, hand off to the analyst.",

)

analyst = Agent(

name="analyst",

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

system_prompt="You are an analyst. Use tools to perform calculations.",

tools=[calculate],

)

swarm = Swarm([researcher, analyst], node_timeout=900.0)

# Collect events

events = await alist(swarm.stream_async("Calculate 10 + 5 and explain 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) >= 1, f"Expected at least 1 node_start event, 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) >= 1, f"Expected at least 1 node_stop event, 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 "message" in handoff, "Handoff event missing message"

assert handoff["from_node_ids"] == ["researcher"], (

f"Expected from_node_ids=['researcher'], got {handoff['from_node_ids']}"

)

assert handoff["to_node_ids"] == ["analyst"], f"Expected to_node_ids=['analyst'], got {handoff['to_node_ids']}"

# Verify node stop event structure

stop_event = node_stop_events[0]

assert "node_id" in stop_event, "Node stop event missing node_id"

assert "node_result" in stop_event, "Node stop event missing node_result"

node_result = stop_event["node_result"]

assert hasattr(node_result, "execution_time"), "NodeResult missing execution_time"

assert node_result.execution_time > 0, "Expected positive execution_time"

# Verify we have events from at least one agent

researcher_events = [e for e in events if e.get("node_id") == "researcher"]

analyst_events = [e for e in events if e.get("node_id") == "analyst"]

assert len(researcher_events) > 0 or len(analyst_events) > 0, "Expected events from at least one agent"

@pytest.mark.asyncio

async def test_swarm_node_result_structure():

"""Test that NodeResult properly contains AgentResult after swarm execution.

This test verifies the merge conflict resolution where AgentResult import

was correctly handled and NodeResult properly wraps AgentResult objects.

"""

from strands.agent.agent_result import AgentResult

from strands.multiagent.base import NodeResult

researcher = Agent(

name="researcher",

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

system_prompt="You are a researcher. Answer the question directly without handing off.",

)

swarm = Swarm([researcher])

# Execute the swarm

result = await swarm.invoke_async("What is 2 + 2?")

# Verify the result structure

assert result.status.value in ["completed", "failed"] # May fail due to credentials

# If execution succeeded, verify the structure

if result.status.value == "completed":

assert len(result.results) == 1

assert "researcher" in result.results

# Verify NodeResult contains AgentResult

node_result = result.results["researcher"]

assert isinstance(node_result, NodeResult)

assert isinstance(node_result.result, AgentResult)

# Verify AgentResult has expected attributes

agent_result = node_result.result

assert hasattr(agent_result, "message")

assert hasattr(agent_result, "stop_reason")

assert hasattr(agent_result, "metrics")

assert agent_result.message is not None

assert agent_result.stop_reason in ["end_turn", "max_tokens", "stop_sequence"]

# Verify metrics are properly accumulated

assert node_result.accumulated_usage["totalTokens"] > 0

assert node_result.accumulated_metrics["latencyMs"] > 0

@pytest.mark.asyncio

async def test_swarm_multiple_handoffs_with_agent_results():

"""Test that multiple handoffs properly preserve AgentResult in each NodeResult.

This test ensures the AgentResult type is correctly used throughout the swarm

execution chain, verifying the import resolution from the merge conflict.

"""

from strands.agent.agent_result import AgentResult

agent1 = Agent(

name="agent1",

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

system_prompt="You are agent1. Hand off to agent2 immediately.",

)

agent2 = Agent(

name="agent2",

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

system_prompt="You are agent2. Hand off to agent3 immediately.",

)

agent3 = Agent(

name="agent3",

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

system_prompt="You are agent3. Complete the task without handing off.",

)

swarm = Swarm([agent1, agent2, agent3])

# Execute the swarm

result = await swarm.invoke_async("Complete this task")

# Verify execution completed or failed gracefully

assert result.status.value in ["completed", "failed"]

# If execution succeeded, verify the structure

if result.status.value == "completed":

assert len(result.node_history) >= 2 # At least 2 agents should have executed

# Verify each NodeResult contains a valid AgentResult

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

assert isinstance(node_result.result, AgentResult), f"Node {node_id} result is not an AgentResult"

assert node_result.result.message is not None, f"Node {node_id} AgentResult has no message"

assert node_result.accumulated_usage["totalTokens"] >= 0, f"Node {node_id} has invalid token usage"

@pytest.mark.asyncio

async def test_swarm_get_agent_results_flattening():

"""Test that get_agent_results() properly extracts AgentResult objects from NodeResults.

This test verifies that the NodeResult.get_agent_results() method correctly

handles AgentResult objects, ensuring the type system works correctly after

the merge conflict resolution.

"""

from strands.agent.agent_result import AgentResult

agent1 = Agent(

name="agent1",

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

system_prompt="You are agent1. Answer directly.",

)

swarm = Swarm([agent1])

# Execute the swarm

result = await swarm.invoke_async("What is the capital of France?")

# Verify execution completed or failed gracefully

assert result.status.value in ["completed", "failed"]

# If execution succeeded, verify the structure

if result.status.value == "completed":

assert "agent1" in result.results

node_result = result.results["agent1"]

# Test get_agent_results() method

agent_results = node_result.get_agent_results()

assert len(agent_results) == 1

assert isinstance(agent_results[0], AgentResult)

assert agent_results[0].message is not None

def test_swarm_resume_from_executing_state(tmpdir, exit_hook, verify_hook):

"""Test swarm resuming from EXECUTING state using BeforeNodeCallEvent hook."""

session_id = f"swarm_resume_{uuid4()}"

# First execution - exit before second node

session_manager = FileSessionManager(session_id=session_id, storage_dir=tmpdir)

researcher = Agent(name="researcher", system_prompt="you are a researcher.")

analyst = Agent(name="analyst", system_prompt="you are an analyst.")

writer = Agent(name="writer", system_prompt="you are a writer.")

swarm = Swarm([researcher, analyst, writer], session_manager=session_manager, hooks=[exit_hook])

try:

swarm("write AI trends and calculate growth in 100 words")

except SystemExit as e:

assert "Controlled exit before analyst" in str(e)

# Verify state was persisted with EXECUTING status and next node

persisted_state = session_manager.read_multi_agent(session_id, swarm.id)

assert persisted_state["status"] == "executing"

assert len(persisted_state["node_history"]) == 1

assert persisted_state["node_history"][0] == "researcher"

assert persisted_state["next_nodes_to_execute"] == ["analyst"]

exit_hook.should_exit = False

researcher2 = Agent(name="researcher", system_prompt="you are a researcher.")

analyst2 = Agent(name="analyst", system_prompt="you are an analyst.")

writer2 = Agent(name="writer", system_prompt="you are a writer.")

new_swarm = Swarm([researcher2, analyst2, writer2], session_manager=session_manager, hooks=[verify_hook])

result = new_swarm("write AI trends and calculate growth in 100 words")

# Verify swarm behavior - should resume from analyst, not restart

assert result.status.value == "completed"

assert verify_hook.first_node == "analyst"

node_ids = [n.node_id for n in result.node_history]

assert "analyst" in node_ids