fix: enhance studio environment management with batch removal and stability checks

Author: Code2Life

cmd/ggo/studio/studio.go

@@ -81,7 +81,10 @@ Examples:
   ggo studio stop my-studio
 
   # Remove an environment
-  ggo studio rm my-studio`,
+  ggo studio rm my-studio
+
+  # Batch remove all environments
+  ggo studio rm --all -f`,
 		PersistentPreRun: func(cmd *cobra.Command, args []string) {
 			// Initialize klog flags if not already initialized
 			klog.InitFlags(nil)
@@ -667,18 +670,61 @@ func newStopCmd() *cobra.Command {
 
 func newRemoveCmd() *cobra.Command {
 	var force bool
+	var all bool
 
 	cmd := &cobra.Command{
 		Use:     "rm <name>",
-		Short:   "Remove a studio environment",
+		Short:   "Remove studio environment(s)",
 		Aliases: []string{"remove", "delete"},
-		Args:    cobra.ExactArgs(1),
+		Args: func(cmd *cobra.Command, args []string) error {
+			if all {
+				if len(args) > 0 {
+					return fmt.Errorf("--all does not accept a studio name")
+				}
+				return nil
+			}
+			return cobra.ExactArgs(1)(cmd, args)
+		},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			ctx := context.Background()
 			mgr := getManager()
 			out := getOutput()
-			name := args[0]
 
+			if all {
+				if !force && !out.IsJSON() {
+					styles := tui.DefaultStyles()
+					fmt.Printf("%s Are you sure you want to remove ALL studio environments? [y/N]: ", styles.Warning.Render("!"))
+					var confirm string
+					fmt.Scanln(&confirm)
+					if confirm != "y" && confirm != "Y" {
+						out.Info("Cancelled")
+						return nil
+					}
+				}
+
+				removedNames, err := mgr.RemoveAll(ctx)
+				if err != nil {
+					cmd.SilenceUsage = true
+					return err
+				}
+				if len(removedNames) == 0 {
+					out.Info("No studio environments found")
+					return nil
+				}
+				for _, removedName := range removedNames {
+					if err := mgr.RemoveSSHConfig(removedName); err != nil {
+						klog.Warningf("Failed to remove SSH config for %s: error=%v", removedName, err)
+					}
+				}
+
+				return out.Render(&cmdutil.ActionData{
+					Success: true,
+					Message: fmt.Sprintf("Removed %d studio environment(s)", len(removedNames)),
+					ID:      "all",
+				})
+			}
+
+			name := args[0]
 			if !force && !out.IsJSON() {
 				styles := tui.DefaultStyles()
 				fmt.Printf("%s Are you sure you want to remove environment %s? [y/N]: ",
@@ -710,6 +756,7 @@ func newRemoveCmd() *cobra.Command {
 	}
 
 	cmd.Flags().BoolVarP(&force, "force", "f", false, "Force remove")
+	cmd.Flags().BoolVar(&all, "all", false, "Remove all studio environments in one batch")
 	return cmd
 }
 

internal/studio/backend_apple.go

@@ -464,7 +464,7 @@ func extractSSHPort(published []appleContainerPublishedPort) int {
 			}
 		}
 	}
-	return 22
+	return 0
 }
 
 func normalizeContainerMemory(memory string) string {

internal/studio/backend_docker.go

@@ -455,7 +455,7 @@ func (b *DockerBackend) Get(ctx context.Context, idOrName string) (*Environment,
 	}
 
 	// Parse SSH port
-	sshPort := 22
+	sshPort := 0
 	if ports, ok := c.NetworkSettings.Ports["22/tcp"]; ok && len(ports) > 0 {
 		if p, err := strconv.Atoi(ports[0].HostPort); err == nil {
 			sshPort = p
@@ -562,7 +562,7 @@ func parseSSHPort(ports string) int {
 			}
 		}
 	}
-	return 22
+	return 0
 }
 
 // findAvailablePort finds an available port in the SSH port range (12000-18000)

internal/studio/backend_ssh_port_ginkgo_test.go

@@ -0,0 +1,32 @@
+package studio
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("SSH port parsing", func() {
+	Describe("parseSSHPort", func() {
+		It("returns mapped host SSH port when 22/tcp is published", func() {
+			port := parseSSHPort("0.0.0.0:15432->22/tcp")
+			Expect(port).To(Equal(15432))
+		})
+
+		It("returns zero when no SSH port mapping exists", func() {
+			Expect(parseSSHPort("0.0.0.0:8888->8888/tcp")).To(BeZero())
+			Expect(parseSSHPort("")).To(BeZero())
+		})
+	})
+
+	Describe("extractSSHPort", func() {
+		It("returns mapped host SSH port for Apple Container", func() {
+			port := extractSSHPort([]appleContainerPublishedPort{{HostPort: 17000, ContainerPort: 22, Proto: "tcp", Count: 1}})
+			Expect(port).To(Equal(17000))
+		})
+
+		It("returns zero when Apple Container has no SSH mapping", func() {
+			port := extractSSHPort([]appleContainerPublishedPort{{HostPort: 8080, ContainerPort: 8080, Proto: "tcp", Count: 1}})
+			Expect(port).To(BeZero())
+		})
+	})
+})

internal/studio/manager.go

@@ -4,18 +4,27 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"net"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime"
 	"sort"
+	"strconv"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/NexusGPU/gpu-go/internal/errors"
 	"github.com/NexusGPU/gpu-go/internal/platform"
 )
 
+var (
+	createStabilityWindow       = 3 * time.Second
+	createStabilityPollInterval = 200 * time.Millisecond
+	createSSHProbeTimeout       = 500 * time.Millisecond
+)
+
 // Manager manages AI studio environments across different backends
 type Manager struct {
 	paths    *platform.Paths
@@ -181,6 +190,12 @@ func (m *Manager) Create(ctx context.Context, opts *CreateOptions) (*Environment
 		return nil, err
 	}
 
+	if err := m.waitForStableRunning(ctx, backend, env); err != nil {
+		return nil, err
+	}
+
+	m.clearUnreachableSSH(ctx, env)
+
 	// Save environment to local state
 	if err := m.saveEnvironment(env); err != nil {
 		// Log but don't fail
@@ -345,6 +360,49 @@ func (m *Manager) Remove(ctx context.Context, idOrName string) error {
 	return m.removeEnvironment(env.ID)
 }
 
+// RemoveAll removes all known environments. When runtimes are offline, stale state entries are still cleaned up.
+func (m *Manager) RemoveAll(ctx context.Context) ([]string, error) {
+	envs, err := m.List(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	seen := make(map[string]struct{}, len(envs))
+	removed := make([]string, 0, len(envs))
+	failed := make([]string, 0)
+
+	for _, env := range envs {
+		if env == nil || env.ID == "" {
+			continue
+		}
+		if _, ok := seen[env.ID]; ok {
+			continue
+		}
+		seen[env.ID] = struct{}{}
+
+		switch env.Status {
+		case StatusUnknown, StatusDeleted:
+			if err := m.removeEnvironment(env.ID); err != nil {
+				failed = append(failed, fmt.Sprintf("%s (%v)", env.Name, err))
+				continue
+			}
+			removed = append(removed, env.Name)
+		default:
+			if err := m.Remove(ctx, env.ID); err != nil {
+				failed = append(failed, fmt.Sprintf("%s (%v)", env.Name, err))
+				continue
+			}
+			removed = append(removed, env.Name)
+		}
+	}
+
+	if len(failed) > 0 {
+		return removed, fmt.Errorf("failed to remove %d environment(s): %s", len(failed), strings.Join(failed, "; "))
+	}
+
+	return removed, nil
+}
+
 // AddSSHConfig adds an SSH config entry for an environment
 func (m *Manager) AddSSHConfig(env *Environment) error {
 	if env.SSHHost == "" || env.SSHPort == 0 {
@@ -447,6 +505,109 @@ func (m *Manager) getSSHConfigPath() string {
 	return filepath.Join(home, ".ssh", "config")
 }
 
+func (m *Manager) waitForStableRunning(ctx context.Context, backend Backend, created *Environment) error {
+	if createStabilityWindow <= 0 {
+		return nil
+	}
+	if created == nil || created.ID == "" {
+		return fmt.Errorf("studio creation returned invalid environment metadata")
+	}
+
+	deadline := time.Now().Add(createStabilityWindow)
+	seenRunning := false
+	var lastStatus EnvironmentStatus
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
+		env, err := backend.Get(ctx, created.ID)
+		if err != nil {
+			return fmt.Errorf(
+				"studio '%s' created but status check failed within %s: %w. Check status with `ggo studio list` and logs with `ggo studio logs %s`",
+				created.Name,
+				createStabilityWindow,
+				err,
+				created.Name,
+			)
+		}
+		lastStatus = env.Status
+
+		switch env.Status {
+		case StatusRunning:
+			seenRunning = true
+		case StatusPending, StatusStarting, StatusPulling:
+			if seenRunning {
+				return fmt.Errorf(
+					"studio '%s' failed to stay running for at least %s (status=%s). Check status with `ggo studio list` and logs with `ggo studio logs %s`",
+					created.Name,
+					createStabilityWindow,
+					env.Status,
+					created.Name,
+				)
+			}
+		default:
+			return fmt.Errorf(
+				"studio '%s' failed to stay running for at least %s (status=%s). Check status with `ggo studio list` and logs with `ggo studio logs %s`",
+				created.Name,
+				createStabilityWindow,
+				env.Status,
+				created.Name,
+			)
+		}
+
+		remaining := time.Until(deadline)
+		if remaining <= 0 {
+			if seenRunning {
+				return nil
+			}
+			return fmt.Errorf(
+				"studio '%s' did not reach running status within %s (last status=%s). Check status with `ggo studio list` and logs with `ggo studio logs %s`",
+				created.Name,
+				createStabilityWindow,
+				lastStatus,
+				created.Name,
+			)
+		}
+
+		sleepFor := createStabilityPollInterval
+		if sleepFor <= 0 || sleepFor > remaining {
+			sleepFor = remaining
+		}
+
+		timer := time.NewTimer(sleepFor)
+		select {
+		case <-ctx.Done():
+			timer.Stop()
+			return ctx.Err()
+		case <-timer.C:
+		}
+	}
+}
+
+func (m *Manager) clearUnreachableSSH(ctx context.Context, env *Environment) {
+	if env == nil || env.SSHHost == "" || env.SSHPort <= 0 {
+		return
+	}
+
+	dialCtx, cancel := context.WithTimeout(ctx, createSSHProbeTimeout)
+	defer cancel()
+
+	dialer := &net.Dialer{}
+	address := net.JoinHostPort(env.SSHHost, strconv.Itoa(env.SSHPort))
+	conn, err := dialer.DialContext(dialCtx, "tcp", address)
+	if err != nil {
+		env.SSHHost = ""
+		env.SSHPort = 0
+		env.SSHUser = ""
+		return
+	}
+	_ = conn.Close()
+}
+
 // State management
 
 func (m *Manager) getStatePath() string {