Quiet Agent UI improvements

TUI/Screens/quietagentworkflowscreen.py

@@ -34,7 +34,7 @@ from textual.app import ComposeResult
 from textual.containers import Horizontal, Vertical
 from textual.reactive import reactive
 from textual.screen import Screen
-from textual.widgets import Button, DataTable, Footer, Header, Static
+from textual.widgets import Button, DataTable, Footer, Header, Input, Static
 
 from models.policy import Policy
 from services.API import AirlockAPIWrapper
@@ -51,16 +51,17 @@ class QuietAgentWorkflowScreen(Screen):
 
     This screen provides a multi-step workflow:
     1. Select initial policy to analyze
-    2. View categorized agents (enforce ready vs. non-enforce ready)
+    2. Configure analysis parameters (history period and quiet time period)
-    3. Select target policies for each category
+    3. View categorized agents (enforce ready vs. non-enforce ready)
-    4. Execute agent migrations
+    4. Select target policies for each category
+    5. Execute agent migrations
 
     Attributes:
         api (AirlockAPIWrapper): API wrapper for Airlock operations
         policies (List[Policy]): List of all available policies
         selected_policy (Optional[Policy]): The initially selected policy to analyze
-        history_days (int): Number of days of history to pull (default: 150)
+        history_days (int): Number of days of history to pull (default: 150, range: 1-365)
-        quiet_days (int): Number of days without execution to be considered quiet (default: 45)
+        quiet_days (int): Number of days without execution to be considered quiet (default: 45, range: 1-365)
         agents_df (Optional[pd.DataFrame]): DataFrame of all agents with analysis results
         enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents ready for enforcement
         non_enforce_ready_df (Optional[pd.DataFrame]): DataFrame of agents not ready for enforcement
@@ -86,7 +87,7 @@ class QuietAgentWorkflowScreen(Screen):
         self.api = api
         self.policies = policies
         self.selected_policy: Optional[Policy] = None
-        self.history_days = 150  # Fixed as per requirements
+        self.history_days = 150  # Default value, user-selectable
         self.quiet_days = 45  # Default value
         self.agents_df: Optional[pd.DataFrame] = None
         self.enforce_ready_df: Optional[pd.DataFrame] = None
@@ -130,7 +131,7 @@ class QuietAgentWorkflowScreen(Screen):
 
         stage_messages = {
             "select_policy": "Step 1: Select Policy to Analyze",
-            "select_quiet_days": "Step 2: Select Quiet Time Period",
+            "select_history_days": "Step 2: Configure Analysis Parameters",
             "analyzing": "Analyzing agent activity...",
             "view_results": "Step 3: Review Categorized Agents",
             "select_enforce_target": "Step 4: Select Target Policy for Enforce Ready Agents",
@@ -161,7 +162,7 @@ class QuietAgentWorkflowScreen(Screen):
             # Initial policy selection for analysis
             self.selected_policy = message.policy
             logger.info(f"Selected policy for analysis: {self.selected_policy.name}")
-            self._show_quiet_days_selection()
+            self._show_history_days_selection()
         elif self.workflow_stage == "select_enforce_target":
             # Target policy selection for enforce ready agents
             self.enforce_ready_target_policy = message.policy
@@ -177,48 +178,167 @@ class QuietAgentWorkflowScreen(Screen):
             )
             self._show_migration_confirmation()
 
-    def _show_quiet_days_selection(self) -> None:
+    def _show_history_days_selection(self) -> None:
-        """Show the quiet days selection screen."""
+        """Show the history days and quiet days selection screen."""
-        self.workflow_stage = "select_quiet_days"
+        self.workflow_stage = "select_history_days"
         content = self.query_one("#content_area", Vertical)
         content.remove_children()
 
         # Create info text
         info_widget = Static(
             f"Policy Selected: {self.selected_policy.name}\n\n"
-            f"History Period: {self.history_days} days\n\n"
+            "Configure Analysis Parameters:",
-            "Select quiet time period (days without untrusted execution):",
+            id="analysis_params_info",
-            id="quiet_days_info",
         )
         info_widget.styles.margin = (0, 0, 2, 0)
         content.mount(info_widget)
 
