Update src/controller/custom_controller.py

src/controller/custom_controller.py

@@ -1,16 +1,17 @@
 import pdb
+
 import pyperclip
 from typing import Optional, Type, Callable, Dict, Any, Union, Awaitable, TypeVar
 from pydantic import BaseModel
 from browser_use.agent.views import ActionResult
 from browser_use.browser.context import BrowserContext
-from browser_use.controller.service import Controller
+from browser_use.controller.service import Controller, DoneAction
 from browser_use.controller.registry.service import Registry, RegisteredAction
 from main_content_extractor import MainContentExtractor
 from browser_use.controller.views import (
     ClickElementAction,
     DoneAction,
-    # ExtractPageContentAction was removed here
+    ExtractPageContentAction,
     GoToUrlAction,
     InputTextAction,
     OpenTabAction,
@@ -34,6 +35,7 @@ logger = logging.getLogger(__name__)
 
 Context = TypeVar('Context')
 
+
 class CustomController(Controller):
     def __init__(self, exclude_actions: list[str] = [],
                  output_model: Optional[Type[BaseModel]] = None,
@@ -79,6 +81,7 @@ class CustomController(Controller):
                 return ActionResult(error=f'File {path} does not exist')
 
             dom_el = await browser.get_dom_element_by_index(index)
+
             file_upload_dom_el = dom_el.get_file_upload_element()
 
             if file_upload_dom_el is None:
@@ -108,9 +111,11 @@ class CustomController(Controller):
             self,
             action: ActionModel,
             browser_context: Optional[BrowserContext] = None,
+            #
             page_extraction_llm: Optional[BaseChatModel] = None,
             sensitive_data: Optional[Dict[str, str]] = None,
             available_file_paths: Optional[list[str]] = None,
+            #
             context: Context | None = None,
     ) -> ActionResult:
         """Execute an action"""
@@ -119,6 +124,7 @@ class CustomController(Controller):
             for action_name, params in action.model_dump(exclude_unset=True).items():
                 if params is not None:
                     if action_name.startswith("mcp"):
+                        # this is a mcp tool
                         logger.debug(f"Invoke MCP tool: {action_name}")
                         mcp_tool = self.registry.registry.actions.get(action_name).function
                         result = await mcp_tool.ainvoke(params)
@@ -152,6 +158,9 @@ class CustomController(Controller):
             self.register_mcp_tools()
 
     def register_mcp_tools(self):
+        """
+        Register the MCP tools used by this controller.
+        """
         if self.mcp_client:
             for server_name in self.mcp_client.server_name_to_tools:
                 for tool in self.mcp_client.server_name_to_tools[server_name]:
@@ -163,9 +172,11 @@ class CustomController(Controller):
                         param_model=create_tool_param_model(tool),
                     )
                     logger.info(f"Add mcp tool: {tool_name}")
+                logger.debug(
+                    f"Registered {len(self.mcp_client.server_name_to_tools[server_name])} mcp tools for {server_name}")
         else:
             logger.warning(f"MCP client not started.")
 
     async def close_mcp_client(self):
         if self.mcp_client:
-            await self.mcp_client.__aexit__(None, None, None)
+            await self.mcp_client.__aexit__(None, None, None)