From 0c20fe7b029d307572ec104c1661c13a19da4003 Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Tue, 20 May 2025 22:36:46 +0800
Subject: [PATCH] feat: add argument --with-uixt to start built-in uixt MCP
 server

---
 cmd/mcphost.go                |  4 ++-
 internal/version/VERSION      |  2 +-
 mcphost/chat_test.go          |  4 +--
 mcphost/dump_test.go          |  2 +-
 mcphost/host.go               | 46 ++++++++++++++++++++++++++++++-----
 mcphost/host_test.go          | 22 ++++++++---------
 mcphost/mcp_server.go         | 25 ++++++++++++++++---
 runner.go                     |  2 +-
 server/main.go                |  2 +-
 uixt/driver_ext_ai.go         |  2 +-
 uixt/driver_ext_screenshot.go |  4 +--
 11 files changed, 84 insertions(+), 31 deletions(-)

diff --git a/cmd/mcphost.go b/cmd/mcphost.go
index 9efc71c9..bf078514 100644
--- a/cmd/mcphost.go
+++ b/cmd/mcphost.go
@@ -15,7 +15,7 @@ var CmdMCPHost = &cobra.Command{
 	Long:  `mcphost is a command-line tool that allows you to interact with MCP tools.`,
 	RunE: func(cmd *cobra.Command, args []string) error {
 		// Create MCP host
-		host, err := mcphost.NewMCPHost(mcpConfigPath)
+		host, err := mcphost.NewMCPHost(mcpConfigPath, withUIXT)
 		if err != nil {
 			return fmt.Errorf("failed to create MCP host: %w", err)
 		}
@@ -40,9 +40,11 @@ var CmdMCPHost = &cobra.Command{
 var (
 	mcpConfigPath string
 	dumpPath      string
+	withUIXT      bool
 )
 
 func init() {
 	CmdMCPHost.Flags().StringVarP(&mcpConfigPath, "mcp-config", "c", "$HOME/.hrp/mcp.json", "path to the MCP config file")
 	CmdMCPHost.Flags().StringVar(&dumpPath, "dump", "", "path to save the exported tools JSON file")
+	CmdMCPHost.Flags().BoolVar(&withUIXT, "with-uixt", false, "start built-in uixt MCP server")
 }
diff --git a/internal/version/VERSION b/internal/version/VERSION
index 1ddd2cc8..57d5ccb4 100644
--- a/internal/version/VERSION
+++ b/internal/version/VERSION
@@ -1 +1 @@
-v5.0.0-beta-2505201803
+v5.0.0-beta-2505202236
diff --git a/mcphost/chat_test.go b/mcphost/chat_test.go
index f7383789..a50ec56f 100644
--- a/mcphost/chat_test.go
+++ b/mcphost/chat_test.go
@@ -9,7 +9,7 @@ import (
 )
 
 func TestRunPromptWithNoToolCall(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", true)
 	require.NoError(t, err)
 
 	chat, err := host.NewChat(context.Background())
@@ -21,7 +21,7 @@ func TestRunPromptWithNoToolCall(t *testing.T) {
 }
 
 func TestRunPromptWithToolCall(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", true)
 	require.NoError(t, err)
 
 	chat, err := host.NewChat(context.Background())
diff --git a/mcphost/dump_test.go b/mcphost/dump_test.go
index 8e6a32e9..1785aa75 100644
--- a/mcphost/dump_test.go
+++ b/mcphost/dump_test.go
@@ -13,7 +13,7 @@ import (
 )
 
 func TestConvertToolsToRecordsFromFile(t *testing.T) {
-	hub, err := NewMCPHost("./testdata/test.mcp.json")
+	hub, err := NewMCPHost("./testdata/test.mcp.json", true)
 	require.NoError(t, err)
 
 	// use ExportToolsToJSON to dump tools to JSON file
diff --git a/mcphost/host.go b/mcphost/host.go
index b567521a..f8375151 100644
--- a/mcphost/host.go
+++ b/mcphost/host.go
@@ -13,6 +13,7 @@ import (
 	"github.com/cloudwego/eino/components/tool"
 	"github.com/cloudwego/eino/schema"
 	"github.com/httprunner/httprunner/v5/internal/version"
+	"github.com/httprunner/httprunner/v5/uixt"
 	"github.com/mark3labs/mcp-go/client"
 	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/pkg/errors"
@@ -24,6 +25,8 @@ type MCPHost struct {
 	mu          sync.RWMutex
 	connections map[string]*Connection
 	config      *MCPConfig
+	withUIXT    bool
+	drivers     map[string]*uixt.XTDriver
 }
 
 // Connection represents a connection to an MCP server
@@ -40,7 +43,7 @@ type MCPTools struct {
 }
 
 // NewMCPHost creates a new MCPHost instance
-func NewMCPHost(configPath string) (*MCPHost, error) {
+func NewMCPHost(configPath string, withUIXT bool) (*MCPHost, error) {
 	config, err := LoadMCPConfig(configPath)
 	if err != nil {
 		return nil, err
@@ -49,6 +52,8 @@ func NewMCPHost(configPath string) (*MCPHost, error) {
 	host := &MCPHost{
 		connections: make(map[string]*Connection),
 		config:      config,
+		drivers:     make(map[string]*uixt.XTDriver),
+		withUIXT:    withUIXT,
 	}
 
 	// Initialize MCP servers
@@ -62,11 +67,13 @@ func NewMCPHost(configPath string) (*MCPHost, error) {
 // InitServers initializes all MCP servers
 func (h *MCPHost) InitServers(ctx context.Context) error {
 	// initialize uixt MCP server
-	h.connections["uixt"] = &Connection{
-		Client: &MCPClient4XTDriver{
-			server: NewMCPServer(),
-		},
-		Config: nil,
+	if h.withUIXT {
+		h.connections["uixt"] = &Connection{
+			Client: &MCPClient4XTDriver{
+				server: NewMCPServer(),
+			},
+			Config: nil,
+		}
 	}
 
 	for name, server := range h.config.MCPServers {
@@ -387,3 +394,30 @@ func handleToolError(result *mcp.CallToolResult) error {
 	}
 	return fmt.Errorf("tool error: unknown error")
 }
+
+// ScreenshotBase64 get screenshot base64 for the given platform and serial
+func (h *MCPHost) ScreenshotBase64(ctx context.Context, platform, serial string) (string, error) {
+	driver, err := h.GetOrCreateDriver(platform, serial)
+	if err != nil {
+		return "", err
+	}
+	return uixt.GetScreenShotBufferBase64(driver)
+}
+
+// GetOrCreateDriver get or create a driver for the given platform and serial
+func (h *MCPHost) GetOrCreateDriver(platform, serial string) (*uixt.XTDriver, error) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	cacheKey := fmt.Sprintf("%s_%s", platform, serial)
+	if driver, ok := h.drivers[cacheKey]; ok {
+		return driver, nil
+	}
+
+	driverExt, err := initDriverExt(platform, serial)
+	if err != nil {
+		return nil, err
+	}
+	// store driver in cache
+	h.drivers[cacheKey] = driverExt
+	return driverExt, nil
+}
diff --git a/mcphost/host_test.go b/mcphost/host_test.go
index 5bc45113..7bf97c42 100644
--- a/mcphost/host_test.go
+++ b/mcphost/host_test.go
@@ -12,20 +12,20 @@ import (
 
 func TestNewMCPHost(t *testing.T) {
 	// Test with valid config file
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 	assert.NotNil(t, host)
 	assert.NotNil(t, host.config)
 	assert.NotEmpty(t, host.config.MCPServers)
 
 	// Test with non-existent config file
-	host, err = NewMCPHost("./testdata/non_existent.json")
+	host, err = NewMCPHost("./testdata/non_existent.json", false)
 	require.Error(t, err, "expected error when config file does not exist")
 	assert.Nil(t, host)
 }
 
 func TestInitServers(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 
 	// Verify connections are established
@@ -35,7 +35,7 @@ func TestInitServers(t *testing.T) {
 }
 
 func TestGetClient(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 
 	// Test getting existing client
@@ -50,7 +50,7 @@ func TestGetClient(t *testing.T) {
 }
 
 func TestGetTools(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 
 	ctx := context.Background()
@@ -81,7 +81,7 @@ func TestGetTools(t *testing.T) {
 }
 
 func TestGetTool(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 
 	ctx := context.Background()
@@ -104,7 +104,7 @@ func TestGetTool(t *testing.T) {
 }
 
 func TestInvokeTool(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 
 	ctx := context.Background()
@@ -132,7 +132,7 @@ func TestInvokeTool(t *testing.T) {
 }
 
 func TestCallEinoTool(t *testing.T) {
-	hub, err := NewMCPHost("./testdata/test.mcp.json")
+	hub, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 
 	ctx := context.Background()
@@ -147,7 +147,7 @@ func TestCallEinoTool(t *testing.T) {
 }
 
 func TestCloseServers(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 
 	// Verify servers are connected
@@ -162,7 +162,7 @@ func TestCloseServers(t *testing.T) {
 }
 
 func TestConcurrentOperations(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 
 	// Test concurrent tool invocations
@@ -193,7 +193,7 @@ func TestConcurrentOperations(t *testing.T) {
 }
 
 func TestDisabledServer(t *testing.T) {
-	host, err := NewMCPHost("./testdata/test.mcp.json")
+	host, err := NewMCPHost("./testdata/test.mcp.json", false)
 	require.NoError(t, err)
 
 	// Verify only enabled servers are connected
diff --git a/mcphost/mcp_server.go b/mcphost/mcp_server.go
index fb3bd50b..f20974c4 100644
--- a/mcphost/mcp_server.go
+++ b/mcphost/mcp_server.go
@@ -126,7 +126,8 @@ func (ums *MCPServer4XTDriver) handleTapXY(ctx context.Context, request mcp.Call
 		return mcp.NewToolResultError("parse parameters error: " + err.Error()), nil
 	}
 	if tapReq.Duration > 0 {
-		err := driverExt.Drag(tapReq.X, tapReq.Y, tapReq.X, tapReq.Y, option.WithDuration(tapReq.Duration))
+		err := driverExt.Drag(tapReq.X, tapReq.Y, tapReq.X, tapReq.Y,
+			option.WithDuration(tapReq.Duration))
 		if err != nil {
 			return mcp.NewToolResultError("Tap failed: " + err.Error()), nil
 		}
@@ -136,7 +137,9 @@ func (ums *MCPServer4XTDriver) handleTapXY(ctx context.Context, request mcp.Call
 			return mcp.NewToolResultError("Tap failed: " + err.Error()), nil
 		}
 	}
-	return mcp.NewToolResultText("Tap successful."), nil
+	return mcp.NewToolResultText(
+		fmt.Sprintf("tap (%f,%f) success", tapReq.X, tapReq.Y),
+	), nil
 }
 
 // handleSwipe handles the swipe tool call.
@@ -153,11 +156,15 @@ func (ums *MCPServer4XTDriver) handleSwipe(ctx context.Context, request mcp.Call
 	if swipeReq.Duration > 0 {
 		actionOptions = append(actionOptions, option.WithDuration(swipeReq.Duration/1000.0))
 	}
-	err = driverExt.Swipe(swipeReq.FromX, swipeReq.FromY, swipeReq.ToX, swipeReq.ToY, actionOptions...)
+	err = driverExt.Swipe(swipeReq.FromX, swipeReq.FromY,
+		swipeReq.ToX, swipeReq.ToY, actionOptions...)
 	if err != nil {
 		return mcp.NewToolResultError("Swipe failed: " + err.Error()), nil
 	}
-	return mcp.NewToolResultText("Swipe successful."), nil
+	return mcp.NewToolResultText(
+		fmt.Sprintf("swipe (%f,%f)->(%f,%f) success",
+			swipeReq.FromX, swipeReq.FromY, swipeReq.ToX, swipeReq.ToY),
+	), nil
 }
 
 // handleScreenShot handles the screenshot tool call.
@@ -198,6 +205,16 @@ func (ums *MCPServer4XTDriver) setupXTDriver(_ context.Context, args map[string]
 		}
 	}
 
+	driverExt, err := initDriverExt(platform, serial)
+	if err != nil {
+		return nil, err
+	}
+	// store driver in cache
+	ums.driverCache.Store(cacheKey, driverExt)
+	return driverExt, nil
+}
+
+func initDriverExt(platform, serial string) (*uixt.XTDriver, error) {
 	// init device
 	var device uixt.IDevice
 	var err error
diff --git a/runner.go b/runner.go
index 9d7446fa..fbb9a027 100644
--- a/runner.go
+++ b/runner.go
@@ -318,7 +318,7 @@ func NewCaseRunner(testcase TestCase, hrpRunner *HRPRunner) (*CaseRunner, error)
 
 	// init MCP servers
 	if config.MCPConfigPath != "" {
-		mcpHost, err := mcphost.NewMCPHost(config.MCPConfigPath)
+		mcpHost, err := mcphost.NewMCPHost(config.MCPConfigPath, false)
 		if err != nil {
 			log.Error().Err(err).Msg("init MCP hub failed")
 			return nil, err
diff --git a/server/main.go b/server/main.go
index 201003c3..8e195f95 100644
--- a/server/main.go
+++ b/server/main.go
@@ -25,7 +25,7 @@ type Router struct {
 }
 
 func (r *Router) InitMCPHost(configPath string) error {
-	mcpHost, err := mcphost.NewMCPHost(configPath)
+	mcpHost, err := mcphost.NewMCPHost(configPath, false)
 	if err != nil {
 		log.Error().Err(err).Msg("init MCP host failed")
 		return err
diff --git a/uixt/driver_ext_ai.go b/uixt/driver_ext_ai.go
index 5af66357..b10fabbc 100644
--- a/uixt/driver_ext_ai.go
+++ b/uixt/driver_ext_ai.go
@@ -119,7 +119,7 @@ func (dExt *XTDriver) AIAssert(assertion string, opts ...option.ActionOption) er
 		return errors.New("LLM service is not initialized")
 	}
 
-	screenShotBase64, err := getScreenShotBufferBase64(dExt.IDriver)
+	screenShotBase64, err := GetScreenShotBufferBase64(dExt.IDriver)
 	if err != nil {
 		return err
 	}
diff --git a/uixt/driver_ext_screenshot.go b/uixt/driver_ext_screenshot.go
index 2236ace4..03dff094 100644
--- a/uixt/driver_ext_screenshot.go
+++ b/uixt/driver_ext_screenshot.go
@@ -222,8 +222,8 @@ func getScreenShotBuffer(driver IDriver) (compressedBufSource *bytes.Buffer, err
 	return compressBufSource, nil
 }
 
-// getScreenShotBufferBase64 takes a screenshot, returns the compressed image buffer in base64 format
-func getScreenShotBufferBase64(driver IDriver) (compressedBufBase64 string, err error) {
+// GetScreenShotBufferBase64 takes a screenshot, returns the compressed image buffer in base64 format
+func GetScreenShotBufferBase64(driver IDriver) (compressedBufBase64 string, err error) {
 	compressBufSource, err := getScreenShotBuffer(driver)
 	if err != nil {
 		return "", err