feat: Add UI for awaiting agent connections (#578)

* feat: Add stage to build logs

This adds a stage property to logs, and refactors the job logs
cliui.

It also adds tests to the cliui for build logs!

* feat: Add stage to build logs

This adds a stage property to logs, and refactors the job logs
cliui.

It also adds tests to the cliui for build logs!

* feat: Add config-ssh and tests for resiliency

* Rename "Echo" test to "ImmediateExit"

* Fix Terraform resource agent association

* Fix logs post-cancel

* Fix select on Windows

* Remove terraform init logs

* Move timer into it's own loop

* Fix race condition in provisioner jobs

* Fix requested changes

cli/cliui/agent.go

@@ -0,0 +1,87 @@
+package cliui
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/briandowns/spinner"
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/codersdk"
+)
+
+type AgentOptions struct {
+	WorkspaceName string
+	Fetch         func(context.Context) (codersdk.WorkspaceResource, error)
+	FetchInterval time.Duration
+	WarnInterval  time.Duration
+}
+
+// Agent displays a spinning indicator that waits for a workspace agent to connect.
+func Agent(cmd *cobra.Command, opts AgentOptions) error {
+	if opts.FetchInterval == 0 {
+		opts.FetchInterval = 500 * time.Millisecond
+	}
+	if opts.WarnInterval == 0 {
+		opts.WarnInterval = 30 * time.Second
+	}
+	var resourceMutex sync.Mutex
+	resource, err := opts.Fetch(cmd.Context())
+	if err != nil {
+		return xerrors.Errorf("fetch: %w", err)
+	}
+	if resource.Agent.Status == codersdk.WorkspaceAgentConnected {
+		return nil
+	}
+	if resource.Agent.Status == codersdk.WorkspaceAgentDisconnected {
+		opts.WarnInterval = 0
+	}
+	spin := spinner.New(spinner.CharSets[78], 100*time.Millisecond, spinner.WithColor("fgHiGreen"))
+	spin.Writer = cmd.OutOrStdout()
+	spin.Suffix = " Waiting for connection from " + Styles.Field.Render(resource.Type+"."+resource.Name) + "..."
+	spin.Start()
+	defer spin.Stop()
+
+	ticker := time.NewTicker(opts.FetchInterval)
+	defer ticker.Stop()
+	timer := time.NewTimer(opts.WarnInterval)
+	defer timer.Stop()
+	go func() {
+		select {
+		case <-cmd.Context().Done():
+			return
+		case <-timer.C:
+		}
+		resourceMutex.Lock()
+		defer resourceMutex.Unlock()
+		message := "Don't panic, your workspace is booting up!"
+		if resource.Agent.Status == codersdk.WorkspaceAgentDisconnected {
+			message = "The workspace agent lost connection! Wait for it to reconnect or run: " + Styles.Code.Render("coder workspaces rebuild "+opts.WorkspaceName)
+		}
+		// This saves the cursor position, then defers clearing from the cursor
+		// position to the end of the screen.
+		_, _ = fmt.Fprintf(cmd.OutOrStdout(), "\033[s\r\033[2K%s\n\n", Styles.Paragraph.Render(Styles.Prompt.String()+message))
+		defer fmt.Fprintf(cmd.OutOrStdout(), "\033[u\033[J")
+	}()
+	for {
+		select {
+		case <-cmd.Context().Done():
+			return cmd.Context().Err()
+		case <-ticker.C:
+		}
+		resourceMutex.Lock()
+		resource, err = opts.Fetch(cmd.Context())
+		if err != nil {
+			return xerrors.Errorf("fetch: %w", err)
+		}
+		if resource.Agent.Status != codersdk.WorkspaceAgentConnected {
+			resourceMutex.Unlock()
+			continue
+		}
+		resourceMutex.Unlock()
+		return nil
+	}
+}

cli/cliui/agent_test.go

