feat(swarm): add AgentBase protocol support

Author: Emaan Khan

tests/strands/multiagent/test_swarm_agentbase.py

@@ -156,24 +156,6 @@ def test_swarm_node_state_management_with_agentbase():
     node.reset_executor_state()  # Should complete without error
 
 
-def test_swarm_node_state_management_with_agent():
-    """Test SwarmNode state management with Agent (hasattr should find attributes)."""
-    # Create Agent with messages/state attributes
-    agent = create_mock_agent("agent_with_state")
-
-    # Add actual messages and state
-    agent.messages = [{"role": "user", "content": [{"text": "Hello"}]}]
-    agent.state = Mock()
-    agent.state.get = Mock(return_value={"key": "value"})
-
-    # Create SwarmNode
-    node = SwarmNode(node_id="test_node", executor=agent)
-
-    # __post_init__ should capture state
-    assert len(node._initial_messages) == 1
-    assert node._initial_state.get()["key"] == "value"
-
-
 @pytest.mark.asyncio
 async def test_swarm_execution_with_agentbase():
     """Test Swarm execution with AgentBase implementations."""
@@ -195,27 +177,6 @@ async def test_swarm_execution_with_agentbase():
     assert agentbase._call_count >= 1
 
 
-@pytest.mark.asyncio
-async def test_swarm_execution_mixed_agents():
-    """Test Swarm execution with mixed Agent and AgentBase nodes."""
-    # Create mixed nodes
-    agent = create_mock_agent("regular_agent", "Agent response")
-    agentbase = MockAgentBase("agentbase_node", "AgentBase response")
-
-    # Create swarm with agent as entry point (it has tools for handoff)
-    swarm = Swarm(nodes=[agent, agentbase], entry_point=agent)
-
-    # Execute swarm
-    result = await swarm.invoke_async("Test mixed execution")
-
-    # Verify execution completed
-    assert result.status == Status.COMPLETED
-    assert len(result.results) >= 1
-
-    # At least the entry point should have executed
-    assert "regular_agent" in result.results
-
-
 def test_swarm_agentbase_without_name_attribute():
     """Test Swarm handles AgentBase without name attribute."""
 
@@ -247,96 +208,6 @@ async def stream_async(self, prompt: Any = None, **kwargs: Any):
     assert "node_0" in swarm.nodes
 
 
-def test_swarm_agentbase_entry_point():
-    """Test Swarm with AgentBase as entry point."""
-    agentbase1 = MockAgentBase("agentbase1")
-    agentbase2 = MockAgentBase("agentbase2")
-
-    # Set agentbase2 as entry point
-    swarm = Swarm(nodes=[agentbase1, agentbase2], entry_point=agentbase2)
-
-    # Verify entry point is set
-    assert swarm.entry_point is agentbase2
-
-    # Verify initial node uses entry point
-    initial_node = swarm._initial_node()
-    assert initial_node.node_id == "agentbase2"
-
-
-def test_swarm_agentbase_entry_point_without_name():
-    """Test entry point handling when AgentBase has no name attribute."""
-
-    class AgentBaseNoName:
-        """AgentBase without name attribute."""
-
-        async def invoke_async(self, prompt: Any = None, **kwargs: Any) -> AgentResult:
-            return AgentResult(
-                message={"role": "assistant", "content": [{"text": "response"}]},
-                stop_reason="end_turn",
-                state={},
-                metrics=None,
-            )
-
-        def __call__(self, prompt: Any = None, **kwargs: Any) -> AgentResult:
-            return asyncio.run(self.invoke_async(prompt, **kwargs))
-
-        async def stream_async(self, prompt: Any = None, **kwargs: Any):
-            result = await self.invoke_async(prompt, **kwargs)
-            yield {"result": result}
-
-    agentbase = AgentBaseNoName()
-
-    # Should create swarm and fallback to first node
-    swarm = Swarm(nodes=[agentbase])
-
-    # Should fallback to first node when entry point has no name
-    initial_node = swarm._initial_node()
-    assert initial_node.node_id == "node_0"
-
-
-def test_swarm_node_hasattr_pattern_for_messages():
-    """Test that SwarmNode uses hasattr for messages (Graph pattern)."""
-    # Create AgentBase without messages
-    agentbase = MockAgentBase("agentbase_node")
-
-    # Create SwarmNode
-    node = SwarmNode(node_id="test_node", executor=agentbase)
-
-    # Should not have captured messages (none exist)
-    assert node._initial_messages == []
-
-    # Add messages attribute dynamically
-    agentbase.messages = [{"role": "user", "content": [{"text": "Test"}]}]
-
-    # Create new node with messages
-    node_with_messages = SwarmNode(node_id="test_node2", executor=agentbase)
-
-    # Should capture messages via hasattr
-    assert len(node_with_messages._initial_messages) == 1
-
-
-def test_swarm_node_hasattr_pattern_for_state():
-    """Test that SwarmNode uses hasattr for state (Graph pattern)."""
-    # Create AgentBase without state
-    agentbase = MockAgentBase("agentbase_node")
-
-    # Create SwarmNode
-    node = SwarmNode(node_id="test_node", executor=agentbase)
-
-    # Should not have captured state (none exists)
-    assert node._initial_state.get() == {}
-
-    # Add state attribute dynamically
-    agentbase.state = Mock()
-    agentbase.state.get = Mock(return_value={"test": "data"})
-
-    # Create new node with state
-    node_with_state = SwarmNode(node_id="test_node2", executor=agentbase)
-
-    # Should capture state via hasattr
-    assert node_with_state._initial_state.get()["test"] == "data"
-
-
 def test_swarm_interrupt_handling_with_agentbase():
     """Test that interrupt handling only saves Agent-specific context."""
     from strands.interrupt import Interrupt
