diff --git a/apps/coordinator/src/context_compaction.py b/apps/coordinator/src/context_compaction.py index 2de6699..e50778a 100644 --- a/apps/coordinator/src/context_compaction.py +++ b/apps/coordinator/src/context_compaction.py @@ -51,6 +51,41 @@ class CompactionResult: ) +@dataclass +class SessionRotation: + """Result of session rotation operation. + + Attributes: + old_agent_id: Identifier of the closed agent session + new_agent_id: Identifier of the newly spawned agent + agent_type: Type of agent (sonnet, haiku, opus, glm) + next_issue_number: Issue number transferred to new agent + context_before_tokens: Token count before rotation + context_before_percent: Usage percentage before rotation + success: Whether rotation succeeded + error_message: Error message if rotation failed + """ + + old_agent_id: str + new_agent_id: str + agent_type: str + next_issue_number: int + context_before_tokens: int + context_before_percent: float + success: bool + error_message: str = "" + + def __repr__(self) -> str: + """String representation.""" + status = "success" if self.success else "failed" + return ( + f"SessionRotation(old={self.old_agent_id!r}, " + f"new={self.new_agent_id!r}, " + f"issue=#{self.next_issue_number}, " + f"status={status})" + ) + + class ContextCompactor: """Handles context compaction to free agent memory. diff --git a/apps/coordinator/src/context_monitor.py b/apps/coordinator/src/context_monitor.py index 04c8835..9c58c28 100644 --- a/apps/coordinator/src/context_monitor.py +++ b/apps/coordinator/src/context_monitor.py @@ -6,7 +6,7 @@ from collections import defaultdict from collections.abc import Callable from typing import Any -from src.context_compaction import CompactionResult, ContextCompactor +from src.context_compaction import CompactionResult, ContextCompactor, SessionRotation from src.models import ContextAction, ContextUsage logger = logging.getLogger(__name__) @@ -164,3 +164,83 @@ class ContextMonitor: logger.error(f"Compaction failed for {agent_id}: {result.error_message}") return result + + async def trigger_rotation( + self, + agent_id: str, + agent_type: str, + next_issue_number: int, + ) -> SessionRotation: + """Trigger session rotation for an agent. + + Spawns fresh agent when context reaches 95% threshold. + + Rotation process: + 1. Get current context usage metrics + 2. Close current agent session + 3. Spawn new agent with same type + 4. Transfer next issue to new agent + 5. Log rotation event with metrics + + Args: + agent_id: Unique identifier for the current agent + agent_type: Type of agent (sonnet, haiku, opus, glm) + next_issue_number: Issue number to transfer to new agent + + Returns: + SessionRotation with rotation details and metrics + """ + logger.warning( + f"Triggering session rotation for agent {agent_id} " + f"(type: {agent_type}, next issue: #{next_issue_number})" + ) + + try: + # Get context usage before rotation + usage = await self.get_context_usage(agent_id) + context_before_tokens = usage.used_tokens + context_before_percent = usage.usage_percent + + logger.info( + f"Agent {agent_id} context before rotation: " + f"{context_before_tokens}/{usage.total_tokens} ({context_before_percent:.1f}%)" + ) + + # Close current session + await self.api_client.close_session(agent_id) + logger.info(f"Closed session for agent {agent_id}") + + # Spawn new agent with same type + spawn_response = await self.api_client.spawn_agent( + agent_type=agent_type, + issue_number=next_issue_number, + ) + new_agent_id = spawn_response["agent_id"] + + logger.info( + f"Session rotation successful: {agent_id} -> {new_agent_id} " + f"(issue #{next_issue_number})" + ) + + return SessionRotation( + old_agent_id=agent_id, + new_agent_id=new_agent_id, + agent_type=agent_type, + next_issue_number=next_issue_number, + context_before_tokens=context_before_tokens, + context_before_percent=context_before_percent, + success=True, + ) + + except Exception as e: + logger.error(f"Session rotation failed for agent {agent_id}: {e}") + return SessionRotation( + old_agent_id=agent_id, + new_agent_id="", + agent_type=agent_type, + next_issue_number=next_issue_number, + context_before_tokens=0, + context_before_percent=0.0, + success=False, + error_message=str(e), + ) diff --git a/apps/coordinator/tests/test_context_monitor.py b/apps/coordinator/tests/test_context_monitor.py index cacfe6f..3c1c263 100644 --- a/apps/coordinator/tests/test_context_monitor.py +++ b/apps/coordinator/tests/test_context_monitor.py @@ -431,6 +431,184 @@ class TestContextMonitor: assert result.success is False assert result.error_message == "API timeout during compaction" + @pytest.mark.asyncio + async def test_trigger_rotation_closes_current_session( + self, mock_claude_api: AsyncMock + ) -> None: + """Should close current agent session when rotation is triggered.""" + monitor = ContextMonitor(api_client=mock_claude_api, poll_interval=0.1) + + # Mock context usage at 96% (above ROTATE_THRESHOLD) + mock_claude_api.get_context_usage.return_value = { + "used_tokens": 192000, + "total_tokens": 200000, + } + + # Mock close_session API + mock_claude_api.close_session = AsyncMock() + + # Trigger rotation + result = await monitor.trigger_rotation( + agent_id="agent-1", + agent_type="sonnet", + next_issue_number=42, + ) + + # Verify session was closed + mock_claude_api.close_session.assert_called_once_with("agent-1") + assert result.success is True + + @pytest.mark.asyncio + async def test_trigger_rotation_spawns_new_agent( + self, mock_claude_api: AsyncMock + ) -> None: + """Should spawn new agent with same type during rotation.""" + monitor = ContextMonitor(api_client=mock_claude_api, poll_interval=0.1) + + # Mock context usage at 96% + mock_claude_api.get_context_usage.return_value = { + "used_tokens": 192000, + "total_tokens": 200000, + } + + # Mock API calls + mock_claude_api.close_session = AsyncMock() + mock_claude_api.spawn_agent = AsyncMock(return_value={"agent_id": "agent-2"}) + + # Trigger rotation + result = await monitor.trigger_rotation( + agent_id="agent-1", + agent_type="opus", + next_issue_number=99, + ) + + # Verify new agent was spawned with same type + mock_claude_api.spawn_agent.assert_called_once_with( + agent_type="opus", + issue_number=99, + ) + assert result.new_agent_id == "agent-2" + assert result.success is True + + @pytest.mark.asyncio + async def test_trigger_rotation_logs_metrics( + self, mock_claude_api: AsyncMock + ) -> None: + """Should log rotation with session IDs and context metrics.""" + monitor = ContextMonitor(api_client=mock_claude_api, poll_interval=0.1) + + # Mock context usage at 97% + mock_claude_api.get_context_usage.return_value = { + "used_tokens": 194000, + "total_tokens": 200000, + } + + # Mock API calls + mock_claude_api.close_session = AsyncMock() + mock_claude_api.spawn_agent = AsyncMock(return_value={"agent_id": "agent-2"}) + + # Trigger rotation + result = await monitor.trigger_rotation( + agent_id="agent-1", + agent_type="haiku", + next_issue_number=123, + ) + + # Verify result contains metrics + assert result.old_agent_id == "agent-1" + assert result.new_agent_id == "agent-2" + assert result.agent_type == "haiku" + assert result.next_issue_number == 123 + assert result.context_before_tokens == 194000 + assert result.context_before_percent == 97.0 + assert result.success is True + + @pytest.mark.asyncio + async def test_trigger_rotation_transfers_issue( + self, mock_claude_api: AsyncMock + ) -> None: + """Should transfer next issue to new agent during rotation.""" + monitor = ContextMonitor(api_client=mock_claude_api, poll_interval=0.1) + + # Mock context usage at 95% + mock_claude_api.get_context_usage.return_value = { + "used_tokens": 190000, + "total_tokens": 200000, + } + + # Mock API calls + mock_claude_api.close_session = AsyncMock() + mock_claude_api.spawn_agent = AsyncMock(return_value={"agent_id": "agent-5"}) + + # Trigger rotation + result = await monitor.trigger_rotation( + agent_id="agent-4", + agent_type="sonnet", + next_issue_number=77, + ) + + # Verify issue was transferred to new agent + assert result.next_issue_number == 77 + mock_claude_api.spawn_agent.assert_called_once_with( + agent_type="sonnet", + issue_number=77, + ) + + @pytest.mark.asyncio + async def test_trigger_rotation_handles_failure( + self, mock_claude_api: AsyncMock + ) -> None: + """Should handle rotation failure and return error details.""" + monitor = ContextMonitor(api_client=mock_claude_api, poll_interval=0.1) + + # Mock context usage + mock_claude_api.get_context_usage.return_value = { + "used_tokens": 190000, + "total_tokens": 200000, + } + + # Mock API failure + mock_claude_api.close_session = AsyncMock(side_effect=Exception("Session close failed")) + + # Trigger rotation + result = await monitor.trigger_rotation( + agent_id="agent-1", + agent_type="sonnet", + next_issue_number=42, + ) + + # Verify failure is reported + assert result.success is False + assert "Session close failed" in result.error_message + + @pytest.mark.asyncio + async def test_rotation_triggered_at_95_percent( + self, mock_claude_api: AsyncMock + ) -> None: + """Should trigger rotation when context reaches exactly 95%.""" + monitor = ContextMonitor(api_client=mock_claude_api, poll_interval=0.1) + + # Mock 95% usage (exactly at ROTATE_THRESHOLD) + mock_claude_api.get_context_usage.return_value = { + "used_tokens": 190000, + "total_tokens": 200000, + } + + # Mock API calls + mock_claude_api.close_session = AsyncMock() + mock_claude_api.spawn_agent = AsyncMock(return_value={"agent_id": "agent-2"}) + + # Trigger rotation + result = await monitor.trigger_rotation( + agent_id="agent-1", + agent_type="sonnet", + next_issue_number=1, + ) + + # Verify rotation was successful at exactly 95% + assert result.success is True + assert result.context_before_percent == 95.0 + class TestIssueMetadata: """Test IssueMetadata model."""