#!/usr/bin/env python3 """ Tests for Sub-Agent Context Management Verifies: 1. Sub-agent context creation and retrieval 2. Phase progression tracking 3. Sibling agent discovery and coordination 4. Context persistence 5. Flow integration """ import pytest import tempfile from pathlib import Path import sys import os # Add parent directory to path for imports sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'lib')) from sub_agent_context import ( SubAgentContext, SubAgentContextManager, FlowPhase, ) from sub_agent_flow_integration import SubAgentFlowIntegrator class TestSubAgentContextCreation: """Test sub-agent context creation""" def setup_method(self): """Setup test fixtures""" self.temp_dir = tempfile.TemporaryDirectory() self.manager = SubAgentContextManager(Path(self.temp_dir.name)) def teardown_method(self): """Cleanup""" self.temp_dir.cleanup() def test_create_sub_agent_context(self): """Test creating a new sub-agent context""" context = self.manager.create_sub_agent_context( parent_task_id="task-123", parent_project="admin", parent_description="Test parent task", parent_context={"key": "value"}, parent_tags=["important", "research"], ) assert context.sub_agent_id is not None assert context.parent_task_id == "task-123" assert context.parent_project == "admin" assert context.parent_description == "Test parent task" assert len(context.phase_progression) == 9 assert context.phase_progression[0].phase_name == "CONTEXT_PREP" def test_phase_progression_initialization(self): """Test that all 9 phases are initialized""" context = self.manager.create_sub_agent_context( parent_task_id="task-456", parent_project="test", parent_description="Phase test", ) phase_names = [p.phase_name for p in context.phase_progression] expected_phases = [ "CONTEXT_PREP", "RECEIVED", "PREDICTING", "ANALYZING", "CONSENSUS_CHECK", "AWAITING_APPROVAL", "STRATEGIZING", "EXECUTING", "LEARNING", ] assert phase_names == expected_phases def test_retrieve_sub_agent_context(self): """Test retrieving sub-agent context""" created = self.manager.create_sub_agent_context( parent_task_id="task-789", parent_project="admin", parent_description="Retrieve test", ) retrieved = self.manager.get_sub_agent_context(created.sub_agent_id) assert retrieved is not None assert retrieved.sub_agent_id == created.sub_agent_id assert retrieved.parent_task_id == "task-789" class TestSiblingDiscovery: """Test sibling agent discovery and awareness""" def setup_method(self): """Setup test fixtures""" self.temp_dir = tempfile.TemporaryDirectory() self.manager = SubAgentContextManager(Path(self.temp_dir.name)) def teardown_method(self): """Cleanup""" self.temp_dir.cleanup() def test_single_sub_agent_no_siblings(self): """Test first sub-agent has no siblings""" context = self.manager.create_sub_agent_context( parent_task_id="parent-1", parent_project="admin", parent_description="First agent", ) assert len(context.sibling_agents) == 0 def test_multiple_sub_agents_discover_siblings(self): """Test multiple sub-agents discover each other as siblings""" # Create first sub-agent agent1 = self.manager.create_sub_agent_context( parent_task_id="parent-2", parent_project="admin", parent_description="Agent 1", ) # Create second sub-agent for same parent agent2 = self.manager.create_sub_agent_context( parent_task_id="parent-2", parent_project="admin", parent_description="Agent 2", ) # Create third sub-agent for same parent agent3 = self.manager.create_sub_agent_context( parent_task_id="parent-2", parent_project="admin", parent_description="Agent 3", ) # Verify sibling relationships assert agent2.sub_agent_id in self.manager.get_sibling_agents(agent1.sub_agent_id) assert agent3.sub_agent_id in self.manager.get_sibling_agents(agent1.sub_agent_id) assert len(self.manager.get_sibling_agents(agent1.sub_agent_id)) == 2 assert agent1.sub_agent_id in self.manager.get_sibling_agents(agent2.sub_agent_id) assert agent3.sub_agent_id in self.manager.get_sibling_agents(agent2.sub_agent_id) assert len(self.manager.get_sibling_agents(agent2.sub_agent_id)) == 2 def test_agents_from_different_parents_not_siblings(self): """Test agents from different parents are not siblings""" agent1 = self.manager.create_sub_agent_context( parent_task_id="parent-a", parent_project="admin", parent_description="Agent 1", ) agent2 = self.manager.create_sub_agent_context( parent_task_id="parent-b", parent_project="admin", parent_description="Agent 2", ) assert agent2.sub_agent_id not in self.manager.get_sibling_agents(agent1.sub_agent_id) assert agent1.sub_agent_id not in self.manager.get_sibling_agents(agent2.sub_agent_id) class TestPhaseProgression: """Test phase progression tracking""" def setup_method(self): """Setup test fixtures""" self.temp_dir = tempfile.TemporaryDirectory() self.manager = SubAgentContextManager(Path(self.temp_dir.name)) self.context = self.manager.create_sub_agent_context( parent_task_id="task-phase", parent_project="admin", parent_description="Phase test", ) def teardown_method(self): """Cleanup""" self.temp_dir.cleanup() def test_update_phase_status(self): """Test updating phase status""" success = self.manager.update_phase( self.context.sub_agent_id, "CONTEXT_PREP", "completed", output="Context prepared", ) assert success is True updated = self.manager.get_sub_agent_context(self.context.sub_agent_id) phase = updated.phase_progression[0] assert phase.status == "completed" assert phase.output == "Context prepared" def test_get_current_phase(self): """Test getting current active phase""" # Initially should be first pending phase current = self.manager.get_current_phase(self.context.sub_agent_id) assert current == "CONTEXT_PREP" # Mark first phase as complete self.manager.update_phase( self.context.sub_agent_id, "CONTEXT_PREP", "completed", ) # Now should be next pending phase current = self.manager.get_current_phase(self.context.sub_agent_id) assert current == "RECEIVED" def test_phase_duration_calculation(self): """Test duration calculation for completed phases""" # Mark phase as in progress self.manager.update_phase( self.context.sub_agent_id, "CONTEXT_PREP", "in_progress", ) # Mark as completed self.manager.update_phase( self.context.sub_agent_id, "CONTEXT_PREP", "completed", output="Done", ) updated = self.manager.get_sub_agent_context(self.context.sub_agent_id) phase = updated.phase_progression[0] assert phase.duration_seconds is not None assert phase.duration_seconds >= 0 def test_phase_progression_sequence(self): """Test progressing through all phases""" sub_agent_id = self.context.sub_agent_id phases = [p.phase_name for p in self.context.phase_progression] for phase_name in phases: self.manager.update_phase( sub_agent_id, phase_name, "completed", output=f"Completed {phase_name}", ) updated = self.manager.get_sub_agent_context(sub_agent_id) all_completed = all(p.status == "completed" for p in updated.phase_progression) assert all_completed is True class TestCoordination: """Test sub-agent coordination and messaging""" def setup_method(self): """Setup test fixtures""" self.temp_dir = tempfile.TemporaryDirectory() self.manager = SubAgentContextManager(Path(self.temp_dir.name)) # Create two sibling agents self.agent1 = self.manager.create_sub_agent_context( parent_task_id="parent-coord", parent_project="admin", parent_description="Agent 1", ) self.agent2 = self.manager.create_sub_agent_context( parent_task_id="parent-coord", parent_project="admin", parent_description="Agent 2", ) def teardown_method(self): """Cleanup""" self.temp_dir.cleanup() def test_send_message_to_sibling(self): """Test sending coordination message to sibling""" success = self.manager.send_message_to_sibling( self.agent1.sub_agent_id, self.agent2.sub_agent_id, "request", {"type": "need_data", "data_type": "context"}, ) assert success is True def test_message_appears_in_both_agents(self): """Test message is visible to both sender and receiver""" self.manager.send_message_to_sibling( self.agent1.sub_agent_id, self.agent2.sub_agent_id, "update", {"status": "ready"}, ) agent1_updated = self.manager.get_sub_agent_context(self.agent1.sub_agent_id) agent2_updated = self.manager.get_sub_agent_context(self.agent2.sub_agent_id) assert len(agent1_updated.coordination_messages) == 1 assert len(agent2_updated.coordination_messages) == 1 assert agent1_updated.coordination_messages[0]["type"] == "update" assert agent2_updated.coordination_messages[0]["type"] == "update" def test_cannot_message_non_sibling(self): """Test cannot send message to non-sibling agent""" # Create agent with different parent agent3 = self.manager.create_sub_agent_context( parent_task_id="parent-other", parent_project="admin", parent_description="Agent 3", ) # Try to send message across parent boundary success = self.manager.send_message_to_sibling( self.agent1.sub_agent_id, agent3.sub_agent_id, "request", {"data": "test"}, ) assert success is False class TestContextPersistence: """Test context persistence to disk""" def test_context_saved_and_loaded(self): """Test contexts are saved to disk and reloaded""" with tempfile.TemporaryDirectory() as temp_dir: manager1 = SubAgentContextManager(Path(temp_dir)) # Create context in first manager context1 = manager1.create_sub_agent_context( parent_task_id="task-persist", parent_project="admin", parent_description="Persistence test", ) sub_agent_id = context1.sub_agent_id # Create new manager pointing to same directory manager2 = SubAgentContextManager(Path(temp_dir)) # Should be able to retrieve context from new manager context2 = manager2.get_sub_agent_context(sub_agent_id) assert context2 is not None assert context2.parent_task_id == "task-persist" assert context2.sub_agent_id == sub_agent_id class TestFlowIntegration: """Test flow integration with sub-agent context""" def setup_method(self): """Setup test fixtures""" self.temp_dir = tempfile.TemporaryDirectory() self.context_manager = SubAgentContextManager(Path(self.temp_dir.name)) self.integrator = SubAgentFlowIntegrator(self.context_manager) def teardown_method(self): """Cleanup""" self.temp_dir.cleanup() def test_execute_sub_agent_flow(self): """Test executing full sub-agent flow""" results = self.integrator.execute_sub_agent_flow( parent_task_id="task-flow", parent_project="admin", parent_description="Flow test", parent_context={"key": "value"}, ) assert results["sub_agent_id"] is not None assert "phases" in results # Should have results for all 9 phases assert len(results["phases"]) == 9 def test_execute_single_phase(self): """Test executing a single phase""" context = self.context_manager.create_sub_agent_context( parent_task_id="task-single", parent_project="admin", parent_description="Single phase test", ) result = self.integrator.execute_phase(context.sub_agent_id, "CONTEXT_PREP") assert result["status"] == "completed" assert "output" in result def test_get_sub_agent_progress(self): """Test getting progress report""" context = self.context_manager.create_sub_agent_context( parent_task_id="task-progress", parent_project="admin", parent_description="Progress test", ) # Execute a phase self.integrator.execute_phase(context.sub_agent_id, "CONTEXT_PREP") self.integrator.execute_phase(context.sub_agent_id, "RECEIVED") progress = self.integrator.get_sub_agent_progress(context.sub_agent_id) assert progress["completed_phases"] == 2 assert progress["in_progress_phases"] == 0 assert progress["total_phases"] == 9 def test_coordinate_sequential_sub_agents(self): """Test sequential coordination of sub-agents""" # Create multiple sub-agents for same parent for i in range(3): self.context_manager.create_sub_agent_context( parent_task_id="task-coord", parent_project="admin", parent_description=f"Agent {i+1}", ) coordination = self.integrator.coordinate_sub_agents( parent_task_id="task-coord", coordination_strategy="sequential", ) assert len(coordination["sub_agents"]) == 3 assert coordination["strategy"] == "sequential" def test_collect_sub_agent_results(self): """Test collecting results from multiple sub-agents""" # Create and execute multiple sub-agents for i in range(2): context = self.context_manager.create_sub_agent_context( parent_task_id="task-collect", parent_project="admin", parent_description=f"Agent {i+1}", ) self.integrator.execute_phase(context.sub_agent_id, "CONTEXT_PREP") results = self.integrator.collect_sub_agent_results("task-collect") assert results["sub_agents_total"] == 2 assert len(results["sub_agents"]) == 2 assert all("progress" in s for s in results["sub_agents"]) class TestContextSummary: """Test context summary generation""" def setup_method(self): """Setup test fixtures""" self.temp_dir = tempfile.TemporaryDirectory() self.manager = SubAgentContextManager(Path(self.temp_dir.name)) def teardown_method(self): """Cleanup""" self.temp_dir.cleanup() def test_get_context_summary(self): """Test getting human-readable summary""" context = self.manager.create_sub_agent_context( parent_task_id="task-summary", parent_project="admin", parent_description="Summary test", parent_tags=["important", "urgent"], ) # Create a sibling self.manager.create_sub_agent_context( parent_task_id="task-summary", parent_project="admin", parent_description="Sibling agent", ) summary = self.manager.get_context_summary(context.sub_agent_id) assert summary is not None assert summary["sub_agent_id"] == context.sub_agent_id assert summary["parent_task_id"] == "task-summary" assert summary["sibling_count"] == 1 assert summary["parent_tags"] == ["important", "urgent"] if __name__ == "__main__": pytest.main([__file__, "-v"])