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

.vscode/settings.json

@@ -8,8 +8,10 @@
     "drpcconn",
     "drpcmux",
     "drpcserver",
+    "Dsts",
     "fatih",
     "goarch",
+    "gographviz",
     "goleak",
     "gossh",
     "hashicorp",

agent/agent_test.go

@@ -39,7 +39,10 @@ func TestAgent(t *testing.T) {
 		t.Cleanup(func() {
 			_ = conn.Close()
 		})
-		client := agent.Conn{conn}
+		client := agent.Conn{
+			Negotiator: api,
+			Conn:       conn,
+		}
 		sshClient, err := client.SSHClient()
 		require.NoError(t, err)
 		session, err := sshClient.NewSession()
@@ -65,7 +68,10 @@ func TestAgent(t *testing.T) {
 		t.Cleanup(func() {
 			_ = conn.Close()
 		})
-		client := &agent.Conn{conn}
+		client := &agent.Conn{
+			Negotiator: api,
+			Conn:       conn,
+		}
 		sshClient, err := client.SSHClient()
 		require.NoError(t, err)
 		session, err := sshClient.NewSession()

agent/conn.go

@@ -8,11 +8,15 @@ import (
 	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/peer"
+	"github.com/coder/coder/peerbroker/proto"
 )
 
 // Conn wraps a peer connection with helper functions to
 // communicate with the agent.
 type Conn struct {
+	// Negotiator is responsible for exchanging messages.
+	Negotiator proto.DRPCPeerBrokerClient
+
 	*peer.Conn
 }
 
@@ -48,3 +52,8 @@ func (c *Conn) SSHClient() (*ssh.Client, error) {
 	}
 	return ssh.NewClient(sshConn, channels, requests), nil
 }
+
+func (c *Conn) Close() error {
+	_ = c.Negotiator.DRPCConn().Close()
+	return c.Conn.Close()
+}

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(),
 	)