refactor: enhance context cancellation handling in mobile UI and driver extensions

step_ui.go

@@ -925,18 +925,18 @@ func runStepMobileUI(s *SessionRunner, step IStep) (stepResult *StepResult, err
 			}
 
 			// call MCP tool to execute action with cancellable context
-			ctx, cancel := context.WithCancel(context.Background())
+			ctx, cancel := context.WithCancelCause(context.Background())
-			defer cancel()
+			defer cancel(nil)
 
-			// Create a goroutine to monitor for interrupt signals
+			// Create a goroutine to monitor for interrupt signals and timeouts
 			go func() {
 				select {
 				case <-s.caseRunner.hrpRunner.interruptSignal:
 					log.Warn().Msg("cancelling action due to interrupt signal")
-					cancel()
+					cancel(code.InterruptError)
 				case <-s.caseRunner.hrpRunner.caseTimeoutTimer.C:
 					log.Warn().Msg("cancelling action due to case timeout")
-					cancel()
+					cancel(code.TimeoutError)
 				case <-ctx.Done():
 					// Context already cancelled
 				}

uixt/driver_ext_ai.go

@@ -26,14 +26,17 @@ func (dExt *XTDriver) StartToGoal(ctx context.Context, prompt string, opts ...op
 		attempt++
 		log.Info().Int("attempt", attempt).Msg("planning attempt")
 
-		// Check for context cancellation (interrupt signal)
+		// Check for context cancellation (interrupt signal or timeout)
 		select {
 		case <-ctx.Done():
+			cause := context.Cause(ctx)
 			log.Warn().
 				Int("attempt", attempt).
 				Int("completed_plannings", len(allPlannings)).
-				Msg("interrupted in StartToGoal")
+				Err(cause).
-			return allPlannings, errors.Wrap(code.InterruptError, "StartToGoal interrupted")
+				Msg("StartToGoal cancelled")
+			// Return the specific error type based on the cancellation cause
+			return allPlannings, errors.Wrap(cause, "StartToGoal cancelled")
 		default:
 		}
 
@@ -85,15 +88,18 @@ func (dExt *XTDriver) StartToGoal(ctx context.Context, prompt string, opts ...op
 			// Check for context cancellation before each action
 			select {
 			case <-ctx.Done():
+				cause := context.Cause(ctx)
 				log.Warn().
 					Int("attempt", attempt).
 					Int("completed_plannings", len(allPlannings)).
 					Int("completed_tool_calls", len(planningResult.SubActions)).
 					Int("total_tool_calls", len(planningResult.ToolCalls)).
-					Msg("interrupted in invokeToolCalls")
+					Err(cause).
+					Msg("invokeToolCalls cancelled")
 				planningResult.Elapsed = time.Since(planningStartTime).Milliseconds()
 				allPlannings = append(allPlannings, planningResult)
-				return allPlannings, errors.Wrap(code.InterruptError, "invokeToolCalls interrupted")
+				// Return the specific error type based on the cancellation cause
+				return allPlannings, errors.Wrap(cause, "invokeToolCalls cancelled")
 			default:
 			}