From 2f223fe45482a27911b81b8f75601d860e316be8 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 7 Jun 2023 12:47:23 +0000
Subject: [PATCH 1/5] feat(cli): Add wait and no-wait support to ssh

Fixes #7768
Refs #7893
---
 cli/cliui/agent.go                   | 10 +++++-----
 cli/cliui/agent_test.go              |  8 ++++----
 cli/portforward.go                   |  1 +
 cli/speedtest.go                     |  1 +
 cli/ssh.go                           | 28 ++++++++++++++++++++++++++--
 cli/testdata/coder_ssh_--help.golden | 13 ++++++++-----
 docs/cli/ssh.md                      | 11 ++++++++++-
 7 files changed, 55 insertions(+), 17 deletions(-)

diff --git a/cli/cliui/agent.go b/cli/cliui/agent.go
index f1f14dfef1f51..43b2375123834 100644
--- 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
index efed058e4fe1f..6d54544a45aee 100644
--- 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
index fd0f2bcfab57f..9d3d63030d13b 100644
--- 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
index 986088e2ea238..c9d63bbf36010 100644
--- 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
index 3a22ab1172c9f..8d0075c6dce7b 100644
--- a/cli/ssh.go
+++ b/cli/ssh.go
@@ -42,6 +42,7 @@ var (
 	autostopNotifyCountdown = []time.Duration{30 * time.Minute}
 )
 
+//nolint:gocyclo
 func (r *RootCmd) ssh() *clibase.Cmd {
 	var (
 		stdio          bool
@@ -49,6 +50,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 		forwardGPG     bool
 		identityAgent  string
 		wsPollInterval time.Duration
+		wait           bool
 		noWait         bool
 		logDir         string
 		logToFile      bool
@@ -66,6 +68,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 +111,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 +152,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 +365,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),
 		},
 		{
diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden
index 036a1a357e3af..63f5d6cf06f75 100644
--- a/cli/testdata/coder_ssh_--help.golden
+++ b/cli/testdata/coder_ssh_--help.golden
@@ -25,15 +25,18 @@ Start a shell into a workspace
           Enable diagnostic logging to file.
 
       --no-wait bool, $CODER_SSH_NO_WAIT
-          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.
+          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.
 
       --stdio bool, $CODER_SSH_STDIO
           Specifies whether to emit SSH output over stdin/stdout.
 
+      --wait bool, $CODER_SSH_WAIT
+          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.
+
       --workspace-poll-interval duration, $CODER_WORKSPACE_POLL_INTERVAL (default: 1m)
           Specifies how often to poll for workspace automated shutdown.
 
diff --git a/docs/cli/ssh.md b/docs/cli/ssh.md
index 7d580b23b08b1..901b7dd91f9b4 100644
--- a/docs/cli/ssh.md
+++ b/docs/cli/ssh.md
@@ -65,7 +65,7 @@ Enable diagnostic logging to file.
 | Type        | <code>bool</code>               |
 | Environment | <code>$CODER_SSH_NO_WAIT</code> |
 
-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.
+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.
 
 ### --stdio
 
@@ -76,6 +76,15 @@ Specifies whether to wait for a workspace to become ready before logging in (onl
 
 Specifies whether to emit SSH output over stdin/stdout.
 
+### --wait
+
+|             |                              |
+| ----------- | ---------------------------- |
+| Type        | <code>bool</code>            |
+| Environment | <code>$CODER_SSH_WAIT</code> |
+
+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.
+
 ### --workspace-poll-interval
 
 |             |                                             |

From 5e8c487c482d1f5d3b3db3d62fb1a97959950494 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 7 Jun 2023 15:15:39 +0000
Subject: [PATCH 2/5] Use proposed --wait=yes|no|auto

---
 cli/ssh.go                           | 37 ++++++++++++++++++----------
 cli/testdata/coder_ssh_--help.golden | 10 ++++----
 2 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/cli/ssh.go b/cli/ssh.go
index 8d0075c6dce7b..5923db624d75a 100644
--- a/cli/ssh.go
+++ b/cli/ssh.go
@@ -50,7 +50,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 		forwardGPG     bool
 		identityAgent  string
 		wsPollInterval time.Duration
-		wait           bool
+		waitEnum       string
 		noWait         bool
 		logDir         string
 		logToFile      bool
@@ -68,10 +68,6 @@ 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 {
@@ -112,14 +108,26 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			}
 
 			// Select the startup script behavior based on template configuration or flags.
-			if !wait && !noWait {
+			var wait bool
+			switch waitEnum {
+			case "yes":
+				wait = true
+			case "no":
+				wait = false
+			case "auto":
 				switch workspaceAgent.StartupScriptBehavior {
 				case codersdk.WorkspaceAgentStartupScriptBehaviorBlocking:
 					wait = true
 				case codersdk.WorkspaceAgentStartupScriptBehaviorNonBlocking:
 					wait = false
+				default:
+					return xerrors.Errorf("unknown startup script behavior %q", workspaceAgent.StartupScriptBehavior)
 				}
-			} else if noWait {
+			default:
+				return xerrors.Errorf("unknown wait value %q", waitEnum)
+			}
+			// The `--no-wait` flag is deprecated, but for now, check it.
+			if noWait {
 				wait = false
 			}
 
@@ -331,6 +339,13 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			return nil
 		},
 	}
+	waitOption := clibase.Option{
+		Flag:        "wait",
+		Env:         "CODER_SSH_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"),
+	}
 	cmd.Options = clibase.OptionSet{
 		{
 			Flag:        "stdio",
@@ -365,17 +380,13 @@ 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),
-		},
+		waitOption,
 		{
 			Flag:        "no-wait",
 			Env:         "CODER_SSH_NO_WAIT",
 			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),
