From 52278db871dcc9c26387b93302b29a4b2b7d86ec Mon Sep 17 00:00:00 2001 From: famez Date: Sat, 28 Feb 2026 20:11:55 +0100 Subject: [PATCH 1/8] feat: Added interactive mode (#1) * Adding interactive mode. * Added content of /loot dir to .gitignore. * When writing just a mode on the TUI, it will automatically switch to that mode (no need to prepend /agent, /assist, /interact or /crew before. --- .gitignore | 3 +- pentestagent/agents/base_agent.py | 80 +++++++++ pentestagent/agents/pa_agent/pa_agent.py | 9 +- pentestagent/agents/prompts/__init__.py | 1 + pentestagent/agents/prompts/pa_interact.jinja | 85 ++++++++++ pentestagent/interface/tui.py | 156 ++++++++++++++++-- 6 files changed, 316 insertions(+), 18 deletions(-) create mode 100644 pentestagent/agents/prompts/pa_interact.jinja diff --git a/.gitignore b/.gitignore index b286f93..2e1b15c 100644 --- a/.gitignore +++ b/.gitignore @@ -93,7 +93,6 @@ tests/test_*.py # Workspaces directory (user data should not be committed) /workspaces/ -loot/token_usage.json +loot/* mcp_examples/kali/mcp_servers.json pentestagent/mcp/mcp_servers.json -loot/notes.json diff --git a/pentestagent/agents/base_agent.py b/pentestagent/agents/base_agent.py index 63a43e0..f865bee 100644 --- a/pentestagent/agents/base_agent.py +++ b/pentestagent/agents/base_agent.py @@ -944,6 +944,86 @@ Call create_plan with the new steps OR feasible=False.""" self.state_manager.transition_to(AgentState.COMPLETE) + async def interact(self, message: str) -> AsyncIterator[AgentMessage]: + """ + Interactive mode + + Args: + message: The user message to respond to + + Yields: + AgentMessage objects + """ + self.state_manager.transition_to(AgentState.THINKING) + 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"] + + + while True: + # Single LLM call with tools available + response = await self.llm.generate( + system_prompt=self.get_system_prompt(mode="interact"), + messages=self._format_messages_for_llm(), + tools=assist_tools, + ) + + # If LLM wants to use tools, execute and return result + if response.tool_calls: + # Build tool calls list + tool_calls = [ + ToolCall( + id=tc.id if hasattr(tc, "id") else str(i), + name=( + tc.function.name + if hasattr(tc, "function") + else tc.get("name", "") + ), + arguments=self._parse_arguments(tc), + ) + for i, tc in enumerate(response.tool_calls) + ] + + # Store in history (minimal content to save tokens) + assistant_msg = AgentMessage( + role="assistant", content=response.content or "", tool_calls=tool_calls + ) + self.conversation_history.append(assistant_msg) + yield assistant_msg + + # NOW execute the tools (this can take a while) + self.state_manager.transition_to(AgentState.EXECUTING) + tool_results = await self._execute_tools(response.tool_calls) + + tool_result_msg = AgentMessage( + role="tool_result", content="", tool_results=tool_results + ) + self.conversation_history.append(tool_result_msg) + yield tool_result_msg + + else: + # Direct response, no tools needed + assistant_msg = AgentMessage( + role="assistant", content=response.content or "" + ) + self.conversation_history.append(assistant_msg) + yield assistant_msg + + # Check finish_reason to determine if we should return control to user + finish_reason = getattr(response, 'finish_reason', None) + if finish_reason == 'length': + # Context limit reached + break + elif finish_reason == 'stop': + # Natural stop - LLM finished responding + break + elif finish_reason == 'tool_calls': + # LLM wants to use tools + pass + + self.state_manager.transition_to(AgentState.COMPLETE) + def _format_tool_results(self, results: List[ToolResult]) -> str: """Format tool results as a simple response.""" parts = [] diff --git a/pentestagent/agents/pa_agent/pa_agent.py b/pentestagent/agents/pa_agent/pa_agent.py index f4449fe..cf43de8 100644 --- a/pentestagent/agents/pa_agent/pa_agent.py +++ b/pentestagent/agents/pa_agent/pa_agent.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, List, Optional from ..base_agent import BaseAgent -from ..prompts import pa_agent, pa_assist +from ..prompts import pa_agent, pa_assist, pa_interact if TYPE_CHECKING: from ...knowledge import RAGEngine @@ -132,7 +132,12 @@ class PentestAgentAgent(BaseAgent): env = self.runtime.environment # Select template based on mode - template = pa_assist if mode == "assist" else pa_agent + if mode == "assist": + template = pa_assist + elif mode == "interact": + template = pa_interact + else: + template = pa_agent return template.render( target=self.target, diff --git a/pentestagent/agents/prompts/__init__.py b/pentestagent/agents/prompts/__init__.py index db12a63..4f67e0e 100644 --- a/pentestagent/agents/prompts/__init__.py +++ b/pentestagent/agents/prompts/__init__.py @@ -24,3 +24,4 @@ def load_prompt(name: str) -> Template: pa_agent = load_prompt("pa_agent") pa_assist = load_prompt("pa_assist") pa_crew = load_prompt("pa_crew") +pa_interact = load_prompt("pa_interact") \ No newline at end of file diff --git a/pentestagent/agents/prompts/pa_interact.jinja b/pentestagent/agents/prompts/pa_interact.jinja new file mode 100644 index 0000000..9c54f4d --- /dev/null +++ b/pentestagent/agents/prompts/pa_interact.jinja @@ -0,0 +1,85 @@ +# PentestAgent – Interactive Red Team Assistant + +You are PentestAgent, an interactive penetration testing assistant. + +## Behavior +Your primary goal is to assist the user conversationally during a penetration test. +You MUST: +- Ask clarifying questions when needed. +- Think out loud (to the user) about strategies, ideas, attack paths, risks, hypotheses. +- Suggest tools, commands, or approaches proactively. +- Use tools ONLY when explicitly needed to progress or when the user requests it. +- Before using a tool, explain why you want to use it and ask the user if they agree. +- After tool outputs, interpret the results for the user, propose next steps, and continue the dialogue. + +## Penetration Testing Authorization +You are operating in an authorized, legal penetration test. +Do NOT ask again for authorization. Permission is granted. + +## Decision Policy for Tools +Follow this rule: + +1. **Talk first.** + - Explain your reasoning. + - Suggest ideas. + +2. **Use a tool ONLY IF:** + - The user explicitly asks to run it, OR + - Running it is the logical next step AND you have told the user your intention. + +3. **After tool use:** + - Summarize findings. + - Offer next steps. + - Continue in conversational mode. + +## Interactivity Rules +You MUST remain conversational: +- Ask questions such as: “Do you want me to run a scan?”, “Which target would you prefer to scan now?”” +- Guide the user through the pentest. +- Propose attack paths, hypotheses, and options. + + +{% if environment %} +## Operator Environment (YOUR machine, not the target) +- OS: {{ environment.os }} ({{ environment.os_version }}) +- Architecture: {{ environment.architecture }} +- Shell: {{ environment.shell }} +{% endif %} + +## Tools +{% for tool in tools %} +- **{{ tool.name }}**: {{ tool.description }} +{% endfor %} + +{% if environment and environment.available_tools %} +## Available CLI Tools +{% for tool in environment.available_tools %} + • {{ tool.name }}{% if tool.category %} ({{ tool.category }}){% endif %} +{% endfor %} +{% endif %} + +{% if environment %} +## Output Directories +- loot/notes.json (working notes) +- loot/reports/ (generated reports) +- loot/artifacts/ (screenshots, captured files) +{% endif %} + +{% if target %} +## Target +{{ target }} +{% endif %} +{% if scope %} +Scope: {{ scope | join(', ') }} +{% endif %} + +{% if rag_context %} +## Context +{{ rag_context }} +{% endif %} + +{% if notes_context %} +## Saved Notes (from previous tasks) +(Notes may be truncated. Use `notes(action='read', key='...')` to see full content if needed.) +{{ notes_context }} +{% endif %} diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py index c6f1210..fc11262 100644 --- a/pentestagent/interface/tui.py +++ b/pentestagent/interface/tui.py @@ -835,6 +835,8 @@ class StatusBar(Static): text.append(" :: Crew ", style="#9a9a9a") elif self.mode == "agent": text.append(" >> Agent ", style="#9a9a9a") + elif self.mode == "interact": + text.append(" >> Interact ", style="#9a9a9a") else: text.append(" >> Assist ", style="#9a9a9a") @@ -1412,7 +1414,7 @@ class PentestAgentTUI(App): self.rag_engine = None # RAG engine # State - self._mode = "assist" # "assist", "agent", or "crew" + self._mode = "assist" # "assist", "agent", "crew" or "interact" self._is_running = False self._is_initializing = True # Block input during init self._should_stop = False @@ -1747,7 +1749,7 @@ class PentestAgentTUI(App): def _show_system_prompt(self) -> None: """Display the current system prompt""" if self.agent: - prompt = self.agent.get_system_prompt() + prompt = self.agent.get_system_prompt(self._mode) self._add_system(f"=== System Prompt ===\n{prompt}") else: self._add_system("Agent not initialized") @@ -2184,14 +2186,21 @@ Be concise. Use the actual data from notes.""" return self._add_user(message) + - # Hide crew sidebar when entering assist mode - self._hide_sidebar() - - # Use assist mode by default if self.agent and not self._is_running: - # Schedule assist run and keep task handle (do not wrap in asyncio.create_task; @work returns a Worker) - self._current_worker = self._run_assist(message) + if self._mode == "assist": + # Schedule assist run and keep task handle (do not wrap in asyncio.create_task; @work returns a Worker) + self._hide_sidebar() + self._current_worker = self._run_assist(message) + elif self._mode == "interact": + self._hide_sidebar() + self._current_worker = self._run_interact(message) + elif self._mode == "agent": + self._hide_sidebar() + self._current_worker = self._run_agent_mode(message) + elif self._mode == "crew": + self._current_worker = self._run_crew_mode(message) async def _handle_command(self, cmd: str) -> None: """Handle slash commands""" @@ -2239,7 +2248,6 @@ Be concise. Use the actual data from notes.""" ) self._add_system(msg) - elif cmd_lower.startswith("/mcp"): await self._parse_mcp_command(cmd_original) elif cmd_lower in ["/quit", "/exit", "/q"]: @@ -2484,14 +2492,24 @@ Be concise. Use the actual data from notes.""" await self._parse_agent_command(cmd_original) elif cmd_original.startswith("/crew"): await self._parse_crew_command(cmd_original) + elif cmd_original.startswith("/interact"): + await self._parse_interact_command(cmd_original) + elif cmd_original.startswith("/assist"): + await self._parse_assist_command(cmd_original) else: self._add_system(f"Unknown command: {cmd}\nType /help for commands.") async def _parse_agent_command(self, cmd: str) -> None: """Parse and execute /agent command""" + self._set_status("idle", "agent") + self._update_header() + self._add_system( + "Changed to agent mode\n" + ) + # Remove /agent prefix - rest = cmd[6:].strip() + rest = cmd[len("/agent"):].strip() if not rest: self._add_system( @@ -2520,8 +2538,15 @@ Be concise. Use the actual data from notes.""" async def _parse_crew_command(self, cmd: str) -> None: """Parse and execute /crew command""" + + self._set_status("idle", "crew") + self._update_header() + self._add_system( + "Changed to crew mode\n" + ) + # Remove /crew prefix - rest = cmd[5:].strip() + rest = cmd[len("/crew"):].strip() if not rest: self._add_system( @@ -2548,6 +2573,48 @@ Be concise. Use the actual data from notes.""" # Schedule crew run and keep task handle (do not wrap in asyncio.create_task; @work returns a Worker) self._current_worker = self._run_crew_mode(target) + async def _parse_interact_command(self, cmd: str) -> None: + # Use interact mode by default + + self._set_status("idle", "interact") + self._update_header() + self._add_system( + "Changed to interact mode\n" + ) + + message = cmd[len("/interact"):].strip() + if not message: + self._add_system( + "Usage: /interact \n" + "Example: /interact Can you help me recon the target site?\n" + ) + return + + if self.agent and not self._is_running: + # Schedule interact run and keep task handle (do not wrap in asyncio.create_task; @work returns a Worker) + self._current_worker = self._run_interact(message) + + async def _parse_assist_command(self, cmd: str) -> None: + # Use assist mode by default + + self._set_status("idle", "assist") + self._update_header() + self._add_system( + "Changed to assist mode\n" + ) + + message = cmd[len("/assist"):].strip() + if not message: + self._add_system( + "Usage: /assist \n" + "Example: /assist Can you help me recon the target site?\n" + ) + return + + if self.agent and not self._is_running: + # Schedule assist run and keep task handle (do not wrap in asyncio.create_task; @work returns a Worker) + self._current_worker = self._run_assist(message) + async def _parse_mcp_command(self, cmd: str) -> None: # Remove /agent prefix rest = cmd[len("/mcp"):].strip() @@ -2756,9 +2823,8 @@ Be concise. Use the actual data from notes.""" # 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()) - mode = getattr(self, '_mode', 'assist') - if mode == 'assist': - mode = 'assist (use /agent or /crew for autonomous modes)' + mode = getattr(self, '_mode', "") + mode += ' (use /assist for single tool execution, /agent or /crew for autonomous modes, /interact for interactive chat)' lines.append( f"+ PentestAgent ready\n Model: {getattr(self, 'model', '')} | Tools: {tools_count} | MCP: {getattr(self, 'mcp_server_count', '')} | RAG: {getattr(self, 'rag_doc_count', '')}\n Runtime: {runtime_str} | Mode: {mode}" ) @@ -3202,6 +3268,68 @@ Be concise. Use the actual data from notes.""" finally: self._is_running = False + @work(thread=False) + async def _run_interact(self, message: str) -> None: + """Run in interact mode""" + if not self.agent: + self._add_system("[!] Agent not ready") + return + + self._is_running = True + self._should_stop = False + self._set_status("thinking", "interact") + + + #Use to track the tool message with the result. + tool_messages_mapping: dict[str, ToolMessage] = {} + + try: + async for response in self.agent.interact(message): + if self._should_stop: + self._add_system("[!] Stopped by user") + break + + self._set_status("processing") + + # Show thinking/plan FIRST if there's content with tool calls + if response.content: + content = response.content.strip() + if response.tool_calls: + self._add_thinking(content) + else: + self._add_assistant(content) + + # Show tool calls (skip 'finish' - internal control) + if response.tool_calls: + for call in response.tool_calls: + args_str = str(call.arguments) + tool_messages_mapping[call.id] = self._add_tool(call.name, args_str) + + # Show tool results (displayed after execution completes) + if response.tool_results: + for result in response.tool_results: + if result.tool_call_id not in tool_messages_mapping: + continue + if result.success: + self._add_tool_result(tool_messages_mapping[result.tool_call_id], + result.tool_name, result.result or "Done" + ) + else: + self._add_tool_result(tool_messages_mapping[result.tool_call_id], + result.tool_name, f"Error: {result.error}" + ) + + self._set_status("idle", "interact") + + except asyncio.CancelledError: + self._add_system("[!] Cancelled") + self._set_status("idle", "interact") + except Exception as e: + self._add_system(f"[!] Error: {e}") + self._set_status("error") + finally: + self._is_running = False + @work(thread=False) async def _run_assist(self, message: str) -> None: """Run in assist mode - single response""" From a1a8cbbe5d1943b65483740328f600206025413e Mon Sep 17 00:00:00 2001 From: famez Date: Sat, 28 Feb 2026 23:20:34 +0100 Subject: [PATCH 2/8] fix: Added policy on pa_interact to read notes. Before starting the pentest to resume. --- pentestagent/agents/prompts/pa_interact.jinja | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/pentestagent/agents/prompts/pa_interact.jinja b/pentestagent/agents/prompts/pa_interact.jinja index 9c54f4d..b650ce8 100644 --- a/pentestagent/agents/prompts/pa_interact.jinja +++ b/pentestagent/agents/prompts/pa_interact.jinja @@ -38,6 +38,49 @@ You MUST remain conversational: - Guide the user through the pentest. - Propose attack paths, hypotheses, and options. +## Notes Handling Policy (Critical) + +- You MUST treat stored notes as historical and append-only. +- You MUST NEVER overwrite, delete, or destructively update previous notes. +- Use the Notes tool **only** through its API (you are agnostic to any underlying file path or `notes.json` file). +- At the start of each session, you MUST: + - `list` all available notes + - `read` relevant notes + - Summarize prior context +- When adding information: + - Always append as a new note (`create`). + - If a key already exists, create a new versioned key rather than updating. +- Destructive actions (`delete`, destructive `update`) are **forbidden** unless the user explicitly instructs otherwise. + +## Previous Session Recovery (Artifacts & Reports) + +At the beginning of every session: + +1. Attempt to inspect `loot/artifacts/` using available tools. + - If you cannot list files, ask the user to provide a directory listing. + - Summarize any useful artifacts (screenshots, dumps, captures, configs, etc.). + +2. Attempt to inspect `loot/reports/`. + - If tools permit, list and summarize existing reports. + - If tools are not available, ask the user to provide the latest report or summary. + +3. After gathering notes, artifacts, and report context, ask the user: + - Whether they want to resume from the previous session context, + - Or start a new phase. + +You MUST NOT modify or delete artifacts or reports unless explicitly instructed. + +## Session Initialization Procedure + +Before performing any pentesting actions: + +1. Use the Notes tool to load existing context (`list` → `read` relevant entries). +2. Inspect artifacts. +3. Inspect reports. +4. Present a concise summary of findings. +5. Ask the user how they want to proceed (resume or start new). + +Proceed only after the user responds. {% if environment %} ## Operator Environment (YOUR machine, not the target) From 1ad2c50f99048875fd2918d09740eef14d514165 Mon Sep 17 00:00:00 2001 From: famez Date: Sat, 28 Feb 2026 23:21:14 +0100 Subject: [PATCH 3/8] fix: Added /interact mode to /help command --- pentestagent/interface/tui.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py index fc11262..5ca2f10 100644 --- a/pentestagent/interface/tui.py +++ b/pentestagent/interface/tui.py @@ -159,14 +159,16 @@ class HelpScreen(ModalScreen): def _get_help_text(self) -> str: header = ( - "[bold]Modes:[/] Assist | Agent | Crew\n" + "[bold]Modes:[/] Assist | Agent | Crew | Interact\n" "[bold]Keys:[/] Enter=Send Up/Down=History Ctrl+Q=Quit\n\n" "[bold]Commands:[/]\n" ) cmds = [ + ("/assist ", "Run in assist mode"), ("/agent ", "Run in agent mode"), ("/crew ", "Run multi-agent crew mode"), + ("/interact ", "Run in interact mode"), ("/target ", "Set target"), ("/prompt", "Show system prompt"), ("/memory", "Show memory stats"), From 64a78f2e9753e9c0efdf756c6304456bc6cd0817 Mon Sep 17 00:00:00 2001 From: famez Date: Sat, 28 Feb 2026 23:22:01 +0100 Subject: [PATCH 4/8] fix(notes): Notes were not properly loading at the beginning. For correct loading with get_all_notes() function. --- pentestagent/interface/tui.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py index 5ca2f10..00f15bb 100644 --- a/pentestagent/interface/tui.py +++ b/pentestagent/interface/tui.py @@ -1486,6 +1486,10 @@ class PentestAgentTUI(App): notify("warning", f"TUI: failed to register notifier callback: {e}") except Exception as ne: logging.getLogger(__name__).exception("Failed to notify operator about notifier registration failure: %s", ne) + + #Added, because the _notes variable is not properly loaded at the beginning and the notes.json is recreated unless you do a /notes command explicitely. + from ..tools.notes import get_all_notes + await get_all_notes() # Call the textual worker - decorator returns a Worker, not a coroutine _ = cast(Any, self._initialize_agent()) From 016dbcf3335256b25c7dae441275cd6080c055a3 Mon Sep 17 00:00:00 2001 From: famez Date: Sat, 28 Feb 2026 23:22:51 +0100 Subject: [PATCH 5/8] feat(notes): Allow listing all the notes (not truncated) --- pentestagent/tools/notes/__init__.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/pentestagent/tools/notes/__init__.py b/pentestagent/tools/notes/__init__.py index 52b1612..10139f0 100644 --- a/pentestagent/tools/notes/__init__.py +++ b/pentestagent/tools/notes/__init__.py @@ -148,12 +148,12 @@ def _validate_note_schema(category: str, metadata: Dict[str, Any]) -> str | None @register_tool( name="notes", - description="Manage persistent notes for key findings. Actions: create, read, update, delete, list.", + description="Manage persistent notes for key findings. Actions: create, read, update, delete, list (2 options, all or the truncated text to 60 characters).", schema=ToolSchema( properties={ "action": { "type": "string", - "enum": ["create", "read", "update", "delete", "list"], + "enum": ["create", "read", "update", "delete", "list_all", "list_truncated"], "description": "The action to perform", }, "key": { @@ -399,7 +399,7 @@ async def notes(arguments: dict, runtime) -> str: _save_notes_unlocked() return f"Deleted note '{key}'" - elif action == "list": + elif action == "list_all" or action == "list_truncated": if not _notes: return "No notes saved" @@ -417,9 +417,13 @@ async def notes(arguments: dict, runtime) -> str: lines.append(f"\n## {cat.title()}") for k, v in by_category[cat]: content = v["content"] - display_val = ( - content if len(content) <= 60 else content[:57] + "..." - ) + + display_val = content + if action == "list_truncated": + display_val = ( + content if len(content) <= 60 else content[:57] + "..." + ) + conf = v.get("confidence", "medium") lines.append(f" [{k}] ({conf}) {display_val}") From 1d483dbe8255c239e58d7887e72790ae4472ddbd Mon Sep 17 00:00:00 2001 From: famez Date: Sat, 28 Feb 2026 23:23:19 +0100 Subject: [PATCH 6/8] Comment modified. --- pentestagent/agents/base_agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pentestagent/agents/base_agent.py b/pentestagent/agents/base_agent.py index f865bee..f0d5c38 100644 --- a/pentestagent/agents/base_agent.py +++ b/pentestagent/agents/base_agent.py @@ -957,7 +957,7 @@ Call create_plan with the new steps OR feasible=False.""" self.state_manager.transition_to(AgentState.THINKING) self.conversation_history.append(AgentMessage(role="user", content=message)) - # Filter out 'finish' tool - not needed for single-shot assist mode + # 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"] From 0dc1db00134f2e9379f97884293c86d52a5abe09 Mon Sep 17 00:00:00 2001 From: famez Date: Sun, 1 Mar 2026 10:18:42 +0100 Subject: [PATCH 7/8] fix(tui): Avoid warning on text editor --- pentestagent/interface/tui.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py index 00f15bb..0a888bb 100644 --- a/pentestagent/interface/tui.py +++ b/pentestagent/interface/tui.py @@ -2226,7 +2226,8 @@ Be concise. Use the actual data from notes.""" elif cmd_lower == "/tools": # Open the interactive tools browser (split-pane). try: - await self.push_screen(ToolsScreen(tools=self.agent.get_tools())) + if self.agent: + await self.push_screen(ToolsScreen(tools=self.agent.get_tools())) except Exception: # Fallback: list tools in the system area if UI push fails from ..runtime.runtime import detect_environment From 7cb9428e6cb33365e7f1655f92cff6f34fd234bb Mon Sep 17 00:00:00 2001 From: famez Date: Sun, 1 Mar 2026 12:19:16 +0100 Subject: [PATCH 8/8] feat(tools): Allow enabling or disabling for the agent. --- pentestagent/agents/base_agent.py | 10 +++-- pentestagent/interface/tui.py | 67 +++++++++++++++++++++++++++++-- 2 files changed, 69 insertions(+), 8 deletions(-) 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(