diff --git a/.gitignore b/.gitignore index c6a733e..b286f93 100644 --- a/.gitignore +++ b/.gitignore @@ -96,3 +96,4 @@ tests/test_*.py loot/token_usage.json mcp_examples/kali/mcp_servers.json pentestagent/mcp/mcp_servers.json +loot/notes.json diff --git a/pentestagent/interface/tui.py b/pentestagent/interface/tui.py index f79d9f3..caed3bc 100644 --- a/pentestagent/interface/tui.py +++ b/pentestagent/interface/tui.py @@ -648,34 +648,69 @@ class ThinkingMessage(Static): class ToolMessage(Static): """Tool execution message""" - # Standard tool icon and color (pa theme) - TOOL_ICON = "$" - TOOL_COLOR = "#9a9a9a" # spirit gray + TOOL_COLOR = "#9a9a9a" + ARG_COLOR = "#6b6b6b" + HINT_COLOR = "#6b6b6b" + + CHEVRON_COLLAPSED = "▶" + CHEVRON_EXPANDED = "▼" + HINT_TEXT = " (click to see result)" + + expanded: bool = reactive(False, layout=True) def __init__(self, tool_name: str, args: str = "", **kwargs): super().__init__(**kwargs) self.tool_name = tool_name self.tool_args = args + self._result_widget: ToolResultMessage | None = None def render(self) -> Text: text = Text() - text.append(f"{self.TOOL_ICON} ", style=self.TOOL_COLOR) - text.append(f"{self.tool_name}", style=self.TOOL_COLOR) - text.append("\n", style="") - # Wrap args + chevron = ( + self.CHEVRON_EXPANDED if self.expanded else self.CHEVRON_COLLAPSED + ) + + # Header line + text.append(f"{chevron} ", style=self.TOOL_COLOR) + text.append(self.tool_name, style=self.TOOL_COLOR) + + # Hint text (only when result exists and is collapsed) + if self._result_widget and not self.expanded: + text.append(self.HINT_TEXT, style=self.HINT_COLOR) + + text.append("\n") + + # Tool arguments if self.tool_args: for line in wrap_text_lines(self.tool_args, width=110): - text.append(f" {line}\n", style="#6b6b6b") + text.append(f" {line}\n", style=self.ARG_COLOR) return text + def attach_result(self, result_widget: "ToolResultMessage") -> None: + """Attach a ToolResultMessage widget below this message.""" + if self._result_widget is not None: + return + + self._result_widget = result_widget + self._result_widget.display = self.expanded + + # Mount directly after this widget + self.mount(self._result_widget, after=self) + + def on_click(self) -> None: + self.expanded = not self.expanded + if self._result_widget: + self._result_widget.display = self.expanded + class ToolResultMessage(Static): """Tool result/output message""" RESULT_ICON = "#" - RESULT_COLOR = "#7a7a7a" + RESULT_COLOR = "#124670" + OUTPUT_COLOR = "#17606d" def __init__(self, tool_name: str, result: str = "", **kwargs): super().__init__(**kwargs) @@ -684,13 +719,14 @@ class ToolResultMessage(Static): def render(self) -> Text: text = Text() + text.append(f"{self.RESULT_ICON} ", style=self.RESULT_COLOR) text.append(f"{self.tool_name} output", style=self.RESULT_COLOR) - text.append("\n", style="") + text.append("\n") if self.result: for line in wrap_text_lines(self.result, width=110): - text.append(f" {line}\n", style="#5a5a5a") + text.append(f" {line}\n", style=self.OUTPUT_COLOR) return text @@ -1685,14 +1721,14 @@ class PentestAgentTUI(App): def _add_thinking(self, content: str) -> None: self._add_message(ThinkingMessage(content)) - def _add_tool(self, name: str, action: str = "") -> None: - self._add_message(ToolMessage(name, action)) + def _add_tool(self, name: str, action: str = "") -> ToolMessage: + tool_message = ToolMessage(name, action) + self._add_message(tool_message) + return tool_message - def _add_tool_result(self, name: str, result: str) -> None: + def _add_tool_result(self, tool_message: ToolMessage, name: str, result: str) -> None: """Display tool execution result""" - # Hide tool output - LLM will synthesize it in its response - # This prevents duplication and keeps the chat clean - pass + tool_message.attach_result(ToolResultMessage(name, result)) def _show_system_prompt(self) -> None: """Display the current system prompt""" @@ -3215,6 +3251,8 @@ Be concise. Use the actual data from notes.""" from ..agents.base_agent import AgentState + last_tool_message: ToolMessage + async for response in self.agent.agent_loop(task): if self._should_stop: self._add_system("[!] Stopped by user") @@ -3241,7 +3279,7 @@ Be concise. Use the actual data from notes.""" for call in response.tool_calls: # Show all tools including finish args_str = str(call.arguments) - self._add_tool(call.name, args_str) + last_tool_message = self._add_tool(call.name, args_str) # Show tool results if response.tool_results: @@ -3251,11 +3289,11 @@ Be concise. Use the actual data from notes.""" continue if result.success: - self._add_tool_result( + self._add_tool_result(last_tool_message, result.tool_name, result.result or "Done" ) else: - self._add_tool_result( + self._add_tool_result(last_tool_message, result.tool_name, f"Error: {result.error}" )