GithubBackup/httprunner
1
0
Fork 0
mirror of https://github.com/httprunner/httprunner.git synced 2026-09-04 23:16:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement MCP hooks integration with anti_risk option

Browse Source
This commit is contained in:
lilong.129
2025-05-27 19:46:08 +08:00
parent f4cc74b3ca
commit 866cc0e4d2
10 changed files with 222 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
uixt/driver_ext_ai.go
+1 -1
View File
@@ -62,7 +62,7 @@ func (dExt *XTDriver) AIAction(text string, opts ...option.ActionOption) error {
},
}
_, err = dExt.Client.CallTool(context.Background(), req)
_, err = dExt.client.CallTool(context.Background(), req)
if err != nil {
return err
}
uixt/driver_handler.go
+138
View File
@@ -1,6 +1,7 @@
package uixt
import (
"context"
"fmt"
"path/filepath"
"time"
@@ -8,6 +9,7 @@ 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"
)
@@ -47,6 +49,14 @@ func (dExt *XTDriver) Call(desc string, fn func(), opts ...option.ActionOption)
func preHandler_TapAbsXY(driver IDriver, options *option.ActionOptions, rawX, rawY float64) (
x, y float64, err error) {
// Call MCP action tool if anti-risk is enabled
if options.AntiRisk {
callMCPActionTool(driver, option.ACTION_TapAbsXY, map[string]any{
"x": rawX,
"y": rawY,
})
}
x, y = options.ApplyTapOffset(rawX, rawY)
// mark UI operation
@@ -143,3 +153,131 @@ func postHandler(driver IDriver, actionType option.ActionName, options *option.A
}
return nil
}
// callMCPActionTool calls MCP tool for the given action
func callMCPActionTool(driver IDriver, actionType option.ActionName, 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)
}
// 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
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
}
// 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")
}
uixt/mcp_server.go
+8 -8
View File
@@ -1071,10 +1071,10 @@ func (t *ToolSwipeCoordinate) Implement() server.ToolHandlerFunc {
params := []float64{unifiedReq.FromX, unifiedReq.FromY, unifiedReq.ToX, unifiedReq.ToY}
opts := []option.ActionOption{}
if unifiedReq.Duration > 0 && unifiedReq.Duration > 0 {
if unifiedReq.Duration > 0 {
opts = append(opts, option.WithDuration(unifiedReq.Duration))
}
if unifiedReq.PressDuration > 0 && unifiedReq.PressDuration > 0 {
if unifiedReq.PressDuration > 0 {
opts = append(opts, option.WithPressDuration(unifiedReq.PressDuration))
}
@@ -1146,10 +1146,10 @@ func (t *ToolSwipeToTapApp) Implement() server.ToolHandlerFunc {
}
// Add numeric options
if unifiedReq.MaxRetryTimes > 0 && unifiedReq.MaxRetryTimes > 0 {
if unifiedReq.MaxRetryTimes > 0 {
opts = append(opts, option.WithMaxRetryTimes(unifiedReq.MaxRetryTimes))
}
if unifiedReq.Index > 0 && unifiedReq.Index > 0 {
if unifiedReq.Index > 0 {
opts = append(opts, option.WithIndex(unifiedReq.Index))
}
@@ -1218,10 +1218,10 @@ func (t *ToolSwipeToTapText) Implement() server.ToolHandlerFunc {
}
// Add numeric options
if unifiedReq.MaxRetryTimes > 0 && unifiedReq.MaxRetryTimes > 0 {
if unifiedReq.MaxRetryTimes > 0 {
opts = append(opts, option.WithMaxRetryTimes(unifiedReq.MaxRetryTimes))
}
if unifiedReq.Index > 0 && unifiedReq.Index > 0 {
if unifiedReq.Index > 0 {
opts = append(opts, option.WithIndex(unifiedReq.Index))
}
@@ -1290,10 +1290,10 @@ func (t *ToolSwipeToTapTexts) Implement() server.ToolHandlerFunc {
}
// Add numeric options
if unifiedReq.MaxRetryTimes > 0 && unifiedReq.MaxRetryTimes > 0 {
if unifiedReq.MaxRetryTimes > 0 {
opts = append(opts, option.WithMaxRetryTimes(unifiedReq.MaxRetryTimes))
}
if unifiedReq.Index > 0 && unifiedReq.Index > 0 {
if unifiedReq.Index > 0 {
opts = append(opts, option.WithIndex(unifiedReq.Index))
}
uixt/option/action.go
+13
View File
@@ -175,6 +175,9 @@ type ActionOptions struct {
ScreenOptions
// Anti-risk options
AntiRisk bool `json:"anti_risk,omitempty" yaml:"anti_risk,omitempty" desc:"Enable anti-risk MCP tool calls"`
// Custom options
Custom map[string]interface{} `json:"custom,omitempty" yaml:"custom,omitempty" desc:"Custom options"`
}
@@ -286,6 +289,10 @@ func (o *ActionOptions) Options() []ActionOption {
options = append(options, WithMatchOne(true))
}
if o.AntiRisk {
options = append(options, WithAntiRisk(true))
}
// custom options
if o.Custom != nil {
for k, v := range o.Custom {
@@ -494,6 +501,12 @@ func WithIgnoreNotFoundError(ignoreError bool) ActionOption {
}
}
func WithAntiRisk(antiRisk bool) ActionOption {
return func(o *ActionOptions) {
o.AntiRisk = antiRisk
}
}
// HTTP API direct usage methods
// ValidateForHTTPAPI validates the request for HTTP API usage
uixt/sdk.go
+26 -5
View File
@@ -15,9 +15,10 @@ import (
func NewXTDriver(driver IDriver, opts ...option.AIServiceOption) (*XTDriver, error) {
driverExt := &XTDriver{
IDriver: driver,
Client: &MCPClient4XTDriver{
client: &MCPClient4XTDriver{
Server: NewMCPServer(),
},
loadedMCPClients: make(map[string]client.MCPClient),
}
services := option.NewAIServiceOptions(opts...)
@@ -47,7 +48,8 @@ type XTDriver struct {
CVService ai.ICVService // OCR/CV
LLMService ai.ILLMService // LLM
Client *MCPClient4XTDriver // MCP Client
client *MCPClient4XTDriver // MCP Client for built-in uixt server
loadedMCPClients map[string]client.MCPClient // External MCP clients
}
// MCPClient4XTDriver is a mock MCP client that only implements the methods used by the host
@@ -80,9 +82,9 @@ func (c *MCPClient4XTDriver) Close() error {
return nil
}
func (dExt *XTDriver) ExecuteAction(action MobileAction) (err error) {
func (dExt *XTDriver) ExecuteAction(ctx context.Context, action MobileAction) (err error) {
// Find the corresponding tool for this action method
tool := dExt.Client.Server.GetToolByAction(action.Method)
tool := dExt.client.Server.GetToolByAction(action.Method)
if tool == nil {
return fmt.Errorf("no tool found for action method: %s", action.Method)
}
@@ -94,7 +96,7 @@ func (dExt *XTDriver) ExecuteAction(action MobileAction) (err error) {
}
// Execute via MCP tool
result, err := dExt.Client.CallTool(context.Background(), req)
result, err := dExt.client.CallTool(ctx, req)
if err != nil {
return fmt.Errorf("MCP tool call failed: %w", err)
}
@@ -139,3 +141,22 @@ func NewDeviceWithDefault(platform, serial string) (device IDevice, err error) {
return device, err
}
// SetMCPClients sets the external MCP clients for the driver
func (dExt *XTDriver) SetMCPClients(clients map[string]client.MCPClient) {
if dExt.loadedMCPClients == nil {
dExt.loadedMCPClients = make(map[string]client.MCPClient)
}
for name, client := range clients {
dExt.loadedMCPClients[name] = client
}
}
// GetMCPClient returns the MCP client for the specified server name
func (dExt *XTDriver) GetMCPClient(serverName string) (client.MCPClient, bool) {
if dExt.loadedMCPClients == nil {
return nil, false
}
client, exists := dExt.loadedMCPClients[serverName]
return client, exists
}