From 2fe5b14d6393e705ed44a68a80adf60002e19d38 Mon Sep 17 00:00:00 2001 From: "lilong.129" Date: Tue, 27 May 2025 21:39:17 +0800 Subject: [PATCH] refactor: integrate and optimize MCP tool calling methods --- internal/version/VERSION | 2 +- step_ui.go | 21 ++++++ uixt/android_driver_adb.go | 4 +- uixt/cache.go | 29 ++++++++ uixt/driver_action.go | 1 + uixt/driver_handler.go | 140 ++++++------------------------------- uixt/option/action.go | 5 ++ uixt/sdk.go | 58 +++++++++++++++ 8 files changed, 140 insertions(+), 120 deletions(-) diff --git a/internal/version/VERSION b/internal/version/VERSION index 7f1cebe6..8630c33f 100644 --- a/internal/version/VERSION +++ b/internal/version/VERSION @@ -1 +1 @@ -v5.0.0-beta-2505272016 +v5.0.0-beta-2505272139 diff --git a/step_ui.go b/step_ui.go index ad13749b..b3d9cafa 100644 --- a/step_ui.go +++ b/step_ui.go @@ -448,6 +448,16 @@ func (s *StepMobile) ClosePopups(opts ...option.ActionOption) *StepMobile { return s } +func (s *StepMobile) Call(name string, fn func(), opts ...option.ActionOption) *StepMobile { + s.obj().Actions = append(s.obj().Actions, uixt.MobileAction{ + Method: option.ACTION_CallFunction, + Params: name, // function description + Fn: fn, + Options: option.NewActionOptions(opts...), + }) + return s +} + // Validate switches to step validation. func (s *StepMobile) Validate() *StepMobileUIValidation { return &StepMobileUIValidation{ @@ -804,6 +814,17 @@ func runStepMobileUI(s *SessionRunner, step IStep) (stepResult *StepResult, err continue } + // call custom function + if action.Method == option.ACTION_CallFunction { + if funcDesc, ok := action.Params.(string); ok { + err := uiDriver.Call(funcDesc, action.Fn, action.GetOptions()...) + if err != nil { + return stepResult, err + } + } + continue + } + err = uiDriver.ExecuteAction(context.Background(), action) actionResult.Elapsed = time.Since(actionStartTime).Milliseconds() stepResult.Actions = append(stepResult.Actions, actionResult) diff --git a/uixt/android_driver_adb.go b/uixt/android_driver_adb.go index 8fa6e175..9e2211cd 100644 --- a/uixt/android_driver_adb.go +++ b/uixt/android_driver_adb.go @@ -281,7 +281,9 @@ func (ad *ADBDriver) AppLaunch(packageName string) (err error) { return errors.Wrap(code.MobileUILaunchAppError, fmt.Sprintf("monkey aborted: %s", strings.TrimSpace(sOutput))) } - return nil + + return postHandler(ad, option.ACTION_SetTouchInfo, + option.NewActionOptions(option.WithAntiRisk(true))) } func (ad *ADBDriver) AppTerminate(packageName string) (successful bool, err error) { diff --git a/uixt/cache.go b/uixt/cache.go index 0bdc56d2..cc4c99b1 100644 --- a/uixt/cache.go +++ b/uixt/cache.go @@ -308,3 +308,32 @@ func RegisterXTDriver(serial string, driver *XTDriver) error { return nil } + +// getXTDriverFromCache gets XTDriver from cache using device UUID +func getXTDriverFromCache(driver IDriver) *XTDriver { + // Get device info to find the corresponding XTDriver + device := driver.GetDevice() + if device == nil { + log.Warn().Msg("Cannot get device from driver for MCP hook") + return nil + } + + // Get device UUID (serial/udid/connectKey/browserID) + deviceUUID := device.UUID() + if deviceUUID == "" { + log.Warn().Msg("Cannot get device UUID for MCP hook") + return nil + } + + // Get XTDriver from cache using device UUID as serial + cachedDrivers := ListCachedDrivers() + for _, cached := range cachedDrivers { + if cached.Serial == deviceUUID { + return cached.Driver + } + } + + log.Warn().Str("uuid", deviceUUID). + Msg("Cannot find cached XTDriver for MCP hook") + return nil +} diff --git a/uixt/driver_action.go b/uixt/driver_action.go index 426f7b77..a313e06d 100644 --- a/uixt/driver_action.go +++ b/uixt/driver_action.go @@ -7,6 +7,7 @@ import ( type MobileAction struct { Method option.ActionName `json:"method,omitempty" yaml:"method,omitempty"` Params interface{} `json:"params,omitempty" yaml:"params,omitempty"` + Fn func() `json:"-" yaml:"-"` // used for function action, not serialized Options *option.ActionOptions `json:"options,omitempty" yaml:"options,omitempty"` option.ActionOptions } diff --git a/uixt/driver_handler.go b/uixt/driver_handler.go index 0abd7d87..51d88a01 100644 --- a/uixt/driver_handler.go +++ b/uixt/driver_handler.go @@ -9,7 +9,6 @@ import ( "github.com/httprunner/httprunner/v5/internal/builtin" "github.com/httprunner/httprunner/v5/internal/config" "github.com/httprunner/httprunner/v5/uixt/option" - "github.com/mark3labs/mcp-go/mcp" "github.com/rs/zerolog/log" ) @@ -51,10 +50,7 @@ func preHandler_TapAbsXY(driver IDriver, options *option.ActionOptions, rawX, ra // Call MCP action tool if anti-risk is enabled if options.AntiRisk { - callMCPActionTool(driver, option.ACTION_TapAbsXY, map[string]any{ - "x": rawX, - "y": rawY, - }) + // TODO } x, y = options.ApplyTapOffset(rawX, rawY) @@ -129,6 +125,13 @@ func preHandler_Swipe(driver IDriver, actionType option.ActionName, } func postHandler(driver IDriver, actionType option.ActionName, options *option.ActionOptions) error { + if options.AntiRisk && actionType == option.ACTION_SetTouchInfo { + arguments := getAntiRisk_SetTouchInfo_Arguments(driver) + if arguments != nil { + callMCPActionTool(driver, "evalpkgs", string(actionType), arguments) + } + } + // save screenshot after action if options.PostMarkOperation { // get compressed screenshot buffer @@ -155,129 +158,30 @@ func postHandler(driver IDriver, actionType option.ActionName, options *option.A } // callMCPActionTool calls MCP tool for the given action -func callMCPActionTool(driver IDriver, actionType option.ActionName, arguments map[string]any) { +func callMCPActionTool(driver IDriver, + serverName, actionType string, arguments map[string]any) { // Get XTDriver from cache dExt := getXTDriverFromCache(driver) if dExt == nil { return } - // Define action to MCP server mapping for pre-hooks - serverMapping := getPreHookServerMapping(actionType) - if serverMapping == nil { - return // No MCP hook configured for this action - } - - callMCPTool(dExt, serverMapping.ServerName, serverMapping.ToolName, arguments, actionType) + dExt.CallMCPTool(context.Background(), + serverName, actionType, arguments) } -// MCPServerMapping defines the mapping between action and MCP server/tool -type MCPServerMapping struct { - ServerName string - ToolName string -} - -// getPreHookServerMapping returns MCP server mapping for pre-hooks -// TODO: You can customize these mappings according to your needs -func getPreHookServerMapping(actionType option.ActionName) *MCPServerMapping { - mappings := map[option.ActionName]*MCPServerMapping{ - option.ACTION_TapAbsXY: { - ServerName: "evalpkgs", - ToolName: "log_pre_action", - }, - // Add more mappings as needed - // option.ACTION_Swipe: { - // ServerName: "monitor", - // ToolName: "start_timer", - // }, - } - return mappings[actionType] -} - -// getXTDriverFromCache gets XTDriver from cache using device UUID -func getXTDriverFromCache(driver IDriver) *XTDriver { - // Get device info to find the corresponding XTDriver +func getAntiRisk_SetTouchInfo_Arguments(driver IDriver) map[string]interface{} { + var deviceModel string device := driver.GetDevice() - if device == nil { - log.Warn().Msg("Cannot get device from driver for MCP hook") - return nil - } - - // Get device UUID (serial/udid/connectKey/browserID) - deviceUUID := device.UUID() - if deviceUUID == "" { - log.Warn().Msg("Cannot get device UUID for MCP hook") - return nil - } - - // Get XTDriver from cache using device UUID as serial - cachedDrivers := ListCachedDrivers() - for _, cached := range cachedDrivers { - if cached.Serial == deviceUUID { - return cached.Driver + if adbDevice, ok := device.(*AndroidDevice); ok { + var err error + deviceModel, err = adbDevice.Model() + if err != nil { + return nil } } - log.Warn().Str("uuid", deviceUUID). - Msg("Cannot find cached XTDriver for MCP hook") - return nil -} - -// callMCPTool calls the specified MCP tool -func callMCPTool(dExt *XTDriver, serverName, toolName string, arguments map[string]any, actionType option.ActionName) { - // Get MCP client - mcpClient, exists := dExt.GetMCPClient(serverName) - if !exists { - log.Debug().Str("server", serverName).Msg("MCP server not found for hook") - return - } - - // Create context with timeout - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - // Prepare arguments - if arguments == nil { - arguments = make(map[string]any) - } - // Add action type and hook type to arguments - arguments["action_type"] = string(actionType) - - // Call MCP tool - req := mcp.CallToolRequest{ - Params: struct { - Name string `json:"name"` - Arguments map[string]any `json:"arguments,omitempty"` - Meta *struct { - ProgressToken mcp.ProgressToken `json:"progressToken,omitempty"` - } `json:"_meta,omitempty"` - }{ - Name: toolName, - Arguments: arguments, - }, - } - - result, err := mcpClient.CallTool(ctx, req) - if err != nil { - log.Debug().Err(err). - Str("server", serverName). - Str("tool", toolName). - Msg("MCP hook call failed") - return - } - - if result.IsError { - log.Debug(). - Str("server", serverName). - Str("tool", toolName). - Interface("content", result.Content). - Msg("MCP hook returned error") - return - } - - log.Debug(). - Str("server", serverName). - Str("tool", toolName). - Str("action", string(actionType)). - Msg("MCP hook called successfully") + return map[string]interface{}{ + "deviceModel": deviceModel, + } } diff --git a/uixt/option/action.go b/uixt/option/action.go index f145261a..47580ea5 100644 --- a/uixt/option/action.go +++ b/uixt/option/action.go @@ -83,6 +83,11 @@ const ( ACTION_UninstallApp ActionName = "uninstall_app" ACTION_DownloadApp ActionName = "download_app" ACTION_Finished ActionName = "finished" + ACTION_CallFunction ActionName = "call_function" + + // anti-risk actions + ACTION_SetTouchInfo ActionName = "set_touch_info" + ACTION_SetTouchInfoList ActionName = "set_touch_info_list" ) const ( diff --git a/uixt/sdk.go b/uixt/sdk.go index 4ce4b05d..2da69e48 100644 --- a/uixt/sdk.go +++ b/uixt/sdk.go @@ -160,3 +160,61 @@ func (dExt *XTDriver) GetMCPClient(serverName string) (client.MCPClient, bool) { client, exists := dExt.loadedMCPClients[serverName] return client, exists } + +// CallMCPTool calls the specified MCP tool +func (dExt *XTDriver) CallMCPTool(ctx context.Context, + serverName, toolName string, arguments map[string]any) (result *mcp.CallToolResult, err error) { + // Get MCP client + + mcpClient, exists := dExt.GetMCPClient(serverName) + if !exists { + log.Warn().Str("server", serverName).Msg("MCP server not found") + return nil, fmt.Errorf("MCP server %s not found", serverName) + } + + // Prepare arguments + if arguments == nil { + arguments = make(map[string]any) + } + + log.Debug().Str("server", serverName).Str("tool", toolName). + Interface("arguments", arguments).Msg("call MCP tool") + + // Call MCP tool + req := mcp.CallToolRequest{ + Params: struct { + Name string `json:"name"` + Arguments map[string]any `json:"arguments,omitempty"` + Meta *struct { + ProgressToken mcp.ProgressToken `json:"progressToken,omitempty"` + } `json:"_meta,omitempty"` + }{ + Name: toolName, + Arguments: arguments, + }, + } + + result, err = mcpClient.CallTool(ctx, req) + if err != nil { + log.Debug().Err(err). + Str("server", serverName). + Str("tool", toolName). + Msg("MCP hook call failed") + return nil, err + } + + if result.IsError { + log.Debug(). + Str("server", serverName). + Str("tool", toolName). + Interface("content", result.Content). + Msg("MCP hook returned error") + return nil, fmt.Errorf("MCP hook returned error") + } + + log.Debug(). + Str("server", serverName). + Str("tool", toolName). + Msg("MCP hook called successfully") + return result, nil +}