fix: handle vscodessh style workspace names in coder ssh

f0ssel · f0ssel · commit 0154e6cd10c3 · 2025-03-28T16:03:43.000Z
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -13,6 +13,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"regexp"
 	"slices"
 	"strconv"
 	"strings"
@@ -57,6 +58,7 @@ var (
 	autostopNotifyCountdown = []time.Duration{30 * time.Minute}
 	// gracefulShutdownTimeout is the timeout, per item in the stack of things to close
 	gracefulShutdownTimeout = 2 * time.Second
+	workspaceNameRe         = regexp.MustCompile(`[/.]+|--`)
 )
 
 func (r *RootCmd) ssh() *serpent.Command {
@@ -200,10 +202,9 @@ func (r *RootCmd) ssh() *serpent.Command {
 				parsedEnv = append(parsedEnv, [2]string{k, v})
 			}
 
-			namedWorkspace := strings.TrimPrefix(inv.Args[0], hostPrefix)
-			// Support "--" as a delimiter between owner and workspace name
-			namedWorkspace = strings.Replace(namedWorkspace, "--", "/", 1)
-
+			workspaceInput := strings.TrimPrefix(inv.Args[0], hostPrefix)
+			// convert workspace name format into owner/workspace.agent
+			namedWorkspace := normalizeWorkspaceInput(workspaceInput)
 			workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, !disableAutostart, namedWorkspace)
 			if err != nil {
 				return err
@@ -1413,3 +1414,28 @@ func collectNetworkStats(ctx context.Context, agentConn *workspacesdk.AgentConn,
 		DownloadBytesSec: int64(downloadSecs),
 	}, nil
 }
+
+// Converts workspace name input to owner/workspace.agent format
+// Possible valid input formats:
+// workspace
+// owner/workspace
+// owner--workspace
+// owner/workspace--agent
+// owner/workspace.agent
+// owner--workspace--agent
+// owner--workspace.agent
+func normalizeWorkspaceInput(input string) string {
+	// Split on "/", "--", and "."
+	parts := workspaceNameRe.Split(input, -1)
+
+	switch len(parts) {
+	case 1:
+		return input // "workspace"
+	case 2:
+		return fmt.Sprintf("%s/%s", parts[0], parts[1]) // "owner/workspace"
+	case 3:
+		return fmt.Sprintf("%s/%s.%s", parts[0], parts[1], parts[2]) // "owner/workspace.agent"
+	default:
+		return input // Fallback
+	}
+}
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -63,8 +63,11 @@ func setupWorkspaceForAgent(t *testing.T, mutations ...func([]*proto.Agent) []*p
 	client, store := coderdtest.NewWithDatabase(t, nil)
 	client.SetLogger(testutil.Logger(t).Named("client"))
 	first := coderdtest.CreateFirstUser(t, client)
-	userClient, user := coderdtest.CreateAnotherUser(t, client, first.OrganizationID)
+	userClient, user := coderdtest.CreateAnotherUserMutators(t, client, first.OrganizationID, nil, func(r *codersdk.CreateUserRequestWithOrgs) {
+		r.Username = "myuser"
+	})
 	r := dbfake.WorkspaceBuild(t, store, database.WorkspaceTable{
+		Name:           "myworkspace",
 		OrganizationID: first.OrganizationID,
 		OwnerID:        user.ID,
 	}).WithAgent(mutations...).Do()
@@ -98,6 +101,46 @@ func TestSSH(t *testing.T) {
 		pty.WriteLine("exit")
 		<-cmdDone
 	})
+	t.Run("WorkspaceNameInput", func(t *testing.T) {
+		t.Parallel()
+
+		cases := []string{
+			"myworkspace",
+			"myuser/myworkspace",
+			"myuser--myworkspace",
+			"myuser/myworkspace--dev",
+			"myuser/myworkspace.dev",
+			"myuser--myworkspace--dev",
+			"myuser--myworkspace.dev",
+		}
+
+		for _, tc := range cases {
+			t.Run(tc, func(t *testing.T) {
+				t.Parallel()
+				ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+				defer cancel()
+
+				client, workspace, agentToken := setupWorkspaceForAgent(t)
+
+				inv, root := clitest.New(t, "ssh", tc)
+				clitest.SetupConfig(t, client, root)
+				pty := ptytest.New(t).Attach(inv)
+
+				cmdDone := tGo(t, func() {
+					err := inv.WithContext(ctx).Run()
+					assert.NoError(t, err)
+				})
+				pty.ExpectMatch("Waiting")
+
+				_ = agenttest.New(t, client.URL, agentToken)
+				coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)
+
+				// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
+				pty.WriteLine("exit")
+				<-cmdDone
+			})
+		}
+	})
 	t.Run("StartStoppedWorkspace", func(t *testing.T) {
 		t.Parallel()
 