@@ -368,75 +239,3 @@ def test_swarm_interrupt_handling_with_agentbase():
     # Should NOT have saved AgentBase context (isinstance check should prevent it)
     # The interrupt is registered but no Agent-specific context is saved
     assert swarm._interrupt_state.activated
-
-
-@pytest.mark.asyncio
-async def test_swarm_agentbase_streaming():
-    """Test that AgentBase streaming works correctly."""
-    # Create AgentBase
-    agentbase = MockAgentBase("streaming_agentbase", "Streaming response")
-
-    # Create swarm
-    swarm = Swarm(nodes=[agentbase])
-
-    # Collect events
-    events = []
-    async for event in swarm.stream_async("Test streaming"):
-        events.append(event)
-
-    # Should have received events
-    assert len(events) > 0
-
-    # Should have node start/stop events
-    node_start_events = [e for e in events if e.get("type") == "multiagent_node_start"]
-    node_stop_events = [e for e in events if e.get("type") == "multiagent_node_stop"]
-
-    assert len(node_start_events) >= 1
-    assert len(node_stop_events) >= 1
-
-    # Should have final result
-    result_events = [e for e in events if "result" in e and e.get("type") != "multiagent_node_stream"]
-    assert len(result_events) == 1
-
-
-def test_swarm_duplicate_agentbase_instances():
-    """Test that Swarm rejects duplicate AgentBase instances."""
-    agentbase = MockAgentBase("agentbase_node")
-
-    # Try to create swarm with same instance twice
-    with pytest.raises(ValueError, match="Duplicate node instance detected"):
-        Swarm(nodes=[agentbase, agentbase])
-
-
-def test_swarm_agentbase_without_session_manager():
-    """Test that AgentBase nodes don't trigger session manager validation."""
-    # Create AgentBase (no _session_manager attribute)
-    agentbase = MockAgentBase("agentbase_node")
-
-    # Should not raise error (validation only applies to Agent instances)
-    swarm = Swarm(nodes=[agentbase])
-    assert len(swarm.nodes) == 1
-
-
-def test_swarm_node_reset_with_agentbase():
-    """Test SwarmNode reset_executor_state with AgentBase."""
-    # Create AgentBase
-    agentbase = MockAgentBase("agentbase_node")
-
-    # Add state attributes that can be reset
-    agentbase.messages = [{"role": "user", "content": [{"text": "Original"}]}]
-    agentbase.state = Mock()
-    agentbase.state.get = Mock(return_value={"original": "data"})
-
-    # Create node
-    node = SwarmNode(node_id="test_node", executor=agentbase)
-
-    # Modify state
-    agentbase.messages.append({"role": "assistant", "content": [{"text": "Modified"}]})
-
-    # Reset state
-    node.reset_executor_state()
-
-    # Should have reset to initial state
-    assert len(agentbase.messages) == 1
-    assert agentbase.messages[0]["content"][0]["text"] == "Original"