+			UseInstead:  []clibase.Option{waitOption},
 		},
 		{
 			Flag:        "log-dir",
diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden
index 63f5d6cf06f75..f9442f6497980 100644
--- a/cli/testdata/coder_ssh_--help.golden
+++ b/cli/testdata/coder_ssh_--help.golden
@@ -28,14 +28,14 @@ Start a shell into a workspace
           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.
-
+ DEPRECATED 
       --stdio bool, $CODER_SSH_STDIO
           Specifies whether to emit SSH output over stdin/stdout.
 
-      --wait bool, $CODER_SSH_WAIT
-          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.
+      --wait yes|no|auto, $CODER_SSH_WAIT (default: auto)
+          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.
 
       --workspace-poll-interval duration, $CODER_WORKSPACE_POLL_INTERVAL (default: 1m)
           Specifies how often to poll for workspace automated shutdown.

From e4d37e12f36ec0fa6ba2260cd0abaf03069cc4c7 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 7 Jun 2023 16:05:04 +0000
Subject: [PATCH 3/5] Fix bug in dbfake

---
 coderd/database/dbfake/databasefake.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/coderd/database/dbfake/databasefake.go b/coderd/database/dbfake/databasefake.go
index 9294335106304..f1a6ffe05bb68 100644
--- a/coderd/database/dbfake/databasefake.go
+++ b/coderd/database/dbfake/databasefake.go
@@ -3060,6 +3060,7 @@ func (q *fakeQuerier) InsertWorkspaceAgent(_ context.Context, arg database.Inser
 		Architecture:             arg.Architecture,
 		OperatingSystem:          arg.OperatingSystem,
 		Directory:                arg.Directory,
+		StartupScriptBehavior:    arg.StartupScriptBehavior,
 		StartupScript:            arg.StartupScript,
 		InstanceMetadata:         arg.InstanceMetadata,
 		ResourceMetadata:         arg.ResourceMetadata,

From 6d0867eead420ab7a5774b523f026a519bf7b584 Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 7 Jun 2023 16:09:26 +0000
Subject: [PATCH 4/5] Clean up description

---
 cli/ssh.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cli/ssh.go b/cli/ssh.go
index 5923db624d75a..05573b3e2afae 100644
--- a/cli/ssh.go
+++ b/cli/ssh.go
@@ -384,7 +384,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 		{
 			Flag:        "no-wait",
 			Env:         "CODER_SSH_NO_WAIT",
-			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.",
+			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.",
 			Value:       clibase.BoolOf(&noWait),
 			UseInstead:  []clibase.Option{waitOption},
 		},

From 24a98ce560d1e3007965cfb9c70d58a70d13289e Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Wed, 7 Jun 2023 16:46:35 +0000
Subject: [PATCH 5/5] Gen

---
 cli/testdata/coder_ssh_--help.golden | 2 +-
 docs/cli/ssh.md                      | 9 +++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/cli/testdata/coder_ssh_--help.golden b/cli/testdata/coder_ssh_--help.golden
index f9442f6497980..4880365501690 100644
--- a/cli/testdata/coder_ssh_--help.golden
+++ b/cli/testdata/coder_ssh_--help.golden
@@ -27,7 +27,7 @@ Start a shell into a workspace
       --no-wait bool, $CODER_SSH_NO_WAIT
           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.
+          behavior as non-blocking.
  DEPRECATED 
       --stdio bool, $CODER_SSH_STDIO
           Specifies whether to emit SSH output over stdin/stdout.
diff --git a/docs/cli/ssh.md b/docs/cli/ssh.md
index 901b7dd91f9b4..e6f03f6c326ae 100644
--- a/docs/cli/ssh.md
+++ b/docs/cli/ssh.md
@@ -65,7 +65,7 @@ Enable diagnostic logging to file.
 | Type        | <code>bool</code>               |
 | Environment | <code>$CODER_SSH_NO_WAIT</code> |
 
-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.
+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.
 
 ### --stdio
 
@@ -79,11 +79,12 @@ Specifies whether to emit SSH output over stdin/stdout.
 ### --wait
 
 |             |                              |
-| ----------- | ---------------------------- |
-| Type        | <code>bool</code>            |
+| ----------- | ---------------------------- | --- | ------------ |
+| Type        | <code>enum[yes               | no  | auto]</code> |
 | Environment | <code>$CODER_SSH_WAIT</code> |
+| Default     | <code>auto</code>            |
 
-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.
+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.
 
 ### --workspace-poll-interval