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

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