feat(cli/vscodessh): add support for --wait and scripts that block login

mafredri · mafredri · commit 39aec7a280ff · 2023-11-01T18:48:53.000+02:00
diff --git a/cli/vscodessh.go b/cli/vscodessh.go
@@ -21,6 +21,7 @@ import (
 	"cdr.dev/slog/sloggers/sloghuman"
 
 	"github.com/coder/coder/v2/cli/clibase"
+	"github.com/coder/coder/v2/cli/cliui"
 	"github.com/coder/coder/v2/codersdk"
 )
 
@@ -36,6 +37,7 @@ func (r *RootCmd) vscodeSSH() *clibase.Cmd {
 		urlFile             string
 		networkInfoDir      string
 		networkInfoInterval time.Duration
+		waitEnum            string
 	)
 	cmd := &clibase.Cmd{
 		// A SSH config entry is added by the VS Code extension that
@@ -97,35 +99,41 @@ func (r *RootCmd) vscodeSSH() *clibase.Cmd {
 			}
 			owner := parts[1]
 			name := parts[2]
+			if len(parts) > 3 {
+				name += "." + parts[3]
+			}
 
-			workspace, err := client.WorkspaceByOwnerAndName(ctx, owner, name, codersdk.WorkspaceOptions{})
+			_, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, owner, name)
 			if err != nil {
-				return xerrors.Errorf("find workspace: %w", err)
+				return xerrors.Errorf("find workspace and agent: %w", err)
 			}
 
-			var agent codersdk.WorkspaceAgent
-			var found bool
-			for _, resource := range workspace.LatestBuild.Resources {
-				if len(resource.Agents) == 0 {
-					continue
-				}
-				for _, resourceAgent := range resource.Agents {
-					// If an agent name isn't included we default to
-					// the first agent!
-					if len(parts) != 4 {
-						agent = resourceAgent
-						found = true
+			// Select the startup script behavior based on template configuration or flags.
+			var wait bool
+			switch waitEnum {
+			case "yes":
+				wait = true
+			case "no":
+				wait = false
+			case "auto":
+				for _, script := range workspaceAgent.Scripts {
+					if script.StartBlocksLogin {
+						wait = true
 						break
 					}
-					if resourceAgent.Name != parts[3] {
-						continue
-					}
-					agent = resourceAgent
-					found = true
-					break
 				}
-				if found {
-					break
+			default:
+				return xerrors.Errorf("unknown wait value %q", waitEnum)
+			}
+
+			err = cliui.Agent(ctx, inv.Stderr, workspaceAgent.ID, cliui.AgentOptions{
+				Fetch:     client.WorkspaceAgent,
+				FetchLogs: client.WorkspaceAgentLogsAfter,
+				Wait:      wait,
+			})
+			if err != nil {
+				if xerrors.Is(err, context.Canceled) {
+					return cliui.Canceled
 				}
 			}
 
@@ -137,7 +145,7 @@ func (r *RootCmd) vscodeSSH() *clibase.Cmd {
 			if r.disableDirect {
 				_, _ = fmt.Fprintln(inv.Stderr, "Direct connections disabled.")
 			}
-			agentConn, err := client.DialWorkspaceAgent(ctx, agent.ID, &codersdk.DialWorkspaceAgentOptions{
+			agentConn, err := client.DialWorkspaceAgent(ctx, workspaceAgent.ID, &codersdk.DialWorkspaceAgentOptions{
 				Logger:         logger,
 				BlockEndpoints: r.disableDirect,
 			})
@@ -229,6 +237,12 @@ func (r *RootCmd) vscodeSSH() *clibase.Cmd {
 			Default:     "5s",
 			Value:       clibase.DurationOf(&networkInfoInterval),
 		},
+		{
+			Flag:        "wait",
+			Description: "Specifies whether or not to wait for the startup script to finish executing. Auto means that the agent startup script behavior configured in the workspace template is used.",
+			Default:     "auto",
+			Value:       clibase.EnumOf(&waitEnum, "yes", "no", "auto"),
+		},
 	}
 	return cmd
 }
