refactor: integrate and optimize MCP tool calling methods

This commit is contained in:
lilong.129
2025-05-27 21:39:17 +08:00
parent 229fd4678c
commit 2fe5b14d63
8 changed files with 140 additions and 120 deletions

View File

@@ -1 +1 @@
v5.0.0-beta-2505272016
v5.0.0-beta-2505272139

View File

@@ -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)

View File

@@ -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) {

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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,
}
}

View File

@@ -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 (

View File

@@ -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
}