feat(cli): Add wait and no-wait support to ssh

mafredri · mafredri · commit ffbdf29a7cbb · 2023-06-07T12:47:23.000Z
Fixes #7768 Refs #7893
diff --git a/cli/cliui/agent.go b/cli/cliui/agent.go
@@ -27,7 +27,7 @@ type AgentOptions struct {
 	Fetch         func(context.Context) (codersdk.WorkspaceAgent, error)
 	FetchInterval time.Duration
 	WarnInterval  time.Duration
-	NoWait        bool // If true, don't wait for the agent to be ready.
+	Wait          bool // If true, wait for the agent to be ready (startup script).
 }
 
 // Agent displays a spinning indicator that waits for a workspace agent to connect.
@@ -96,7 +96,7 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
 	// we do this just before starting the spinner to avoid needless
 	// spinning.
 	if agent.Status == codersdk.WorkspaceAgentConnected &&
-		agent.StartupScriptBehavior == codersdk.WorkspaceAgentStartupScriptBehaviorBlocking && opts.NoWait {
+		agent.StartupScriptBehavior == codersdk.WorkspaceAgentStartupScriptBehaviorBlocking && !opts.Wait {
 		showMessage()
 		return nil
 	}
@@ -140,7 +140,7 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
 			// NOTE(mafredri): Once we have access to the workspace agent's
 			// startup script logs, we can show them here.
 			// https://github.com/coder/coder/issues/2957
-			if agent.StartupScriptBehavior == codersdk.WorkspaceAgentStartupScriptBehaviorBlocking && !opts.NoWait {
+			if agent.StartupScriptBehavior == codersdk.WorkspaceAgentStartupScriptBehaviorBlocking && opts.Wait {
 				switch agent.LifecycleState {
 				case codersdk.WorkspaceAgentLifecycleReady:
 					return nil
@@ -183,7 +183,7 @@ func waitingMessage(agent codersdk.WorkspaceAgent, opts AgentOptions) (m *messag
 		Prompt: "Don't panic, your workspace is booting up!",
 	}
 	defer func() {
-		if agent.Status == codersdk.WorkspaceAgentConnected && opts.NoWait {
+		if agent.Status == codersdk.WorkspaceAgentConnected && !opts.Wait {
 			m.Spin = ""
 		}
 		if m.Spin != "" {
@@ -225,7 +225,7 @@ func waitingMessage(agent codersdk.WorkspaceAgent, opts AgentOptions) (m *messag
 	case codersdk.WorkspaceAgentConnected:
 		m.Spin = fmt.Sprintf("Waiting for %s to become ready...", DefaultStyles.Field.Render(agent.Name))
 		m.Prompt = "Don't panic, your workspace agent has connected and the workspace is getting ready!"
-		if opts.NoWait {
+		if !opts.Wait {
 			m.Prompt = "Your workspace is still getting ready, it may be in an incomplete state."
 		}
 
diff --git a/cli/cliui/agent_test.go b/cli/cliui/agent_test.go
@@ -140,7 +140,7 @@ func TestAgent_StartupTimeout(t *testing.T) {
 				},
 				FetchInterval: time.Millisecond,
 				WarnInterval:  time.Millisecond,
-				NoWait:        false,
+				Wait:          true,
 			})
 			return err
 		},
@@ -199,7 +199,7 @@ func TestAgent_StartErrorExit(t *testing.T) {
 				},
 				FetchInterval: time.Millisecond,
 				WarnInterval:  60 * time.Second,
-				NoWait:        false,
+				Wait:          true,
 			})
 			return err
 		},
@@ -255,7 +255,7 @@ func TestAgent_NoWait(t *testing.T) {
 				},
 				FetchInterval: time.Millisecond,
 				WarnInterval:  time.Second,
-				NoWait:        true,
+				Wait:          false,
 			})
 			return err
 		},
@@ -325,7 +325,7 @@ func TestAgent_StartupScriptBehaviorNonBlocking(t *testing.T) {
 				},
 				FetchInterval: time.Millisecond,
 				WarnInterval:  time.Second,
-				NoWait:        false,
+				Wait:          true,
 			})
 			return err
 		},
diff --git a/cli/portforward.go b/cli/portforward.go
@@ -92,6 +92,7 @@ func (r *RootCmd) portForward() *clibase.Cmd {
 				Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
 					return client.WorkspaceAgent(ctx, workspaceAgent.ID)
 				},
+				Wait: false,
 			})
 			if err != nil {
 				return xerrors.Errorf("await agent: %w", err)
diff --git a/cli/speedtest.go b/cli/speedtest.go
@@ -45,6 +45,7 @@ func (r *RootCmd) speedtest() *clibase.Cmd {
 				Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
 					return client.WorkspaceAgent(ctx, workspaceAgent.ID)
 				},
+				Wait: false,
 			})
 			if err != nil && !xerrors.Is(err, cliui.AgentStartError) {
 				return xerrors.Errorf("await agent: %w", err)
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -49,6 +49,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 		forwardGPG     bool
 		identityAgent  string
 		wsPollInterval time.Duration
+		wait           bool
 		noWait         bool
 		logDir         string
 		logToFile      bool
@@ -66,6 +67,10 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			ctx, cancel := context.WithCancel(inv.Context())
 			defer cancel()
 
+			if wait && noWait {
+				return xerrors.New("cannot specify both --wait and --no-wait")
+			}
+
 			logger := slog.Make() // empty logger
 			defer func() {
 				if retErr != nil {
@@ -105,6 +110,18 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 				return err
 			}
 
+			// Select the startup script behavior based on template configuration or flags.
+			if !wait && !noWait {
+				switch workspaceAgent.StartupScriptBehavior {
+				case codersdk.WorkspaceAgentStartupScriptBehaviorBlocking:
+					wait = true
+				case codersdk.WorkspaceAgentStartupScriptBehaviorNonBlocking:
+					wait = false
+				}
+			} else if noWait {
+				wait = false
+			}
+
 			templateVersion, err := client.TemplateVersion(ctx, workspace.LatestBuild.TemplateVersionID)
 			if err != nil {
 				return err
@@ -134,7 +151,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 				Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
 					return client.WorkspaceAgent(ctx, workspaceAgent.ID)
 				},
-				NoWait: noWait,
+				Wait: wait,
 			})
 			if err != nil {
 				if xerrors.Is(err, context.Canceled) {
@@ -347,10 +364,16 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			Default:     "1m",
 			Value:       clibase.DurationOf(&wsPollInterval),
 		},
+		{
+			Flag:        "wait",
+			Env:         "CODER_SSH_WAIT",
+			Description: "Wait for the the startup script to finish executing. This is the default if the template has configured the agent startup script behavior as blocking. Can not be used together with --no-wait.",
+			Value:       clibase.BoolOf(&wait),
+		},
 		{
 			Flag:        "no-wait",
 			Env:         "CODER_SSH_NO_WAIT",
-			Description: "Specifies whether to wait for a workspace to become ready before logging in (only applicable when the startup script behavior is blocking). Note that the workspace agent may still be in the process of executing the startup script and the workspace may be in an incomplete state.",
+			Description: "Enter workspace immediately after the agent has connected. This is the default if the template has configured the agent startup script behavior as non-blocking. Can not be used together with --wait.",
 			Value:       clibase.BoolOf(&noWait),
 		},
 		{
