diff --git a/pentestagent/agents/base_agent.py b/pentestagent/agents/base_agent.py index f0d5c38..84c8b6e 100644 --- a/pentestagent/agents/base_agent.py +++ b/pentestagent/agents/base_agent.py @@ -241,7 +241,7 @@ class BaseAgent(ABC): response = await self.llm.generate( system_prompt=self.get_system_prompt(), messages=self._format_messages_for_llm(), - tools=self.tools, + tools=[t for t in self.tools if t.t.enabled], #Only, enabled tools. ) # Case 1: Empty response (Error) @@ -870,7 +870,8 @@ Call create_plan with the new steps OR feasible=False.""" self.conversation_history.append(AgentMessage(role="user", content=message)) # Filter out 'finish' tool - not needed for single-shot assist mode - assist_tools = [t for t in self.tools if t.name != "finish"] + # Pass to the agent only the enabled tools + assist_tools = [t for t in self.tools if t.name != "finish" and t.enabled] # Single LLM call with tools available response = await self.llm.generate( @@ -958,7 +959,8 @@ Call create_plan with the new steps OR feasible=False.""" self.conversation_history.append(AgentMessage(role="user", content=message)) # Filter out 'finish' tool - not needed for interact mode, we read finish_reason - assist_tools = [t for t in self.tools if t.name != "finish"] + # Pass to the agent only enabled tools. + interact_tools = [t for t in self.tools if t.name != "finish" and t.enabled] while True: @@ -966,7 +968,7 @@ Call create_plan with the new steps OR feasible=False.""" response = await self.llm.generate( system_prompt=self.get_system_prompt(mode="interact"), messages=self._format_messages_for_llm(), - tools=assist_tools, + tools=interact_tools, ) # If LLM wants to use tools, execute and return result diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py index 0a888bb..a9d349f 100644 --- a/pentestagent/interface/tui.py +++ b/pentestagent/interface/tui.py @@ -343,10 +343,14 @@ class ToolsScreen(ModalScreen): CSS = """ ToolsScreen { align: center middle; } """ + from ..tools import Tool - def __init__(self, tools: List[Any]) -> None: + def __init__(self, tools: List[Tool], tui: "PentestAgentTUI") -> None: + from ..tools import Tool super().__init__() self.tools = tools + self.selected_tool: Optional[Tool] = None + self.tui = tui def compose(self) -> ComposeResult: # Build a split view: left tree, right description @@ -357,6 +361,7 @@ class ToolsScreen(ModalScreen): yield Tree("TOOLS", id="tools-tree") with Vertical(id="tools-right"): + yield Button("Enabled: 🔴", id="tool-toggle-enabled") yield Static("Description", id="tools-desc-title") yield ScrollableContainer(Static("Select a tool to view details.", id="tools-desc"), id="tools-desc-scroll") @@ -384,6 +389,10 @@ class ToolsScreen(ModalScreen): name = getattr(t, "name", str(t)) root.add(name, data={"tool": t}) + # Hide toggle button initially + toggle_btn = self.query_one("#tool-toggle-enabled", Button) + toggle_btn.display = False + try: tree.focus() except Exception as e: @@ -400,10 +409,12 @@ class ToolsScreen(ModalScreen): node = event.node try: tool = node.data.get("tool") if node.data else None - name = node.label or (getattr(tool, "name", str(tool)) if tool else "Unknown") + self.selected_tool = tool + name = getattr(tool, "name", str(tool)) if tool else "Unknown" # Prefer Tool.description (registered tools use this), then fall back desc = None + tool_enabled = False if tool is not None: desc = getattr(tool, "description", None) if not desc: @@ -412,16 +423,29 @@ class ToolsScreen(ModalScreen): or getattr(tool, "help_text", None) or getattr(tool, "__doc__", None) ) + + tool_enabled = getattr(tool, "enabled", False) + if not desc: desc = "No description available." + # Update right-hand description pane try: desc_widget = self.query_one("#tools-desc", Static) + toggle_btn = self.query_one("#tool-toggle-enabled", Button) text = Text() text.append(f"{name}\n", style="bold #d4d4d4") text.append(str(desc), style="#d4d4d4") desc_widget.update(text) + + if tool: + enabled_icon = "🟢" if tool_enabled else "🔴" + toggle_btn.display = True + toggle_btn.label = f"Enabled: {enabled_icon}" + else: + toggle_btn.display = False + except Exception as e: logging.getLogger(__name__).exception("Failed to update tool description pane: %s", e) try: @@ -444,6 +468,39 @@ class ToolsScreen(ModalScreen): self.app.pop_screen() + # ------------------------------------------------------------ + + @on(Button.Pressed, "#tool-toggle-enabled") + async def toggle_enabled(self) -> None: + + try: + + tool = self.selected_tool + if tool is None: + return + + tool.enabled = not tool.enabled + + await self._refresh_tool_widget() + + except Exception as e: + logging.getLogger(__name__).exception("Toggle failed: %s", e) + try: + from ..interface.notifier import notify + # Make best effort naming + tool_name = getattr(tool, "name", None) or (tool.get("name") if isinstance(tool, dict) else str(tool)) + notify("error", f"Failed to toggle tool {tool_name}: {e}") + except Exception: + pass + return + + async def _refresh_tool_widget(self) -> None: + if self.selected_tool: + # Simulate reselecting same node + dummy = type("Node", (), {"data": {"tool": self.selected_tool}}) + event = type("Evt", (), {"node": dummy}) + self.on_tool_selected(event) + self.tui._update_header() class MCPScreen(ModalScreen): """Interactive MCP browser — split-pane layout.""" @@ -2227,7 +2284,7 @@ Be concise. Use the actual data from notes.""" # Open the interactive tools browser (split-pane). try: if self.agent: - await self.push_screen(ToolsScreen(tools=self.agent.get_tools())) + await self.push_screen(ToolsScreen(tools=self.agent.get_tools(), tui=self)) except Exception: # Fallback: list tools in the system area if UI push fails from ..runtime.runtime import detect_environment @@ -2829,7 +2886,9 @@ Be concise. Use the actual data from notes.""" else: # try to recreate a compact model/runtime line runtime_str = "Docker" if getattr(self, "use_docker", False) else "Local" - tools_count = len(self.agent.get_tools()) + tools_count = 0 + if self.agent: + tools_count = len([t for t in self.agent.get_tools() if t.enabled]) mode = getattr(self, '_mode', "") mode += ' (use /assist for single tool execution, /agent or /crew for autonomous modes, /interact for interactive chat)' lines.append(