aroniscunt/Browser_Web_UI_Automation

1import asyncio
2import pdb
3import sys
4import time
5
6sys.path.append(".")
7
8from dotenv import load_dotenv
9
10load_dotenv()
11
12
13async def test_mcp_client():
14    from src.utils.mcp_client import setup_mcp_client_and_tools, create_tool_param_model
15
16    test_server_config = {
17        "mcpServers": {
18            # "markitdown": {
19            #     "command": "docker",
20            #     "args": [
21            #         "run",
22            #         "--rm",
23            #         "-i",
24            #         "markitdown-mcp:latest"
25            #     ]
26            # },
27            "desktop-commander": {
28                "command": "npx",
29                "args": [
30                    "-y",
31                    "@wonderwhy-er/desktop-commander"
32                ]
33            },
34            # "filesystem": {
35            #     "command": "npx",
36            #     "args": [
37            #         "-y",
38            #         "@modelcontextprotocol/server-filesystem",
39            #         "/Users/xxx/ai_workspace",
40            #     ]
41            # },
42        }
43    }
44
45    mcp_tools, mcp_client = await setup_mcp_client_and_tools(test_server_config)
46
47    for tool in mcp_tools:
48        tool_param_model = create_tool_param_model(tool)
49        print(tool.name)
50        print(tool.description)
51        print(tool_param_model.model_json_schema())
52    pdb.set_trace()
53
54
55async def test_controller_with_mcp():
56    import os
57    from src.controller.custom_controller import CustomController
58    from browser_use.controller.registry.views import ActionModel
59
60    mcp_server_config = {
61        "mcpServers": {
62            # "markitdown": {
63            #     "command": "docker",
64            #     "args": [
65            #         "run",
66            #         "--rm",
67            #         "-i",
68            #         "markitdown-mcp:latest"
69            #     ]
70            # },
71            "desktop-commander": {
72                "command": "npx",
73                "args": [
74                    "-y",
75                    "@wonderwhy-er/desktop-commander"
76                ]
77            },
78            # "filesystem": {
79            #     "command": "npx",
80            #     "args": [
81            #         "-y",
82            #         "@modelcontextprotocol/server-filesystem",
83            #         "/Users/xxx/ai_workspace",
84            #     ]
85            # },
86        }
87    }
88
89    controller = CustomController()
90    await controller.setup_mcp_client(mcp_server_config)
91    action_name = "mcp.desktop-commander.execute_command"
92    action_info = controller.registry.registry.actions[action_name]
93    param_model = action_info.param_model
94    print(param_model.model_json_schema())
95    params = {"command": f"python ./tmp/test.py"
96              }
97    validated_params = param_model(**params)
98    ActionModel_ = controller.registry.create_action_model()
99    # Create ActionModel instance with the validated parameters
100    action_model = ActionModel_(**{action_name: validated_params})
101    result = await controller.act(action_model)
102    result = result.extracted_content
103    print(result)
104    if result and "Command is still running. Use read_output to get more output." in result and "PID" in \
105            result.split("\n")[0]:
106        pid = int(result.split("\n")[0].split("PID")[-1].strip())
107        action_name = "mcp.desktop-commander.read_output"
108        action_info = controller.registry.registry.actions[action_name]
109        param_model = action_info.param_model
110        print(param_model.model_json_schema())
111        params = {"pid": pid}
112        validated_params = param_model(**params)
113        action_model = ActionModel_(**{action_name: validated_params})
114        output_result = ""
115        while True:
116            time.sleep(1)
117            result = await controller.act(action_model)
118            result = result.extracted_content
119            if result:
120                pdb.set_trace()
121                output_result = result
122                break
123        print(output_result)
124        pdb.set_trace()
125    await controller.close_mcp_client()
126    pdb.set_trace()
127
128
129if __name__ == '__main__':
130    # asyncio.run(test_mcp_client())
131    asyncio.run(test_controller_with_mcp())
132