From ef44c04ecde7bb28089b98a1d1bf388f3466e43a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BD=99=E6=B3=93=E9=93=AE?= <yuhongzheng@bytedance.com>
Date: Fri, 15 Aug 2025 15:23:45 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E5=8A=A8=E6=80=81?=
 =?UTF-8?q?=E6=94=B9=E5=8F=98=E4=B8=8A=E4=BC=A0=E5=9B=BE=E7=89=87?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 uixt/driver_ext_ai.go | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/uixt/driver_ext_ai.go b/uixt/driver_ext_ai.go
index c849e6a6..71d3f56c 100644
--- a/uixt/driver_ext_ai.go
+++ b/uixt/driver_ext_ai.go
@@ -298,13 +298,14 @@ func (dExt *XTDriver) PlanNextAction(ctx context.Context, prompt string, opts ..
 	// Parse action options to get ResetHistory setting
 	options := option.NewActionOptions(opts...)
 	resetHistory := options.ResetHistory
+	actionOptions := option.NewActionOptions(opts...)
+	screenOptions := []option.ActionOption{option.WithScreenShotFileName("ai_action"), option.WithScreenShotBase64(true)}
+	if actionOptions.ScreenShotWithUpload {
+		screenOptions = append(screenOptions, option.WithScreenShotUpload(true))
+	}
 
 	// Step 1: Take screenshot and convert to base64
-	screenResult, err := dExt.GetScreenResult(
-		option.WithScreenShotFileName("ai_planning"),
-		option.WithScreenShotBase64(true),
-		option.WithScreenShotUpload(true),
-	)
+	screenResult, err := dExt.GetScreenResult(screenOptions...)
 	if err != nil {
 		return nil, err
 	}
@@ -315,6 +316,12 @@ func (dExt *XTDriver) PlanNextAction(ctx context.Context, prompt string, opts ..
 
 	// Step 2: Call model
 	modelCallStartTime := time.Now()
+	var imageURL string
+	if screenResult.UploadedURL != "" {
+		imageURL = screenResult.UploadedURL
+	} else {
+		imageURL = screenResult.ImagePath
+	}
 	planningOpts := &ai.PlanningOptions{
 		UserInstruction: prompt,
 		Message: &schema.Message{
@@ -323,7 +330,7 @@ func (dExt *XTDriver) PlanNextAction(ctx context.Context, prompt string, opts ..
 				{
 					Type: schema.ChatMessagePartTypeImageURL,
 					ImageURL: &schema.ChatMessageImageURL{
-						URL: screenResult.UploadedURL,
+						URL: imageURL,
 					},
 				},
 			},