From ae0d28c26deebbab9331c1479955d3e18be3426f Mon Sep 17 00:00:00 2001
From: "lilong.129" <lilong.129@bytedance.com>
Date: Mon, 30 Jun 2025 13:55:17 +0800
Subject: [PATCH] feat: add action timeout for StartToGoal

---
 uixt/driver_ext_ai.go | 36 ++++++++++++++++--------------------
 uixt/option/action.go |  6 +++---
 2 files changed, 19 insertions(+), 23 deletions(-)

diff --git a/uixt/driver_ext_ai.go b/uixt/driver_ext_ai.go
index aacde7dd..68a4303c 100644
--- a/uixt/driver_ext_ai.go
+++ b/uixt/driver_ext_ai.go
@@ -18,7 +18,17 @@ import (
 
 func (dExt *XTDriver) StartToGoal(ctx context.Context, prompt string, opts ...option.ActionOption) ([]*PlanningExecutionResult, error) {
 	options := option.NewActionOptions(opts...)
-	log.Info().Int("max_retry_times", options.MaxRetryTimes).Msg("StartToGoal")
+	log.Info().Int("max_retry_times", options.MaxRetryTimes).
+		Int("timeout_seconds", options.Timeout).
+		Msg("StartToGoal")
+
+	// Create timeout context for entire StartToGoal process if Timeout is specified
+	if options.Timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, time.Duration(options.Timeout)*time.Second)
+		defer cancel()
+		log.Info().Int("timeout_seconds", options.Timeout).Msg("StartToGoal timeout configured for entire process")
+	}
 
 	var allPlannings []*PlanningExecutionResult
 	var attempt int
@@ -119,25 +129,11 @@ func (dExt *XTDriver) StartToGoal(ctx context.Context, prompt string, opts ...op
 					planningResult.SubActions = append(planningResult.SubActions, subActionResult)
 				}()
 
-				// Create action context with timeout if specified
-				actionCtx := ctx
-				if options.Timeout > 0 {
-					var cancel context.CancelFunc
-					actionCtx, cancel = context.WithTimeout(ctx, time.Duration(options.Timeout)*time.Second)
-					defer cancel()
-				}
-
-				// Execute the tool call with timeout
-				if err := dExt.invokeToolCall(actionCtx, toolCall, opts...); err != nil {
-					// Check if the error is due to timeout
-					if errors.Is(err, context.DeadlineExceeded) {
-						log.Warn().
-							Str("action", toolCall.Function.Name).
-							Int("timeout_seconds", options.Timeout).
-							Msg("action timeout exceeded, continuing to next action")
-						subActionResult.Error = errors.New("action timeout exceeded")
-						return nil // Continue to next action instead of failing the entire StartToGoal
-					}
+				if err := dExt.invokeToolCall(ctx, toolCall, opts...); err != nil {
+					log.Warn().
+						Str("action", toolCall.Function.Name).
+						Err(err).
+						Msg("invoke tool call failed")
 					subActionResult.Error = err
 					return err
 				}
diff --git a/uixt/option/action.go b/uixt/option/action.go
index f23ecad8..15305edf 100644
--- a/uixt/option/action.go
+++ b/uixt/option/action.go
@@ -200,7 +200,7 @@ type ActionOptions struct {
 	PressDuration float64         `json:"press_duration,omitempty" yaml:"press_duration,omitempty" desc:"Press duration in seconds"`
 	Steps         int             `json:"steps,omitempty" yaml:"steps,omitempty" desc:"Number of steps for action"`
 	Direction     interface{}     `json:"direction,omitempty" yaml:"direction,omitempty" desc:"Direction for swipe operations or custom coordinates"`
-	Timeout       int             `json:"timeout,omitempty" yaml:"timeout,omitempty" desc:"Timeout in seconds"`
+	Timeout       int             `json:"timeout,omitempty" yaml:"timeout,omitempty" desc:"Timeout in seconds for action execution"`
 	Frequency     int             `json:"frequency,omitempty" yaml:"frequency,omitempty" desc:"Action frequency"`
 
 	ScreenOptions
@@ -533,9 +533,9 @@ func WithMaxRetryTimes(maxRetryTimes int) ActionOption {
 	}
 }
 
-func WithTimeout(timeout int) ActionOption {
+func WithTimeout(seconds int) ActionOption {
 	return func(o *ActionOptions) {
-		o.Timeout = timeout
+		o.Timeout = seconds
 	}
 }