from typing import List, Optional from textual.app import ComposeResult from textual.screen import Screen from models.agent import Agent from widgets.agentmoveoperations import AgentMoveOperations from widgets.multiagentselector import MultiAgentSelector from widgets.resultsdisplay import ResultsDisplay class MoveAgentWorkflowScreen(Screen): """Screen that handles the agent movement workflow.""" def __init__(self, all_agents: Optional[List[Agent]]): super().__init__() self.all_agents = all_agents self.selected_agents = None def compose(self) -> ComposeResult: """Start with the multi-agent selector.""" yield MultiAgentSelector(self.all_agents) def on_multi_agent_selector_agents_selected( self, message: MultiAgentSelector.AgentsSelected ) -> None: """Handle selected agents - switch to operations screen.""" self.selected_agents = message.selected_agents # Remove the MultiAgentSelector selector = self.query_one(MultiAgentSelector) selector.remove() # Mount the AgentMoveOperations with the selected Agent objects self.mount(AgentMoveOperations(self.selected_agents)) def on_agent_move_operations_operation_complete( self, message: AgentMoveOperations.OperationComplete ) -> None: """Handle completion of move operation - transition to results screen.""" # Format successful results success_lines = [] for agent, result in message.successful: success_lines.append(f"✓ {agent.hostname}") # Format unsuccessful results failure_lines = [] for agent, error in message.unsuccessful: failure_lines.append(f"✗ {agent.hostname}: {error}") successful_text = "\n".join(success_lines) if success_lines else "(none)" unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)" # Remove the operations widget ops_widget = self.query_one(AgentMoveOperations) ops_widget.remove() # Mount the results display self.mount( ResultsDisplay(message.operation, successful_text, unsuccessful_text) )