From 9080cd9b781db8f86e395b1be349018cdcde3eb4 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 1 Nov 2023 18:48:53 +0200
Subject: [PATCH] feat(cli/vscodessh): add support for --wait and scripts that
 block login

---
 cli/vscodessh.go | 64 +++++++++++++++++++++++++++++++-----------------
 1 file changed, 41 insertions(+), 23 deletions(-)

diff --git a/cli/vscodessh.go b/cli/vscodessh.go
index 5a61946cdad3d..350cbde2c3652 100644
--- 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/cli/cliutil"
 	"github.com/coder/coder/v2/codersdk"
 )
@@ -38,6 +39,7 @@ func (r *RootCmd) vscodeSSH() *clibase.Cmd {
 		logDir              string
 		networkInfoDir      string
 		networkInfoInterval time.Duration
+		waitEnum            string
 	)
 	cmd := &clibase.Cmd{
 		// A SSH config entry is added by the VS Code extension that
@@ -99,35 +101,45 @@ func (r *RootCmd) vscodeSSH() *clibase.Cmd {
 			}
 			owner := parts[1]
 			name := parts[2]
+			if len(parts) > 3 {
+				name += "." + parts[3]
+			}
+
+			// Set autostart to false because it's assumed the VS Code extension
+			// will call this command after the workspace is started.
+			autostart := false
 
-			workspace, err := client.WorkspaceByOwnerAndName(ctx, owner, name, codersdk.WorkspaceOptions{})
+			_, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, autostart, 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
 				}
 			}
 
@@ -152,7 +164,7 @@ func (r *RootCmd) vscodeSSH() *clibase.Cmd {
 			if r.disableDirect {
 				logger.Info(ctx, "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,
 			})
@@ -249,6 +261,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
 }