package uixt import ( "context" "fmt" "strings" "github.com/mark3labs/mcp-go/client" "github.com/mark3labs/mcp-go/mcp" "github.com/pkg/errors" "github.com/rs/zerolog/log" "github.com/httprunner/httprunner/v5/uixt/ai" "github.com/httprunner/httprunner/v5/uixt/option" ) func NewXTDriver(driver IDriver, opts ...option.AIServiceOption) (*XTDriver, error) { driverExt := &XTDriver{ IDriver: driver, client: &MCPClient4XTDriver{ Server: NewMCPServer(), }, loadedMCPClients: make(map[string]client.MCPClient), } services := option.NewAIServiceOptions(opts...) var err error // default to vedem CV service if services.CVService == "" { log.Warn().Msg("no CV service config provided, not using CV service") } else { driverExt.CVService, err = ai.NewCVService(services.CVService) if err != nil { log.Error().Err(err).Msg("init vedem image service failed") return nil, err } } // Handle LLM service initialization if services.LLMConfig != nil { // Use advanced LLM configuration if provided driverExt.LLMService, err = ai.NewLLMServiceWithOptionConfig(services.LLMConfig) if err != nil { return nil, errors.Wrap(err, "init llm service with config failed") } } else if services.LLMService != "" { // Fallback to simple LLM service if no config provided driverExt.LLMService, err = ai.NewLLMService(services.LLMService) if err != nil { return nil, errors.Wrap(err, "init llm service failed") } } else { log.Warn().Msg("no LLM service config provided") } // Register uixt MCP tools to LLM service if it exists if driverExt.LLMService != nil { mcpTools := driverExt.client.Server.ListTools() einoTools := ai.ConvertMCPToolsToEinoToolInfos(mcpTools, "uixt") if err := driverExt.LLMService.RegisterTools(einoTools); err != nil { log.Warn().Err(err).Msg("failed to register uixt tools") } } return driverExt, nil } // XTDriver = IDriver + AI type XTDriver struct { IDriver CVService ai.ICVService // OCR/CV LLMService ai.ILLMService // LLM 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 type MCPClient4XTDriver struct { client.MCPClient Server *MCPServer4XTDriver } func (c *MCPClient4XTDriver) ListTools(ctx context.Context, req mcp.ListToolsRequest) (*mcp.ListToolsResult, error) { tools := c.Server.ListTools() return &mcp.ListToolsResult{Tools: tools}, nil } func (c *MCPClient4XTDriver) CallTool(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { actionName := strings.TrimPrefix(req.Params.Name, "uixt__") actionTool := c.Server.GetToolByAction(option.ActionName(actionName)) if actionTool == nil { return mcp.NewToolResultError(fmt.Sprintf("action %s for tool not found", actionName)), nil } handler := actionTool.Implement() return handler(ctx, req) } func (c *MCPClient4XTDriver) Initialize(ctx context.Context, req mcp.InitializeRequest) (*mcp.InitializeResult, error) { // no need to initialize for local server return &mcp.InitializeResult{}, nil } func (c *MCPClient4XTDriver) Close() error { // no need to close for local server return nil } // GetToolByAction implements ActionToolProvider interface func (c *MCPClient4XTDriver) GetToolByAction(actionName option.ActionName) ActionTool { return c.Server.GetToolByAction(actionName) } func (dExt *XTDriver) ExecuteAction(ctx context.Context, action option.MobileAction) (SessionData, error) { // Find the corresponding tool for this action method tool := dExt.client.Server.GetToolByAction(action.Method) if tool == nil { return SessionData{}, fmt.Errorf("no tool found for action method: %s", action.Method) } // Use the tool's own conversion method req, err := tool.ConvertActionToCallToolRequest(action) if err != nil { return SessionData{}, fmt.Errorf("failed to convert action to MCP tool call: %w", err) } // Execute via MCP tool result, err := dExt.client.CallTool(ctx, req) if err != nil { return SessionData{}, fmt.Errorf("MCP tool call failed: %w", err) } // Check if the tool execution had business logic errors if result.IsError { var errMsg string if len(result.Content) > 0 { errMsg = fmt.Sprintf("invoke tool %s failed: %v", tool.Name(), result.Content) } else { errMsg = fmt.Sprintf("invoke tool %s failed", tool.Name()) } err := errors.New(errMsg) return SessionData{}, err } // For regular actions, collect session data and return it directly sessionData := dExt.GetSession().GetData(true) // reset after getting data log.Debug().Str("tool", string(tool.Name())). Interface("result", result.Content). Msg("executed action via MCP tool") return sessionData, nil } // NewDeviceWithDefault is a helper function to create a device with default options func NewDeviceWithDefault(platform, serial string) (device IDevice, err error) { if serial == "" { return nil, fmt.Errorf("serial is empty") } switch strings.ToLower(platform) { case "android": device, err = NewAndroidDevice(option.WithSerialNumber(serial)) case "ios": device, err = NewIOSDevice( option.WithUDID(serial), option.WithWDAPort(8700), option.WithWDAMjpegPort(8800), option.WithResetHomeOnStartup(false), ) case "browser": device, err = NewBrowserDevice(option.WithBrowserID(serial)) case "harmony": device, err = NewHarmonyDevice(option.WithConnectKey(serial)) default: return nil, fmt.Errorf("unsupported platform: %s", platform) } 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 } // 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 }