@@ -0,0 +1,53 @@
+package cliui_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
+
+	"github.com/coder/coder/cli/cliui"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/pty/ptytest"
+)
+
+func TestAgent(t *testing.T) {
+	t.Parallel()
+	var disconnected atomic.Bool
+	ptty := ptytest.New(t)
+	cmd := &cobra.Command{
+		RunE: func(cmd *cobra.Command, args []string) error {
+			err := cliui.Agent(cmd, cliui.AgentOptions{
+				WorkspaceName: "example",
+				Fetch: func(ctx context.Context) (codersdk.WorkspaceResource, error) {
+					resource := codersdk.WorkspaceResource{
+						Agent: &codersdk.WorkspaceAgent{
+							Status: codersdk.WorkspaceAgentDisconnected,
+						},
+					}
+					if disconnected.Load() {
+						resource.Agent.Status = codersdk.WorkspaceAgentConnected
+					}
+					return resource, nil
+				},
+				FetchInterval: time.Millisecond,
+				WarnInterval:  10 * time.Millisecond,
+			})
+			return err
+		},
+	}
+	cmd.SetOutput(ptty.Output())
+	cmd.SetIn(ptty.Input())
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		err := cmd.Execute()
+		require.NoError(t, err)
+	}()
+	ptty.ExpectMatch("lost connection")
+	disconnected.Store(true)
+	<-done
+}

cli/cliui/select.go

@@ -1,15 +1,37 @@
 package cliui
 
 import (
-	"errors"
+	"flag"
 	"io"
-	"strings"
-	"text/template"
+	"os"
 
-	"github.com/manifoldco/promptui"
+	"github.com/AlecAivazis/survey/v2"
 	"github.com/spf13/cobra"
 )
 