-        # Create button container and mount it first
+        # Create input container
-        button_container = Vertical(id="quiet_days_buttons")
+        input_container = Vertical(id="analysis_params_input_container")
-        button_container.styles.height = "auto"
+        input_container.styles.height = "auto"
-        content.mount(button_container)
+        content.mount(input_container)
 
-        # Now add buttons to the mounted container
+        # History days label
-        for days in [15, 30, 45, 60]:
+        history_label = Static("History Period (days of execution history to pull):")
-            btn = Button(
+        history_label.styles.margin = (0, 0, 1, 0)
-                f"{days} days {'(Default)' if days == 45 else ''}",
+        input_container.mount(history_label)
-                id=f"quiet_days_{days}",
+
-                classes="quiet_day_btn",
+        # Add history days input field
+        history_input = Input(
+            placeholder="Enter days (1-365, default: 150)",
+            value="150",
+            id="history_days_input",
+        )
+        history_input.styles.width = "50"
+        history_input.styles.margin = (0, 0, 2, 0)
+        input_container.mount(history_input)
+
+        # Quiet days label
+        quiet_label = Static(
+            "Quiet Time Period (days without execution to be considered quiet):"
+        )
+        quiet_label.styles.margin = (0, 0, 1, 0)
+        input_container.mount(quiet_label)
+
+        # Add quiet days input field
+        quiet_input = Input(
+            placeholder="Enter days (1-365, default: 45)",
+            value="45",
+            id="quiet_days_input",
+        )
+        quiet_input.styles.width = "50"
+        quiet_input.styles.margin = (0, 0, 2, 0)
+        input_container.mount(quiet_input)
+
+        # Add submit button
+        submit_btn = Button(
+            "Continue",
+            id="analysis_params_submit",
+            variant="primary",
+        )
+        submit_btn.styles.width = "50"
+        submit_btn.styles.margin = (1, 0, 0, 0)
+        input_container.mount(submit_btn)
+
+        # Focus the first input field
+        history_input.focus()
+
+    def _validate_and_submit_history_days(self) -> None:
+        """Validate and submit the history days and quiet days inputs."""
+        try:
+            history_input = self.query_one("#history_days_input", Input)
+            quiet_input = self.query_one("#quiet_days_input", Input)
+
+            history_value = history_input.value.strip()
+            quiet_value = quiet_input.value.strip()
+
+            # Validate history days
+            if not history_value:
+                self.app.notify(
+                    "Please enter a history period value", severity="error", timeout=3
+                )
+                history_input.focus()
+                return
+
+            try:
+                history_days = int(history_value)
+            except ValueError:
+                self.app.notify(
+                    "Please enter a valid number for history period",
+                    severity="error",
+                    timeout=3,
+                )
+                history_input.focus()
+                return
+
+            if history_days < 1 or history_days > 365:
+                self.app.notify(
+                    "History period must be between 1 and 365 days",
+                    severity="error",
+                    timeout=3,
+                )
+                history_input.focus()
+                return
+
+            # Validate quiet days
+            if not quiet_value:
+                self.app.notify(
+                    "Please enter a quiet time period value",
+                    severity="error",
+                    timeout=3,
+                )
+                quiet_input.focus()
+                return
+
+            try:
+                quiet_days = int(quiet_value)
+            except ValueError:
+                self.app.notify(
+                    "Please enter a valid number for quiet time period",
+                    severity="error",
+                    timeout=3,
+                )
+                quiet_input.focus()
+                return
+
+            if quiet_days < 1 or quiet_days > 365:
+                self.app.notify(
+                    "Quiet time period must be between 1 and 365 days",
+                    severity="error",
+                    timeout=3,
+                )
+                quiet_input.focus()
+                return
+
+            # Check that quiet days doesn't exceed history days
+            if quiet_days > history_days:
+                self.app.notify(
+                    "Quiet time period cannot exceed history period",
+                    severity="error",
+                    timeout=3,
+                )
+                quiet_input.focus()
+                return
+
+            # All validation passed
+            self.history_days = history_days
+            self.quiet_days = quiet_days
+            logger.info(
+                f"Selected history days: {history_days}, quiet days: {quiet_days}"
             )
-            btn.styles.width = "100%"
+            self._start_analysis()
-            btn.styles.margin = (0, 0, 1, 0)
+
-            button_container.mount(btn)
+        except Exception as e:
+            logger.error(f"Error validating analysis parameters: {e}")
+            self.app.notify(f"Error: {str(e)}", severity="error", timeout=3)
 
     def on_button_pressed(self, event: Button.Pressed) -> None:
         """Handle button press events."""
         button_id = event.button.id
 
-        # Quiet days selection buttons
+        # Analysis parameters submit button
-        if button_id and button_id.startswith("quiet_days_"):
+        if button_id == "analysis_params_submit":
-            days = int(button_id.split("_")[-1])
+            self._validate_and_submit_history_days()
-            self.quiet_days = days
-            logger.info(f"Selected quiet days: {days}")
-            self._start_analysis()
             return
 
         # Navigation buttons
@@ -258,46 +378,44 @@ class QuietAgentWorkflowScreen(Screen):
             self._show_policy_selection()
             return
 
+    def on_input_submitted(self, event: Input.Submitted) -> None:
+        """Handle input submission (Enter key pressed)."""
+        if event.input.id in ["history_days_input", "quiet_days_input"]:
+            self._validate_and_submit_history_days()
+
     def _start_analysis(self) -> None:
         """Start the agent activity analysis."""
-        self.workflow_stage = "analyzing"
+        # Show notification that analysis is starting
-        content = self.query_one("#content_area", Vertical)
-        content.remove_children()
-
-        # Show analyzing message with detailed steps
-        analyzing_msg = Static(
-            f"Analyzing Agent Activity\n"
-            f"{'=' * 50}\n\n"
-            f"Policy: {self.selected_policy.name}\n"
-            f"History Period: {self.history_days} days\n"
-            f"Quiet Threshold: {self.quiet_days} days\n\n"
-            f"Progress:\n"
-            f"Step 1/4: Fetching agents from policy...\n"
-            f"Step 2/4: Pulling execution history (this may take a moment)...\n"
-            f"Step 3/4: Analyzing activity patterns...\n"
-            f"Step 4/4: Categorizing agents...\n\n"
-            f"Please wait - this operation cannot be cancelled.",
-            id="analyzing_message",
-        )
-        analyzing_msg.styles.margin = (2, 1)
-        content.mount(analyzing_msg)
-
-        # Show notification
         self.app.notify(
             "Starting analysis - this may take several minutes for large policies",
             severity="information",
             timeout=5,
         )
 
-        # Perform the analysis asynchronously
+        # Clear the screen to provide a blank canvas for Rust progress output
-        self.call_later(self._perform_analysis)
+        # (Rust output displays over the TUI, so we clear everything except header/footer)
+        try:
+            # Clear title
+            title_widget = self.query_one("#workflow_title", Static)
+            title_widget.update("")
 
-    def _perform_analysis(self) -> None:
+            # Clear status
+            status_widget = self.query_one("#workflow_status", Static)
+            status_widget.update("")
+
+            # Clear content area
+            content = self.query_one("#content_area", Vertical)
+            content.remove_children()
+        except Exception as e:
+            logger.debug(f"Could not clear screen for analysis: {e}")
+
+        # Delay the analysis start to ensure UI refresh completes first
+        # This prevents Rust output from starting before the screen is cleared
+        self.set_timer(0.2, self._perform_analysis_worker)
+
+    def _perform_analysis_worker(self) -> None:
         """Perform the actual agent activity analysis."""
         try:
-            # Update status: Fetching agents
-            self._update_analysis_status("Step 1/4: Fetching agents from policy...")
-
             # Get agents in the selected policy
             agents = self.api.agents_find_by_group(self.selected_policy.groupid)
 
@@ -310,32 +428,11 @@ class QuietAgentWorkflowScreen(Screen):
                 self._show_policy_selection()
                 return
 
-            agent_count = len(agents)
-            self.app.notify(
-                f"Found {agent_count} agents - fetching execution history...",
-                severity="information",
-                timeout=3,
-            )
-
-            # Update status: Pulling execution history
-            self._update_analysis_status(
-                f"Step 2/4: Pulling execution history for {agent_count} agents...\n"
-                f"(This may take several minutes - progress shown in terminal)"
-            )
-
             # Get execution history (this shows progress bars in terminal via airlock_libs)
             policy_exec_history = getPolicyInfo(
                 self.api, self.selected_policy, [1, 2, 6, 7], self.history_days
             )
 
-            # Update status: Analyzing patterns
-            self._update_analysis_status("Step 3/4: Analyzing activity patterns...")
-            self.app.notify(
-                "History retrieved - analyzing patterns...",
-                severity="information",
-                timeout=2,
-            )
-
             if policy_exec_history.empty:
                 logger.info(
                     "No execution history found for the selected policy and time range."
@@ -381,9 +478,6 @@ class QuietAgentWorkflowScreen(Screen):
                     lambda x: True if pd.isna(x) or x > self.quiet_days else False
                 )
 
-            # Update status: Categorizing
-            self._update_analysis_status("Step 4/4: Categorizing agents...")
-
             # Sort agents
             agents = agents.sort_values(
                 by=["execution_count", "hostname"], ascending=[True, True]
@@ -394,7 +488,7 @@ class QuietAgentWorkflowScreen(Screen):
 
             # Categorize agents into DataFrames
             self.enforce_ready_df = agents[agents["enforce_ready"]].copy()
-            self.non_enforce_ready_df = agents[not agents["enforce_ready"]].copy()
+            self.non_enforce_ready_df = agents[~agents["enforce_ready"]].copy()
 
             logger.info(
                 f"Analysis complete: {len(self.enforce_ready_df)} enforce ready, "
@@ -416,27 +510,6 @@ class QuietAgentWorkflowScreen(Screen):
             self.app.notify(f"Analysis failed: {str(e)}", severity="error", timeout=5)
             self._show_policy_selection()
 
-    def _update_analysis_status(self, status_text: str) -> None:
-        """Update the analysis status message."""
-        try:
-            analyzing_msg = self.query_one("#analyzing_message", Static)
-
-            # Build updated message
-            updated_text = (
-                f"Analyzing Agent Activity\n"
-                f"{'=' * 50}\n\n"
-                f"Policy: {self.selected_policy.name}\n"
-                f"History Period: {self.history_days} days\n"
-                f"Quiet Threshold: {self.quiet_days} days\n\n"
-                f"Progress:\n"
-                f"{status_text}\n\n"
-                f"Please wait - this operation cannot be cancelled."
-            )
-
-            analyzing_msg.update(updated_text)
-        except Exception as e:
-            logger.debug(f"Could not update analysis status: {e}")
-
     def _show_results(self) -> None:
         """Show the categorized results."""
         self.workflow_stage = "view_results"
@@ -821,7 +894,7 @@ class QuietAgentWorkflowScreen(Screen):
         # Depending on stage, go back to previous stage or exit
         if self.workflow_stage in ["select_policy", "view_results", "complete"]:
             self.app.pop_screen()
-        elif self.workflow_stage == "select_quiet_days":
+        elif self.workflow_stage == "select_history_days":
             self._show_policy_selection()
         elif self.workflow_stage == "select_enforce_target":
             self._show_results()