Add --reset flag to sandbox exec and reconnection hint on 2nd exec (#484)

Author: robinaugh

cmd/rwx/sandbox.go

@@ -207,6 +207,7 @@ CONFIG FILE
 			Json:           useJson,
 			Sync:           !sandboxNoSync,
 			InitParameters: initParams,
+			Reset:          sandboxReset,
 		})
 		if err != nil {
 			return err
@@ -388,6 +389,7 @@ var (
 	sandboxOpen       bool
 	sandboxWait       bool
 	sandboxNoSync     bool
+	sandboxReset      bool
 	sandboxInitParams []string
 )
 
@@ -411,6 +413,7 @@ func init() {
 	sandboxExecCmd.Flags().StringVar(&sandboxRunID, "id", "", "Use specific run ID")
 	sandboxExecCmd.Flags().BoolVar(&sandboxOpen, "open", false, "Open the run in a browser")
 	sandboxExecCmd.Flags().BoolVar(&sandboxNoSync, "no-sync", false, "Skip syncing local changes before execution")
+	sandboxExecCmd.Flags().BoolVar(&sandboxReset, "reset", false, "Reset the sandbox before executing")
 	sandboxExecCmd.Flags().StringArrayVar(&sandboxInitParams, "init", []string{}, "initialization parameters for the sandbox run, available in the `init` context. Can be specified multiple times")
 
 	// stop flags

internal/cli/sandbox_storage.go

@@ -40,14 +40,15 @@ func DecodeCliState(encoded string) (*CliState, error) {
 }
 
 type SandboxSession struct {
-	RunID       string     `json:"runId"`
-	ConfigFile  string     `json:"configFile"`
-	ScopedToken string     `json:"scopedToken,omitempty"`
-	RunURL      string     `json:"runUrl,omitempty"`
-	ConfigHash  string     `json:"configHash,omitempty"`
-	CreatedAt   *time.Time `json:"createdAt,omitempty"`
-	LastExecAt  *time.Time `json:"lastExecAt,omitempty"`
-	ExecCount   int        `json:"execCount,omitempty"`
+	RunID         string     `json:"runId"`
+	ConfigFile    string     `json:"configFile"`
+	ScopedToken   string     `json:"scopedToken,omitempty"`
+	RunURL        string     `json:"runUrl,omitempty"`
+	ConfigHash    string     `json:"configHash,omitempty"`
+	CreatedAt     *time.Time `json:"createdAt,omitempty"`
+	LastExecAt    *time.Time `json:"lastExecAt,omitempty"`
+	ExecCount     int        `json:"execCount,omitempty"`
+	ResetNagShown bool       `json:"resetNagShown,omitempty"`
 }
 
 // HashConfigFile returns a hex-encoded SHA-256 hash of the file at the given path.

internal/cli/service_sandbox.go

@@ -52,6 +52,7 @@ type ExecSandboxConfig struct {
 	Json           bool
 	Sync           bool
 	InitParameters map[string]string
+	Reset          bool
 }
 
 type ListSandboxesConfig struct {
@@ -434,6 +435,9 @@ func (s Service) ExecSandbox(cfg ExecSandboxConfig) (*ExecSandboxResult, error)
 	var scopedToken string
 	var sessionRunURL string
 	var storedConfigHash string
+	var execCount int
+	var resetNagShown bool
+	isNewSandbox := false
 
 	// Sandbox selection priority:
 	// 1. --id flag
@@ -501,6 +505,8 @@ func (s Service) ExecSandbox(cfg ExecSandboxConfig) (*ExecSandboxResult, error)
 					scopedToken = session.ScopedToken
 					sessionRunURL = session.RunURL
 					storedConfigHash = session.ConfigHash
+					execCount = session.ExecCount
+					resetNagShown = session.ResetNagShown
 				}
 			}
 		} else {
@@ -533,6 +539,8 @@ func (s Service) ExecSandbox(cfg ExecSandboxConfig) (*ExecSandboxResult, error)
 				scopedToken = activeSessions[0].ScopedToken
 				sessionRunURL = activeSessions[0].RunURL
 				storedConfigHash = activeSessions[0].ConfigHash
+				execCount = activeSessions[0].ExecCount
+				resetNagShown = activeSessions[0].ResetNagShown
 				found = true
 			} else if len(activeSessions) > 1 {
 				UnlockSandboxStorage(lockFile)
@@ -606,10 +614,40 @@ func (s Service) ExecSandbox(cfg ExecSandboxConfig) (*ExecSandboxResult, error)
 			}
 		}
 
+		if found && cfg.Reset {
+			connInfo, connErr := s.APIClient.GetSandboxConnectionInfo(runID, scopedToken)
+			if connErr == nil && connInfo.Sandboxable {
+				if sshErr := s.connectSSH(&connInfo); sshErr == nil {
+					_, _ = s.SSHClient.ExecuteCommand("__rwx_sandbox_end__")
+					s.SSHClient.Close()
+				} else {
+					if cancelErr := s.APIClient.CancelRun(runID, scopedToken); cancelErr != nil {
+						fmt.Fprintf(s.Stderr, "Warning: failed to cancel run %s: %v\n", runID, cancelErr)
+					}
+				}
+			} else if connErr == nil && !connInfo.Polling.Completed {
+				if cancelErr := s.APIClient.CancelRun(runID, scopedToken); cancelErr != nil {
+					fmt.Fprintf(s.Stderr, "Warning: failed to cancel run %s: %v\n", runID, cancelErr)
+				}
+			}
+			s.waitForSandboxCompletion(runID, scopedToken)
+			storage.DeleteSession(branch, configFile)
+			if saveErr := storage.Save(); saveErr != nil {
+				fmt.Fprintf(s.Stderr, "Warning: unable to save sandbox sessions: %v\n", saveErr)
+			}
+			found = false
+			runID = ""
+			configFile = ""
+			scopedToken = ""
+			sessionRunURL = ""
+			storedConfigHash = ""
+		}
+
 		if !found {
 			// Pass the lock to StartSandbox so the "no session found → create
 			// new sandbox → persist session" sequence is atomic. StartSandbox
 			// will release it after the initial session is saved.
+			isNewSandbox = true
 			startResult, err := s.StartSandbox(StartSandboxConfig{
 				ConfigFile:     cfgFile,
 				RwxDirectory:   cfg.RwxDirectory,
@@ -637,6 +675,20 @@ func (s Service) ExecSandbox(cfg ExecSandboxConfig) (*ExecSandboxResult, error)
 		}
 	}
 
+	if !isNewSandbox && execCount >= 1 && !resetNagShown {
+		fmt.Fprintf(s.Stderr, "Reconnecting to existing sandbox. To re-run setup tasks, use: rwx sandbox exec --reset -- <command>\n")
+		if nagLock, nagLockErr := s.lockSandboxStorageWithInfo(cfg.Json); nagLockErr == nil {
+			if nagStorage, nagLoadErr := LoadSandboxStorage(); nagLoadErr == nil {
+				if nagSession, ok := nagStorage.GetSession(branch, configFile); ok {
+					nagSession.ResetNagShown = true
+					nagStorage.SetSession(branch, configFile, *nagSession)
+					_ = nagStorage.Save()
+				}
+			}
+			UnlockSandboxStorage(nagLock)
+		}
+	}
+
 	// Warn if the sandbox definition has changed since this sandbox was started
 	if storedConfigHash != "" && configFile != "" {
 		currentHash := HashConfigFile(configFile)