coder · dwahler · May 17, 2022 · May 9, 2022 · May 9, 2022 · May 9, 2022
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -24,20 +24,22 @@ import (
 	"github.com/coder/coder/coderd/autobuild/schedule"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/cryptorand"
 )
 
 var autostopPollInterval = 30 * time.Second
 var autostopNotifyCountdown = []time.Duration{30 * time.Minute}
 
 func ssh() *cobra.Command {
 	var (
-		stdio bool
+		stdio   bool
+		shuffle bool
 	)
 	cmd := &cobra.Command{
 		Annotations: workspaceCommand,
 		Use:         "ssh <workspace>",
 		Short:       "SSH into a workspace",
-		Args:        cobra.MinimumNArgs(1),
+		Args:        cobra.ArbitraryArgs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			client, err := createClient(cmd)
 			if err != nil {
@@ -48,10 +50,38 @@ func ssh() *cobra.Command {
 				return err
 			}
 
-			workspaceParts := strings.Split(args[0], ".")
-			workspace, err := client.WorkspaceByOwnerAndName(cmd.Context(), organization.ID, codersdk.Me, workspaceParts[0])
-			if err != nil {
-				return err
+			var workspace codersdk.Workspace
+			var workspaceParts []string
+			if shuffle {
+				err := cobra.ExactArgs(0)(cmd, args)
+				if err != nil {
+					return err
+				}
+
+				workspaces, err := client.WorkspacesByOwner(cmd.Context(), organization.ID, codersdk.Me)
+				if err != nil {
+					return err
+				}
+				if len(workspaces) == 0 {
+					return xerrors.New("no workspaces to shuffle")
+				}
+
+				idx, err := cryptorand.Intn(len(workspaces))
+				if err != nil {
+					return err
+				}
+				workspace = workspaces[idx]
+			} else {
+				err := cobra.MinimumNArgs(1)(cmd, args)
+				if err != nil {
+					return err
+				}
+
+				workspaceParts = strings.Split(args[0], ".")
+				workspace, err = client.WorkspaceByOwnerAndName(cmd.Context(), organization.ID, codersdk.Me, workspaceParts[0])
+				if err != nil {
+					return err
+				}
 			}
 
 			if workspace.LatestBuild.Transition != database.WorkspaceTransitionStart {
@@ -96,9 +126,17 @@ func ssh() *cobra.Command {
 			}
 			if agent.ID == uuid.Nil {
 				if len(agents) > 1 {
-					return xerrors.New("you must specify the name of an agent")
+					if !shuffle {
+						return xerrors.New("you must specify the name of an agent")
+					}
+					idx, err := cryptorand.Intn(len(agents))
+					if err != nil {
+						return err
+					}
+					agent = agents[idx]
+				} else {
+					agent = agents[0]
 				}
-				agent = agents[0]
 			}
 			// OpenSSH passes stderr directly to the calling TTY.
 			// This is required in "stdio" mode so a connecting indicator can be displayed.
@@ -189,6 +227,8 @@ func ssh() *cobra.Command {
 		},
 	}
 	cliflag.BoolVarP(cmd.Flags(), &stdio, "stdio", "", "CODER_SSH_STDIO", false, "Specifies whether to emit SSH output over stdin/stdout.")
+	cliflag.BoolVarP(cmd.Flags(), &shuffle, "shuffle", "", "CODER_SSH_SHUFFLE", false, "Specifies whether to choose a random workspace")
+	_ = cmd.Flags().MarkHidden("shuffle")
 
 	return cmd
 }