diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py index 0a9da76..b5834be 100644 --- a/pentestagent/interface/tui.py +++ b/pentestagent/interface/tui.py @@ -442,57 +442,66 @@ class ToolsScreen(ModalScreen): self.app.pop_screen() + class MCPScreen(ModalScreen): - """Interactive MCP browser — split-pane layout. + """Interactive MCP browser — split-pane layout.""" - Left pane: tree of MCP servers. Right pane: full description (scrollable). - - """ - - BINDINGS = [Binding("escape", "dismiss", "Close"), Binding("q", "dismiss", "Close")] + BINDINGS = [ + Binding("escape", "dismiss", "Close"), + Binding("q", "dismiss", "Close"), + ] CSS = """ MCPScreen { align: center middle; } """ + from ..mcp import MCPManager def __init__(self, mcp_manager: MCPManager) -> None: super().__init__() self.mcp_manager = mcp_manager + self.selected_server = None + self.selected_tool = None + + # ------------------------------------------------------------ def compose(self) -> ComposeResult: - # Build a split view: left tree, right description with Container(id="mcp-container"): with Horizontal(id="mcp-split"): + + # LEFT SIDE with Vertical(id="mcp-left"): yield Static("MCP Servers", id="mcp-title") yield Tree("MCP Servers", id="mcp-tree") + # RIGHT SIDE with Vertical(id="mcp-right"): yield Static("Description", id="mcp-desc-title") - yield ScrollableContainer(Static("Select a MCP server to view details.", id="mcp-desc"), id="mcp-desc-scroll") + # ---- Toggle Button ---- + yield Button("Enabled: 🔴", id="mcp-toggle-enabled") + + yield ScrollableContainer( + Static("Select a MCP server to view details.", id="mcp-desc"), + id="mcp-desc-scroll" + ) + + yield Center(Button("Close", id="mcp-close")) + # ------------------------------------------------------------ + def on_mount(self) -> None: try: tree = self.query_one("#mcp-tree", Tree) except Exception as e: logging.getLogger(__name__).exception("Failed to query MCP tree: %s", e) - try: - from ..interface.notifier import notify - - notify("warning", f"TUI: failed to initialize MCP tree: {e}") - except Exception as e: - logging.getLogger(__name__).exception("Failed to notify operator about MCP tree init failure: %s", e) return root = tree.root root.allow_expand = True root.show_root = False - # Populate tool nodes - servers = self.mcp_manager.get_all_servers() for server in servers: @@ -500,69 +509,107 @@ class MCPScreen(ModalScreen): for tool in server.tools: server_node.add(tool['name'], data={"tool": tool}) + # Hide toggle button initially + toggle_btn = self.query_one("#mcp-toggle-enabled", Button) + toggle_btn.display = False + try: tree.focus() except Exception as e: logging.getLogger(__name__).exception("Failed to focus MCP tree: %s", e) - try: - from ..interface.notifier import notify - notify("warning", f"TUI: failed to focus MCP tree: {e}") - except Exception as e: - logging.getLogger(__name__).exception("Failed to notify operator about MCP tree focus failure: %s", e) + # ------------------------------------------------------------ @on(Tree.NodeSelected, "#mcp-tree") def on_mcp_selected(self, event: Tree.NodeSelected) -> None: node = event.node + + self.selected_server = node.data.get("server") if node.data else None + self.selected_tool = node.data.get("tool") if node.data else None + + desc_widget = self.query_one("#mcp-desc", Static) + toggle_btn = self.query_one("#mcp-toggle-enabled", Button) + + text = Text() + + if self.selected_server is not None: + mcp = self.selected_server + + text.append(f"{mcp.name}\n", style="bold #d4d4d4") + text.append(f"{mcp.config.description}\n\n", style="#d4d4d4") + + text.append(f"Command: {mcp.config.command}\n", style="#9a9a9a") + text.append(f"Args: {mcp.config.args}\n\n", style="#9a9a9a") + + enabled_icon = "🟢" if mcp.config.enabled else "🔴" + connected_icon = "🟢" if mcp.connected else "🔴" + + text.append(f"Enabled: {enabled_icon}\n", style="#9a9a9a") + text.append(f"Connected: {connected_icon}\n", style="#9a9a9a") + + if mcp.last_error: + text.append(f"\nLast error: {mcp.last_error}\n", style="#e91b1b") + + desc_widget.update(text) + + toggle_btn.display = True + toggle_btn.label = f"Enabled: {enabled_icon}" + + elif self.selected_tool is not None: + text.append(f"{self.selected_tool['description']}\n", style="#d4d4d4") + desc_widget.update(text) + toggle_btn.display = False + + else: + desc_widget.update("Choose a server.") + toggle_btn.display = False + + # ------------------------------------------------------------ + + @on(Button.Pressed, "#mcp-toggle-enabled") + def toggle_enabled(self) -> None: + mcp = self.selected_server + if mcp is None: + return + try: - mcp = node.data.get("server") if node.data else None - tool = node.data.get("tool") if node.data else None + new_value = not mcp.config.enabled - # Update right-hand description pane - try: - desc_widget = self.query_one("#mcp-desc", Static) - text = Text() - if mcp is not None: - text.append(f"{mcp.name}\n", style="bold #d4d4d4") - text.append(f"{mcp.config.description}\n", style="#d4d4d4") - text.append(f"Command: {mcp.config.command}\n", style="#9a9a9a") - text.append(f"Args: {mcp.config.args}\n", style="#9a9a9a") - - enabled_icon = "🟢" if mcp.config.enabled else "🔴" - text.append(f"Enabled: {enabled_icon}\n", style="#9a9a9a") - - connected_icon = "🟢" if mcp.connected else "🔴" - text.append(f"Connected: {connected_icon}\n", style="#9a9a9a") - if mcp.last_error: - text.append(f"Last error message: {mcp.last_error}\n", style="#e91b1b") - elif tool is not None: - text.append(f"{tool['description']}\n", style="#d4d4d4") - else: - text.append(f"Choose a server\n", style="#d4d4d4") - desc_widget.update(text) + if new_value: + self.mcp_manager.enable(mcp.name) + else: + self.mcp_manager.disable(mcp.name) - except Exception as e: - logging.getLogger(__name__).exception("Failed to update mcp description pane: %s", e) - try: - from ..interface.notifier import notify - - notify("warning", f"TUI: failed to update mcp description: {e}") - except Exception as e: - logging.getLogger(__name__).exception("Failed to notify operator about mcp desc update failure: %s", e) except Exception as e: - logging.getLogger(__name__).exception("Unhandled error in on_mcp_selected: %s", e) + logging.getLogger(__name__).exception("Toggle failed: %s", e) try: from ..interface.notifier import notify + notify("error", f"Failed to toggle {mcp.name}: {e}") + except Exception: + pass + return - notify("warning", f"TUI: error handling mcp selection: {e}") - except Exception as e: - logging.getLogger(__name__).exception("Failed to notify operator about mcp selection error: %s", e) + # Refresh UI + self._refresh_description() + + # ------------------------------------------------------------ + + def _refresh_description(self): + """Rebuild right-hand side using current selection.""" + if self.selected_server: + # Simulate reselecting same node + dummy = type("Node", (), {"data": {"server": self.selected_server}}) + event = type("Evt", (), {"node": dummy}) + self.on_mcp_selected(event) + + # ------------------------------------------------------------ @on(Button.Pressed, "#mcp-close") def close_mcp(self) -> None: self.app.pop_screen() + # ----- Main Chat Message Widgets -----