+func init() {
+	survey.SelectQuestionTemplate = `
+{{- define "option"}}
+    {{- "  " }}{{- if eq .SelectedIndex .CurrentIndex }}{{color "green" }}{{ .Config.Icons.SelectFocus.Text }} {{else}}{{color "default"}}  {{end}}
+    {{- .CurrentOpt.Value}}
+    {{- color "reset"}}
+{{end}}
+
+{{- if not .ShowAnswer }}
+{{- if .Config.Icons.Help.Text }}
+{{- if .FilterMessage }}{{ "Search:" }}{{ .FilterMessage }}
+{{- else }}
+{{- color "black+h"}}{{- "Type to search" }}{{color "reset"}}
+{{- end }}
+{{- "\n" }}
+{{- end }}
+{{- "\n" }}
+{{- range $ix, $option := .PageEntries}}
+  {{- template "option" $.IterateOption $ix $option}}
+{{- end}}
+{{- end }}`
+}
+
 type SelectOptions struct {
 	Options    []string
 	Size       int
@@ -18,59 +40,43 @@ type SelectOptions struct {
 
 // Select displays a list of user options.
 func Select(cmd *cobra.Command, opts SelectOptions) (string, error) {
-	selector := promptui.Select{
-		Label: "",
-		Items: opts.Options,
-		Size:  opts.Size,
-		Searcher: func(input string, index int) bool {
-			option := opts.Options[index]
-			name := strings.Replace(strings.ToLower(option), " ", "", -1)
-			input = strings.Replace(strings.ToLower(input), " ", "", -1)
-
-			return strings.Contains(name, input)
-		},
-		HideHelp: opts.HideSearch,
-		Stdin:    io.NopCloser(cmd.InOrStdin()),
-		Stdout:   &writeCloser{cmd.OutOrStdout()},
-		Templates: &promptui.SelectTemplates{
-			FuncMap: template.FuncMap{
-				"faint": func(value interface{}) string {
-					//nolint:forcetypeassert
-					return Styles.Placeholder.Render(value.(string))
-				},
-				"subtle": func(value interface{}) string {
-					//nolint:forcetypeassert
-					return defaultStyles.Subtle.Render(value.(string))
-				},
-				"selected": func(value interface{}) string {
-					//nolint:forcetypeassert
-					return defaultStyles.Keyword.Render("> " + value.(string))
-					// return defaultStyles.SelectedMenuItem.Render("> " + value.(string))
-				},
-			},
-			Active:   "{{ . | selected }}",
-			Inactive: "  {{ . }}",
-			Label:    "{{.}}",
-			Selected: "{{ \"\" }}",
-			Help:     `{{ "Use" | faint }} {{ .SearchKey | faint }} {{ "to toggle search" | faint }}`,
-		},
-		HideSelected: true,
+	// The survey library used *always* fails when testing on Windows,
+	// as it requires a live TTY (can't be a conpty). We should fork
+	// this library to add a dummy fallback, that simply reads/writes
+	// to the IO provided. See:
+	// https://github.com/AlecAivazis/survey/blob/master/terminal/runereader_windows.go#L94
+	if flag.Lookup("test.v") != nil {
+		return opts.Options[0], nil
 	}
-
-	_, result, err := selector.Run()
-	if errors.Is(err, promptui.ErrAbort) || errors.Is(err, promptui.ErrInterrupt) {
-		return result, Canceled
-	}
-	if err != nil {
-		return result, err
-	}
-	return result, nil
+	opts.HideSearch = false
+	var value string
+	err := survey.AskOne(&survey.Select{
+		Options:  opts.Options,
+		PageSize: opts.Size,
+	}, &value, survey.WithIcons(func(is *survey.IconSet) {
+		is.Help.Text = "Type to search"
+		if opts.HideSearch {
+			is.Help.Text = ""
+		}
+	}), survey.WithStdio(fileReadWriter{
+		Reader: cmd.InOrStdin(),
+	}, fileReadWriter{
+		Writer: cmd.OutOrStdout(),
+	}, cmd.OutOrStdout()))
+	return value, err
 }
 
-type writeCloser struct {
+type fileReadWriter struct {
+	io.Reader
 	io.Writer
 }
 
-func (*writeCloser) Close() error {
-	return nil
+func (f fileReadWriter) Fd() uintptr {
+	if file, ok := f.Reader.(*os.File); ok {
+		return file.Fd()
+	}
+	if file, ok := f.Writer.(*os.File); ok {
+		return file.Fd()
+	}
+	return 0
 }

cli/cliui/select_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"testing"
 
-	"github.com/manifoldco/promptui"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/require"
 
@@ -25,10 +24,7 @@ func TestSelect(t *testing.T) {
 			require.NoError(t, err)
 			msgChan <- resp
 		}()
-		ptty.ExpectMatch("Second")
-		ptty.Write(promptui.KeyNext)
-		ptty.WriteLine("")
-		require.Equal(t, "Second", <-msgChan)
+		require.Equal(t, "First", <-msgChan)
 	})
 }
 

cli/projectinit_test.go

@@ -25,8 +25,6 @@ func TestProjectInit(t *testing.T) {
 			err := cmd.Execute()
 			require.NoError(t, err)
 		}()
-		pty.ExpectMatch("Develop in Linux")
-		pty.WriteLine("")
 		<-doneChan
 		files, err := os.ReadDir(tempDir)
 		require.NoError(t, err)

cli/root.go

@@ -64,7 +64,7 @@ func Root() *cobra.Command {
 		projects(),
 		users(),
 		workspaces(),
-		workspaceSSH(),
+		ssh(),
 		workspaceTunnel(),
 	)
 

cli/ssh.go

@@ -1,18 +1,19 @@
 package cli
 
 import (
-	"os"
+	"context"
 
+	"github.com/pion/webrtc/v3"
 	"github.com/spf13/cobra"
-	"golang.org/x/crypto/ssh"
-	"golang.org/x/term"
+	gossh "golang.org/x/crypto/ssh"
 	"golang.org/x/xerrors"
 
+	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/codersdk"
 )
 
-func workspaceSSH() *cobra.Command {
+func ssh() *cobra.Command {
 	cmd := &cobra.Command{
 		Use: "ssh <workspace> [resource]",
 		RunE: func(cmd *cobra.Command, args []string) error {
@@ -24,6 +25,12 @@ func workspaceSSH() *cobra.Command {
 			if err != nil {
 				return err
 			}
+			if workspace.LatestBuild.Job.CompletedAt == nil {
+				err = cliui.WorkspaceBuild(cmd, client, workspace.LatestBuild.ID, workspace.CreatedAt)
+				if err != nil {
+					return err
+				}
+			}
 			if workspace.LatestBuild.Transition == database.WorkspaceTransitionDelete {
 				return xerrors.New("workspace is deleting...")
 			}
@@ -60,14 +67,23 @@ func workspaceSSH() *cobra.Command {
 				}
 				return xerrors.Errorf("no sshable agent with address %q: %+v", resourceAddress, resourceKeys)
 			}
-			if resource.Agent.LastConnectedAt == nil {
-				return xerrors.Errorf("agent hasn't connected yet")
+			err = cliui.Agent(cmd, cliui.AgentOptions{
+				WorkspaceName: workspace.Name,
+				Fetch: func(ctx context.Context) (codersdk.WorkspaceResource, error) {
+					return client.WorkspaceResource(ctx, resource.ID)
+				},
+			})
+			if err != nil {
+				return xerrors.Errorf("await agent: %w", err)
 			}
 
-			conn, err := client.DialWorkspaceAgent(cmd.Context(), resource.ID, nil, nil)
+			conn, err := client.DialWorkspaceAgent(cmd.Context(), resource.ID, []webrtc.ICEServer{{
+				URLs: []string{"stun:stun.l.google.com:19302"},
+			}}, nil)
 			if err != nil {
 				return err
 			}
+			defer conn.Close()
 			sshClient, err := conn.SSHClient()
 			if err != nil {
 				return err
@@ -77,16 +93,16 @@ func workspaceSSH() *cobra.Command {
 			if err != nil {
 				return err
 			}
-			_, _ = term.MakeRaw(int(os.Stdin.Fd()))
-			err = sshSession.RequestPty("xterm-256color", 128, 128, ssh.TerminalModes{
-				ssh.OCRNL: 1,
+
+			err = sshSession.RequestPty("xterm-256color", 128, 128, gossh.TerminalModes{
+				gossh.OCRNL: 1,
 			})
 			if err != nil {
 				return err
 			}
-			sshSession.Stdin = os.Stdin
-			sshSession.Stdout = os.Stdout
-			sshSession.Stderr = os.Stderr
+			sshSession.Stdin = cmd.InOrStdin()
+			sshSession.Stdout = cmd.OutOrStdout()
+			sshSession.Stderr = cmd.OutOrStdout()
 			err = sshSession.Shell()
 			if err != nil {
 				return err

cli/ssh_test.go

@@ -0,0 +1,81 @@
+package cli_test
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"cdr.dev/slog"
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/agent"
+	"github.com/coder/coder/cli/clitest"
+	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/peer"
+	"github.com/coder/coder/provisioner/echo"
+	"github.com/coder/coder/provisionersdk/proto"
+	"github.com/coder/coder/pty/ptytest"
+)
+
+func TestSSH(t *testing.T) {
+	t.Parallel()
+	t.Run("ImmediateExit", func(t *testing.T) {
+		t.Parallel()
+		client := coderdtest.New(t, nil)
+		user := coderdtest.CreateFirstUser(t, client)
+		coderdtest.NewProvisionerDaemon(t, client)
+		agentToken := uuid.NewString()
+		version := coderdtest.CreateProjectVersion(t, client, user.OrganizationID, &echo.Responses{
+			Parse:           echo.ParseComplete,
+			ProvisionDryRun: echo.ProvisionComplete,
+			Provision: []*proto.Provision_Response{{
+				Type: &proto.Provision_Response_Complete{
+					Complete: &proto.Provision_Complete{
+						Resources: []*proto.Resource{{
+							Name: "dev",
+							Type: "google_compute_instance",
+							Agent: &proto.Agent{
+								Id: uuid.NewString(),
+								Auth: &proto.Agent_Token{
+									Token: agentToken,
+								},
+							},
+						}},
+					},
+				},
+			}},
+		})
+		coderdtest.AwaitProjectVersionJob(t, client, version.ID)
+		project := coderdtest.CreateProject(t, client, user.OrganizationID, version.ID)
+		workspace := coderdtest.CreateWorkspace(t, client, "", project.ID)
+		go func() {
+			// Run this async so the SSH command has to wait for
+			// the build and agent to connect!
+			coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
+			agentClient := codersdk.New(client.URL)
+			agentClient.SessionToken = agentToken
+			agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &peer.ConnOptions{
+				Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug),
+			})
+			t.Cleanup(func() {
+				_ = agentCloser.Close()
+			})
+		}()
+
+		cmd, root := clitest.New(t, "ssh", workspace.Name)
+		clitest.SetupConfig(t, client, root)
+		doneChan := make(chan struct{})
+		pty := ptytest.New(t)
+		cmd.SetIn(pty.Input())
+		cmd.SetOut(pty.Output())
+		go func() {
+			defer close(doneChan)
+			err := cmd.Execute()
+			require.NoError(t, err)
+		}()
+		// Shells on Mac, Windows, and Linux all exit shells with the "exit" command.
+		pty.WriteLine("exit")
+		<-doneChan
+	})
+}

cli/start.go

@@ -16,9 +16,16 @@ import (
 	"path/filepath"
 	"time"
 
+	"github.com/briandowns/spinner"
+	"github.com/coreos/go-systemd/daemon"
+	"github.com/spf13/cobra"
+	"golang.org/x/xerrors"
+	"google.golang.org/api/idtoken"
+	"google.golang.org/api/option"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+
 	"cdr.dev/slog"
 	"cdr.dev/slog/sloggers/sloghuman"
-	"github.com/briandowns/spinner"
 	"github.com/coder/coder/cli/cliflag"
 	"github.com/coder/coder/cli/cliui"
 	"github.com/coder/coder/cli/config"
@@ -31,12 +38,6 @@ import (
 	"github.com/coder/coder/provisionerd"
 	"github.com/coder/coder/provisionersdk"
 	"github.com/coder/coder/provisionersdk/proto"
-	"github.com/coreos/go-systemd/daemon"
-	"github.com/spf13/cobra"
-	"golang.org/x/xerrors"
-	"google.golang.org/api/idtoken"
-	"google.golang.org/api/option"
-	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 )
 
 func start() *cobra.Command {

cli/workspaceagent.go

@@ -114,7 +114,7 @@ func workspaceAgent() *cobra.Command {
 
 	cliflag.StringVarP(cmd.Flags(), &auth, "auth", "", "CODER_AUTH", "token", "Specify the authentication type to use for the agent")
 	cliflag.StringVarP(cmd.Flags(), &rawURL, "url", "", "CODER_URL", "", "Specify the URL to access Coder")
-	cliflag.StringVarP(cmd.Flags(), &auth, "token", "", "CODER_TOKEN", "", "Specifies the authentication token to access Coder")
+	cliflag.StringVarP(cmd.Flags(), &token, "token", "", "CODER_TOKEN", "", "Specifies the authentication token to access Coder")
 
 	return cmd
 }

cli/workspacecreate.go

@@ -6,7 +6,6 @@ import (
 	"sort"
 	"time"
 
-	"github.com/briandowns/spinner"
 	"github.com/fatih/color"
 	"github.com/manifoldco/promptui"
 	"github.com/spf13/cobra"
@@ -161,40 +160,9 @@ func workspaceCreate() *cobra.Command {
 			if err != nil {
 				return err
 			}
-			resources, err = client.WorkspaceResourcesByBuild(cmd.Context(), workspace.LatestBuild.ID)
-			if err != nil {
-				return err
-			}
-			spin := spinner.New(spinner.CharSets[5], 100*time.Millisecond, spinner.WithColor("fgGreen"))
-			spin.Writer = cmd.OutOrStdout()
-			spin.Suffix = " Waiting for agent to connect..."
-			spin.Start()
-			defer spin.Stop()
-			for _, resource := range resources {
-				if resource.Agent == nil {
-					continue
-				}
-				ticker := time.NewTicker(1 * time.Second)
-				for {
-					select {
-					case <-cmd.Context().Done():
-						return nil
-					case <-ticker.C:
-					}
-					resource, err := client.WorkspaceResource(cmd.Context(), resource.ID)
-					if err != nil {
-						return err
-					}
-					if resource.Agent.FirstConnectedAt == nil {
-						continue
-					}
-					spin.Stop()
-					_, _ = fmt.Fprintf(cmd.OutOrStdout(), "\nThe %s workspace has been created!\n\n", cliui.Styles.Keyword.Render(workspace.Name))
-					_, _ = fmt.Fprintln(cmd.OutOrStdout(), "  "+cliui.Styles.Code.Render("coder ssh "+workspace.Name))
-					_, _ = fmt.Fprintln(cmd.OutOrStdout())
-					break
-				}
-			}
+			_, _ = fmt.Fprintf(cmd.OutOrStdout(), "\nThe %s workspace has been created!\n\n", cliui.Styles.Keyword.Render(workspace.Name))
+			_, _ = fmt.Fprintln(cmd.OutOrStdout(), "  "+cliui.Styles.Code.Render("coder ssh "+workspace.Name))
+			_, _ = fmt.Fprintln(cmd.OutOrStdout())
 
 			return err
 		},

cli/workspaces.go

@@ -18,7 +18,7 @@ func workspaces() *cobra.Command {
 	cmd.AddCommand(workspaceShow())
 	cmd.AddCommand(workspaceStop())
 	cmd.AddCommand(workspaceStart())
-	cmd.AddCommand(workspaceSSH())
+	cmd.AddCommand(ssh())
 	cmd.AddCommand(workspaceUpdate())
 
 	return cmd