feat(cli): Add support for `delay_login_until_ready` (#5851)

cli/cliui/agent.go

@@ -10,16 +10,21 @@ import (
 	"time"
 
 	"github.com/briandowns/spinner"
+	"github.com/muesli/reflow/indent"
+	"github.com/muesli/reflow/wordwrap"
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/codersdk"
 )
 
+var AgentStartError = xerrors.New("agent startup exited with non-zero exit status")
+
 type AgentOptions struct {
 	WorkspaceName string
 	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.
 }
 
 // Agent displays a spinning indicator that waits for a workspace agent to connect.
@@ -36,48 +41,33 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
 		return xerrors.Errorf("fetch: %w", err)
 	}
 
-	if agent.Status == codersdk.WorkspaceAgentConnected {
+	// Fast path if the agent is ready (avoid showing connecting prompt).
+	// We don't take the fast path for opts.NoWait yet because we want to
+	// show the message.
+	if agent.Status == codersdk.WorkspaceAgentConnected &&
+		(!agent.DelayLoginUntilReady || agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady) {
 		return nil
 	}
 
+	ctx, cancel := signal.NotifyContext(ctx, os.Interrupt)
+	defer cancel()
+
 	spin := spinner.New(spinner.CharSets[78], 100*time.Millisecond, spinner.WithColor("fgHiGreen"))
 	spin.Writer = writer
 	spin.ForceOutput = true
-	spin.Suffix = " Waiting for connection from " + Styles.Field.Render(agent.Name) + "..."
+	spin.Suffix = waitingMessage(agent, opts).Spin
-	spin.Start()
-	defer spin.Stop()
 
-	ctx, cancelFunc := context.WithCancel(ctx)
+	waitMessage := &message{}
-	defer cancelFunc()
-	stopSpin := make(chan os.Signal, 1)
-	signal.Notify(stopSpin, os.Interrupt)
-	defer signal.Stop(stopSpin)
-	go func() {
-		select {
-		case <-ctx.Done():
-			return
-		case <-stopSpin:
-		}
-		cancelFunc()
-		signal.Stop(stopSpin)
-		spin.Stop()
-		// nolint:revive
-		os.Exit(1)
-	}()
-
-	var waitMessage string
-	messageAfter := time.NewTimer(opts.WarnInterval)
-	defer messageAfter.Stop()
 	showMessage := func() {
 		resourceMutex.Lock()
 		defer resourceMutex.Unlock()
 
-		m := waitingMessage(agent)
+		m := waitingMessage(agent, opts)
-		if m == waitMessage {
+		if m.Prompt == waitMessage.Prompt {
 			return
 		}
 		moveUp := ""
-		if waitMessage != "" {
+		if waitMessage.Prompt != "" {
 			// If this is an update, move a line up
 			// to keep it tidy and aligned.
 			moveUp = "\033[1A"
@@ -86,20 +76,43 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
 
 		// Stop the spinner while we write our message.
 		spin.Stop()
+		spin.Suffix = waitMessage.Spin
 		// Clear the line and (if necessary) move up a line to write our message.
-		_, _ = fmt.Fprintf(writer, "\033[2K%s%s\n\n", moveUp, Styles.Paragraph.Render(Styles.Prompt.String()+waitMessage))
+		_, _ = fmt.Fprintf(writer, "\033[2K%s\n%s\n", moveUp, waitMessage.Prompt)
 		select {
 		case <-ctx.Done():
 		default:
 			// Safe to resume operation.
-			spin.Start()
+			if spin.Suffix != "" {
+				spin.Start()
+			}
 		}
 	}
+
+	// Fast path for showing the error message even when using no wait,
+	// we do this just before starting the spinner to avoid needless
+	// spinning.
+	if agent.Status == codersdk.WorkspaceAgentConnected &&
+		agent.DelayLoginUntilReady && opts.NoWait {
+		showMessage()
+		return nil
+	}
+
+	// Start spinning after fast paths are handled.
+	if spin.Suffix != "" {
+		spin.Start()
+	}
+	defer spin.Stop()
+
+	warnAfter := time.NewTimer(opts.WarnInterval)
+	defer warnAfter.Stop()
+	warningShown := make(chan struct{})
 	go func() {
 		select {
 		case <-ctx.Done():
-		case <-messageAfter.C:
+			close(warningShown)
-			messageAfter.Stop()
+		case <-warnAfter.C:
+			close(warningShown)
 			showMessage()
 		}
 	}()
@@ -121,6 +134,29 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
 		resourceMutex.Unlock()
 		switch agent.Status {
 		case codersdk.WorkspaceAgentConnected:
+			// 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.DelayLoginUntilReady && !opts.NoWait {
+				switch agent.LifecycleState {
+				case codersdk.WorkspaceAgentLifecycleReady:
+					return nil
+				case codersdk.WorkspaceAgentLifecycleStartTimeout:
+					showMessage()
+				case codersdk.WorkspaceAgentLifecycleStartError:
+					showMessage()
+					return AgentStartError
+				default:
+					select {
+					case <-warningShown:
+						showMessage()
+					default:
+						// This state is normal, we don't want
+						// to show a message prematurely.
+					}
+				}
+				continue
+			}
 			return nil
 		case codersdk.WorkspaceAgentTimeout, codersdk.WorkspaceAgentDisconnected:
 			showMessage()
@@ -128,19 +164,78 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
 	}
 }
 
-func waitingMessage(agent codersdk.WorkspaceAgent) string {
+type message struct {
-	var m string
+	Spin         string
+	Prompt       string
+	Troubleshoot bool
+}
+
+func waitingMessage(agent codersdk.WorkspaceAgent, opts AgentOptions) (m *message) {
+	m = &message{
+		Spin:   fmt.Sprintf("Waiting for connection from %s...", Styles.Field.Render(agent.Name)),
+		Prompt: "Don't panic, your workspace is booting up!",
+	}
+	defer func() {
+		if opts.NoWait {
+			m.Spin = ""
+		}
+		if m.Spin != "" {
+			m.Spin = " " + m.Spin
+		}
+
+		// We don't want to wrap the troubleshooting URL, so we'll handle word
+		// wrapping ourselves (vs using lipgloss).
+		w := wordwrap.NewWriter(Styles.Paragraph.GetWidth() - Styles.Paragraph.GetMarginLeft()*2)
+		w.Breakpoints = []rune{' ', '\n'}
+
+		_, _ = fmt.Fprint(w, m.Prompt)
+		if m.Troubleshoot {
+			if agent.TroubleshootingURL != "" {
+				_, _ = fmt.Fprintf(w, " See troubleshooting instructions at:\n%s", agent.TroubleshootingURL)
+			} else {
+				_, _ = fmt.Fprint(w, " Wait for it to (re)connect or restart your workspace.")
+			}
+		}
+		_, _ = fmt.Fprint(w, "\n")
+
+		// We want to prefix the prompt with a caret, but we want text on the
+		// following lines to align with the text on the first line (i.e. added
+		// spacing).
+		ind := " " + Styles.Prompt.String()
+		iw := indent.NewWriter(1, func(w io.Writer) {
+			_, _ = w.Write([]byte(ind))
+			ind = "   " // Set indentation to space after initial prompt.
+		})
+		_, _ = fmt.Fprint(iw, w.String())
+		m.Prompt = iw.String()
+	}()
+
 	switch agent.Status {
 	case codersdk.WorkspaceAgentTimeout:
-		m = "The workspace agent is having trouble connecting."
+		m.Prompt = "The workspace agent is having trouble connecting."
 	case codersdk.WorkspaceAgentDisconnected:
-		m = "The workspace agent lost connection!"
+		m.Prompt = "The workspace agent lost connection!"
+	case codersdk.WorkspaceAgentConnected:
+		m.Spin = fmt.Sprintf("Waiting for %s to become ready...", Styles.Field.Render(agent.Name))
+		m.Prompt = "Don't panic, your workspace agent has connected and the workspace is getting ready!"
+		if opts.NoWait {
+			m.Prompt = "Your workspace is still getting ready, it may be in an incomplete state."
+		}
+
+		switch agent.LifecycleState {
+		case codersdk.WorkspaceAgentLifecycleStartTimeout:
+			m.Prompt = "The workspace is taking longer than expected to get ready, the agent startup script is still executing."
+		case codersdk.WorkspaceAgentLifecycleStartError:
+			m.Spin = ""
+			m.Prompt = "The workspace ran into a problem while getting ready, the agent startup script exited with non-zero status."
+		default:
+			// Not a failure state, no troubleshooting necessary.
+			return m
+		}
 	default:
 		// Not a failure state, no troubleshooting necessary.
-		return "Don't panic, your workspace is booting up!"
+		return m
 	}
-	if agent.TroubleshootingURL != "" {
+	m.Troubleshoot = true
-		return fmt.Sprintf("%s See troubleshooting instructions at: %s", m, agent.TroubleshootingURL)
+	return m
-	}
-	return fmt.Sprintf("%s Wait for it to (re)connect or restart your workspace.", m)
 }

cli/cliui/agent_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/atomic"
 
 	"github.com/coder/coder/cli/cliui"
@@ -17,13 +18,17 @@ import (
 
 func TestAgent(t *testing.T) {
 	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
 	var disconnected atomic.Bool
 	ptty := ptytest.New(t)
 	cmd := &cobra.Command{
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			err := cliui.Agent(cmd.Context(), cmd.OutOrStdout(), cliui.AgentOptions{
 				WorkspaceName: "example",
-				Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
+				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
 					agent := codersdk.WorkspaceAgent{
 						Status: codersdk.WorkspaceAgentDisconnected,
 					}
@@ -46,33 +51,34 @@ func TestAgent(t *testing.T) {
 		err := cmd.Execute()
 		assert.NoError(t, err)
 	}()
-	ptty.ExpectMatch("lost connection")
+	ptty.ExpectMatchContext(ctx, "lost connection")
 	disconnected.Store(true)
 	<-done
 }
 
-func TestAgentTimeoutWithTroubleshootingURL(t *testing.T) {
+func TestAgent_TimeoutWithTroubleshootingURL(t *testing.T) {
 	t.Parallel()
 
-	ctx, _ := testutil.Context(t)
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
 
 	wantURL := "https://coder.com/troubleshoot"
 
 	var connected, timeout atomic.Bool
 	cmd := &cobra.Command{
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: func(cmd *cobra.Command, _ []string) error {
 			err := cliui.Agent(cmd.Context(), cmd.OutOrStdout(), cliui.AgentOptions{
 				WorkspaceName: "example",
-				Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
+				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
 					agent := codersdk.WorkspaceAgent{
 						Status:             codersdk.WorkspaceAgentConnecting,
-						TroubleshootingURL: "https://coder.com/troubleshoot",
+						TroubleshootingURL: wantURL,
 					}
 					switch {
+					case !connected.Load() && timeout.Load():
+						agent.Status = codersdk.WorkspaceAgentTimeout
 					case connected.Load():
 						agent.Status = codersdk.WorkspaceAgentConnected
-					case timeout.Load():
-						agent.Status = codersdk.WorkspaceAgentTimeout
 					}
 					return agent, nil
 				},
@@ -85,15 +91,264 @@ func TestAgentTimeoutWithTroubleshootingURL(t *testing.T) {
 	ptty := ptytest.New(t)
 	cmd.SetOutput(ptty.Output())
 	cmd.SetIn(ptty.Input())
-	done := make(chan struct{})
+	done := make(chan error, 1)
 	go func() {
-		defer close(done)
+		done <- cmd.ExecuteContext(ctx)
-		err := cmd.ExecuteContext(ctx)
-		assert.NoError(t, err)
 	}()
-	ptty.ExpectMatch("Don't panic")
+	ptty.ExpectMatchContext(ctx, "Don't panic, your workspace is booting")
 	timeout.Store(true)
-	ptty.ExpectMatch(wantURL)
+	ptty.ExpectMatchContext(ctx, wantURL)
 	connected.Store(true)
-	<-done
+	require.NoError(t, <-done)
+}
+
+func TestAgent_StartupTimeout(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	wantURL := "https://coder.com/this-is-a-really-long-troubleshooting-url-that-should-not-wrap"
+
+	var status, state atomic.String
+	setStatus := func(s codersdk.WorkspaceAgentStatus) { status.Store(string(s)) }
+	setState := func(s codersdk.WorkspaceAgentLifecycle) { state.Store(string(s)) }
+	cmd := &cobra.Command{
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			err := cliui.Agent(cmd.Context(), cmd.OutOrStdout(), cliui.AgentOptions{
+				WorkspaceName: "example",
+				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
+					agent := codersdk.WorkspaceAgent{
+						Status:               codersdk.WorkspaceAgentConnecting,
+						DelayLoginUntilReady: true,
+						LifecycleState:       codersdk.WorkspaceAgentLifecycleCreated,
+						TroubleshootingURL:   wantURL,
+					}
+
+					if s := status.Load(); s != "" {
+						agent.Status = codersdk.WorkspaceAgentStatus(s)
+					}
+					if s := state.Load(); s != "" {
+						agent.LifecycleState = codersdk.WorkspaceAgentLifecycle(s)
+					}
+					return agent, nil
+				},
+				FetchInterval: time.Millisecond,
+				WarnInterval:  time.Millisecond,
+				NoWait:        false,
+			})
+			return err
+		},
+	}
+
+	ptty := ptytest.New(t)
+	cmd.SetOutput(ptty.Output())
+	cmd.SetIn(ptty.Input())
+	done := make(chan error, 1)
+	go func() {
+		done <- cmd.ExecuteContext(ctx)
+	}()
+	setStatus(codersdk.WorkspaceAgentConnecting)
+	ptty.ExpectMatchContext(ctx, "Don't panic, your workspace is booting")
+	setStatus(codersdk.WorkspaceAgentConnected)
+	setState(codersdk.WorkspaceAgentLifecycleStarting)
+	ptty.ExpectMatchContext(ctx, "workspace is getting ready")
+	setState(codersdk.WorkspaceAgentLifecycleStartTimeout)
+	ptty.ExpectMatchContext(ctx, "is taking longer")
+	ptty.ExpectMatchContext(ctx, wantURL)
+	setState(codersdk.WorkspaceAgentLifecycleReady)
+	require.NoError(t, <-done)
+}
+
+func TestAgent_StartErrorExit(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	wantURL := "https://coder.com/this-is-a-really-long-troubleshooting-url-that-should-not-wrap"
+
+	var status, state atomic.String
+	setStatus := func(s codersdk.WorkspaceAgentStatus) { status.Store(string(s)) }
+	setState := func(s codersdk.WorkspaceAgentLifecycle) { state.Store(string(s)) }
+	cmd := &cobra.Command{
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			err := cliui.Agent(cmd.Context(), cmd.OutOrStdout(), cliui.AgentOptions{
+				WorkspaceName: "example",
+				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
+					agent := codersdk.WorkspaceAgent{
+						Status:               codersdk.WorkspaceAgentConnecting,
+						DelayLoginUntilReady: true,
+						LifecycleState:       codersdk.WorkspaceAgentLifecycleCreated,
+						TroubleshootingURL:   wantURL,
+					}
+
+					if s := status.Load(); s != "" {
+						agent.Status = codersdk.WorkspaceAgentStatus(s)
+					}
+					if s := state.Load(); s != "" {
+						agent.LifecycleState = codersdk.WorkspaceAgentLifecycle(s)
+					}
+					return agent, nil
+				},
+				FetchInterval: time.Millisecond,
+				WarnInterval:  60 * time.Second,
+				NoWait:        false,
+			})
+			return err
+		},
+	}
+
+	ptty := ptytest.New(t)
+	cmd.SetOutput(ptty.Output())
+	cmd.SetIn(ptty.Input())
+	done := make(chan error, 1)
+	go func() {
+		done <- cmd.ExecuteContext(ctx)
+	}()
+	setStatus(codersdk.WorkspaceAgentConnected)
+	setState(codersdk.WorkspaceAgentLifecycleStarting)
+	ptty.ExpectMatchContext(ctx, "to become ready...")
+	setState(codersdk.WorkspaceAgentLifecycleStartError)
+	ptty.ExpectMatchContext(ctx, "ran into a problem")
+	err := <-done
+	require.ErrorIs(t, err, cliui.AgentStartError, "lifecycle start_error should exit with error")
+}
+
+func TestAgent_NoWait(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	wantURL := "https://coder.com/this-is-a-really-long-troubleshooting-url-that-should-not-wrap"
+
+	var status, state atomic.String
+	setStatus := func(s codersdk.WorkspaceAgentStatus) { status.Store(string(s)) }
+	setState := func(s codersdk.WorkspaceAgentLifecycle) { state.Store(string(s)) }
+	cmd := &cobra.Command{
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			err := cliui.Agent(cmd.Context(), cmd.OutOrStdout(), cliui.AgentOptions{
+				WorkspaceName: "example",
+				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
+					agent := codersdk.WorkspaceAgent{
+						Status:               codersdk.WorkspaceAgentConnecting,
+						DelayLoginUntilReady: true,
+						LifecycleState:       codersdk.WorkspaceAgentLifecycleCreated,
+						TroubleshootingURL:   wantURL,
+					}
+
+					if s := status.Load(); s != "" {
+						agent.Status = codersdk.WorkspaceAgentStatus(s)
+					}
+					if s := state.Load(); s != "" {
+						agent.LifecycleState = codersdk.WorkspaceAgentLifecycle(s)
+					}
+					return agent, nil
+				},
+				FetchInterval: time.Millisecond,
+				WarnInterval:  time.Second,
+				NoWait:        true,
+			})
+			return err
+		},
+	}
+
+	ptty := ptytest.New(t)
+	cmd.SetOutput(ptty.Output())
+	cmd.SetIn(ptty.Input())
+	done := make(chan error, 1)
+	go func() {
+		done <- cmd.ExecuteContext(ctx)
+	}()
+	setStatus(codersdk.WorkspaceAgentConnecting)
+	ptty.ExpectMatchContext(ctx, "Don't panic, your workspace is booting")
+
+	setStatus(codersdk.WorkspaceAgentConnected)
+	require.NoError(t, <-done, "created - should exit early")
+
+	setState(codersdk.WorkspaceAgentLifecycleStarting)
+	go func() { done <- cmd.ExecuteContext(ctx) }()
+	require.NoError(t, <-done, "starting - should exit early")
+
+	setState(codersdk.WorkspaceAgentLifecycleStartTimeout)
+	go func() { done <- cmd.ExecuteContext(ctx) }()
+	require.NoError(t, <-done, "start timeout - should exit early")
+
+	setState(codersdk.WorkspaceAgentLifecycleStartError)
+	go func() { done <- cmd.ExecuteContext(ctx) }()
+	require.NoError(t, <-done, "start error - should exit early")
+
+	setState(codersdk.WorkspaceAgentLifecycleReady)
+	go func() { done <- cmd.ExecuteContext(ctx) }()
+	require.NoError(t, <-done, "ready - should exit early")
+}
+
+func TestAgent_DelayLoginUntilReadyDisabled(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	wantURL := "https://coder.com/this-is-a-really-long-troubleshooting-url-that-should-not-wrap"
+
+	var status, state atomic.String
+	setStatus := func(s codersdk.WorkspaceAgentStatus) { status.Store(string(s)) }
+	setState := func(s codersdk.WorkspaceAgentLifecycle) { state.Store(string(s)) }
+	cmd := &cobra.Command{
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			err := cliui.Agent(cmd.Context(), cmd.OutOrStdout(), cliui.AgentOptions{
+				WorkspaceName: "example",
+				Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
+					agent := codersdk.WorkspaceAgent{
+						Status:               codersdk.WorkspaceAgentConnecting,
+						DelayLoginUntilReady: false,
+						LifecycleState:       codersdk.WorkspaceAgentLifecycleCreated,
+						TroubleshootingURL:   wantURL,
+					}
+
+					if s := status.Load(); s != "" {
+						agent.Status = codersdk.WorkspaceAgentStatus(s)
+					}
+					if s := state.Load(); s != "" {
+						agent.LifecycleState = codersdk.WorkspaceAgentLifecycle(s)
+					}
+					return agent, nil
+				},
+				FetchInterval: time.Millisecond,
+				WarnInterval:  time.Second,
+				NoWait:        false,
+			})
+			return err
+		},
+	}
+
+	ptty := ptytest.New(t)
+	cmd.SetOutput(ptty.Output())
+	cmd.SetIn(ptty.Input())
+	done := make(chan error, 1)
+	go func() {
+		done <- cmd.ExecuteContext(ctx)
+	}()
+	setStatus(codersdk.WorkspaceAgentConnecting)
+	ptty.ExpectMatchContext(ctx, "Don't panic, your workspace is booting")
+
+	setStatus(codersdk.WorkspaceAgentConnected)
+	require.NoError(t, <-done, "created - should exit early")
+
+	setState(codersdk.WorkspaceAgentLifecycleStarting)
+	go func() { done <- cmd.ExecuteContext(ctx) }()
+	require.NoError(t, <-done, "starting - should exit early")
+
+	setState(codersdk.WorkspaceAgentLifecycleStartTimeout)
+	go func() { done <- cmd.ExecuteContext(ctx) }()
+	require.NoError(t, <-done, "start timeout - should exit early")
+
+	setState(codersdk.WorkspaceAgentLifecycleStartError)
+	go func() { done <- cmd.ExecuteContext(ctx) }()
+	require.NoError(t, <-done, "start error - should exit early")
+
+	setState(codersdk.WorkspaceAgentLifecycleReady)
+	go func() { done <- cmd.ExecuteContext(ctx) }()
+	require.NoError(t, <-done, "ready - should exit early")
 }

cli/ssh.go

@@ -46,6 +46,7 @@ func ssh() *cobra.Command {
 		forwardGPG     bool
 		identityAgent  string
 		wsPollInterval time.Duration
+		noWait         bool
 	)
 	cmd := &cobra.Command{
 		Annotations: workspaceCommand,
@@ -90,8 +91,15 @@ func ssh() *cobra.Command {
 				Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
 					return client.WorkspaceAgent(ctx, workspaceAgent.ID)
 				},
+				NoWait: noWait,
 			})
 			if err != nil {
+				if xerrors.Is(err, context.Canceled) {
+					return cliui.Canceled
+				}
+				if xerrors.Is(err, cliui.AgentStartError) {
+					return xerrors.New("Agent startup script exited with non-zero status, use --no-wait to login anyway.")
+				}
 				return xerrors.Errorf("await agent: %w", err)
 			}
 
@@ -242,6 +250,7 @@ func ssh() *cobra.Command {
 	cliflag.BoolVarP(cmd.Flags(), &forwardGPG, "forward-gpg", "G", "CODER_SSH_FORWARD_GPG", false, "Specifies whether to forward the GPG agent. Unsupported on Windows workspaces, but supports all clients. Requires gnupg (gpg, gpgconf) on both the client and workspace. The GPG agent must already be running locally and will not be started for you. If a GPG agent is already running in the workspace, it will be attempted to be killed.")
 	cliflag.StringVarP(cmd.Flags(), &identityAgent, "identity-agent", "", "CODER_SSH_IDENTITY_AGENT", "", "Specifies which identity agent to use (overrides $SSH_AUTH_SOCK), forward agent must also be enabled")
 	cliflag.DurationVarP(cmd.Flags(), &wsPollInterval, "workspace-poll-interval", "", "CODER_WORKSPACE_POLL_INTERVAL", workspacePollInterval, "Specifies how often to poll for workspace automated shutdown.")
+	cliflag.BoolVarP(cmd.Flags(), &noWait, "no-wait", "", "CODER_SSH_NO_WAIT", false, "Specifies whether to wait for a workspace to become ready before logging in (only applicable when the delay login until ready option is enabled). 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.")
 	return cmd
 }
 

cli/ssh_test.go

@@ -28,6 +28,7 @@ import (
 
 	"github.com/coder/coder/agent"
 	"github.com/coder/coder/cli/clitest"
+	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/coderd/coderdtest"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/provisioner/echo"
@@ -136,7 +137,7 @@ func TestSSH(t *testing.T) {
 
 		cmdDone := tGo(t, func() {
 			err := cmd.ExecuteContext(ctx)
-			assert.ErrorIs(t, err, context.Canceled)
+			assert.ErrorIs(t, err, cliui.Canceled)
 		})
 		pty.ExpectMatch(wantURL)
 		cancel()

cli/testdata/coder_ssh_--help.golden

@@ -20,6 +20,13 @@ Flags:
                                            $SSH_AUTH_SOCK), forward agent must also be
                                            enabled.
                                            Consumes $CODER_SSH_IDENTITY_AGENT
+      --no-wait                            Specifies whether to wait for a workspace to become
+                                           ready before logging in (only applicable when the
+                                           delay login until ready option is enabled). 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.
+                                           Consumes $CODER_SSH_NO_WAIT
       --stdio                              Specifies whether to emit SSH output over
                                            stdin/stdout.
                                            Consumes $CODER_SSH_STDIO

go.mod

@@ -274,7 +274,7 @@ require (
 	github.com/mitchellh/go-ps v1.0.0 // indirect
 	github.com/moby/term v0.0.0-20210619224110-3f7ff695adc6 // indirect
 	github.com/muesli/ansi v0.0.0-20211031195517-c9f0611b6c70 // indirect
-	github.com/muesli/reflow v0.3.0 // indirect
+	github.com/muesli/reflow v0.3.0
 	github.com/muesli/termenv v0.11.1-0.20220212125758-44cd13922739 // indirect
 	github.com/niklasfasching/go-org v1.6.5 // indirect
 	github.com/nu7hatch/gouuid v0.0.0-20131221200532-179d4d0c4d8d // indirect