Merge pull request #26 from famez/main

MCP tools stopped working after last commits as well as some other code (missing functions that caused exceptions and so on).

Dockerfile

@@ -58,7 +58,7 @@ USER pentestagent
 RUN playwright install
 
 # Copy application code
-COPY . .
+COPY --chown=pentestagent:pentestagent . .
 
 # Expose any needed ports
 EXPOSE 8080

pentestagent/agents/base_agent.py

@@ -835,6 +835,9 @@ Call create_plan with the new steps OR feasible=False."""
         self.state_manager.reset()
         self.conversation_history.clear()
 
+    def add_tools(self, tools : List["Tool"]):
+        self.tools.extend(tools)
+
     async def assist(self, message: str) -> AsyncIterator[AgentMessage]:
         """
         Assist mode - single LLM call, single tool execution if needed.

pentestagent/interface/tui.py

@@ -1312,40 +1312,13 @@ class PentestAgentTUI(App):
                 except Exception as e:
                     self._add_system(f"[!] RAG: {e}")
                     self.rag_engine = None
-
-            # MCP: automatic install/start has been removed. Operators should
-            # install and run any external MCP adapters themselves (for
-            # example under `third_party/`) and then configure
-            # `mcp_servers.json` accordingly. A minimal example adapter is
-            # available at `pentestagent/mcp/example_adapter.py`.
-            try:
-                from ..mcp import MCPManager
-
-                self.mcp_manager = MCPManager()
-                # Start background connect without registering tools into the
-                # TUI process and suppress noisy prints. This keeps the MCP
-                # connection and control socket available while keeping the
-                # TUI's tool list unchanged for the operator.
-                try:
-                    loop = asyncio.get_running_loop()
-                    loop.create_task(self.mcp_manager.connect_all(register=True, quiet=True))
-                except RuntimeError:
-                    # No running loop (unlikely in Textual worker), run in thread
-                    try:
-                        asyncio.run(self.mcp_manager.connect_all(register=False, quiet=True))
-                    except Exception:
-                        pass
-                mcp_server_count = len(self.mcp_manager.list_configured_servers())
-            except Exception:
-                self.mcp_manager = None
-                mcp_server_count = 0
-
+            
             # Runtime - Docker or Local
             if self.use_docker:
                 self._add_system("+ Starting Docker container...")
-                self.runtime = DockerRuntime(mcp_manager=self.mcp_manager)
+                self.runtime = DockerRuntime()
             else:
-                self.runtime = LocalRuntime(mcp_manager=self.mcp_manager)
+                self.runtime = LocalRuntime()
             await self.runtime.start()
 
             # LLM
@@ -1382,6 +1355,37 @@ class PentestAgentTUI(App):
                 rag_engine=self.rag_engine,
             )
 
+            
+            try:
+                from ..mcp import MCPManager
+
+                self.mcp_manager = MCPManager()
+                # Start background connect registering tools on the background to not block.
+
+                async def load_mcp() -> None:
+                    tools = await self.mcp_manager.connect_all()
+                    self.agent.add_tools(tools)
+                    for tool in tools:
+                        register_tool_instance(tool)
+                    self.all_tools = get_all_tools()
+                    self._update_header()
+
+                try:
+                    loop = asyncio.get_running_loop()
+                    loop.create_task(load_mcp())
+                except RuntimeError:
+                    # No running loop (unlikely in Textual worker), run in thread
+                    try:
+                        asyncio.run(load_mcp())
+                    except Exception:
+                        pass
+                mcp_server_count = len(self.mcp_manager.list_configured_servers())
+            except Exception:
+                self.mcp_manager = None
+                mcp_server_count = 0
+
+            
+
             self._set_status("idle", "assist")
             self._is_initializing = False  # Allow input now
 
@@ -1393,6 +1397,8 @@ class PentestAgentTUI(App):
             runtime_str = "Docker" if self.use_docker else "Local"
             # Update persistent header instead of adding an ad-hoc system
             # message to the chat. This keeps the info visible at top.
+            self.mcp_server_count = mcp_server_count
+            self.rag_doc_count = rag_doc_count
             try:
                 self._update_header(model_line=(
                     f"+ PentestAgent ready\n"
@@ -2437,8 +2443,11 @@ 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(getattr(self, "all_tools", [])) if hasattr(self, "all_tools") else 0
+                mode = getattr(self, '_mode', 'assist')
+                if mode == 'assist':
+                    mode = 'assist (use /agent or /crew for autonomous modes)'
                 lines.append(
-                    f"+ PentestAgent ready\n  Model: {getattr(self, 'model', '')} | Tools: {tools_count} | RAG: {getattr(self, 'rag_doc_count', '')}\n  Runtime: {runtime_str} | Mode: {getattr(self, '_mode', 'assist')}"
+                    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}"
                 )
             # Ensure target line is present/updated
             if target is None:

pentestagent/mcp/manager.py

@@ -159,7 +159,7 @@ class MCPManager:
             for n, s in servers.items()
         ]
 
-    async def connect_all(self) -> List[Any]:
+    async def connect_all(self) -> List["Tool"]:
         servers_config = self._load_config()
         all_tools = []
         for name, config in servers_config.items():