Skip to content

Commit 82dfd6c

Browse files
kylecarbskylecarbs
authored
feat: Add UI for awaiting agent connections (coder#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
1 parent 620c889 commit 82dfd6c

File tree

26 files changed

+536
-228
lines changed

26 files changed

+536
-228
lines changed

‎.vscode/settings.json

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,10 @@
88
"drpcconn",
99
"drpcmux",
1010
"drpcserver",
11+
"Dsts",
1112
"fatih",
1213
"goarch",
14+
"gographviz",
1315
"goleak",
1416
"gossh",
1517
"hashicorp",

‎agent/agent_test.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,10 @@ func TestAgent(t *testing.T) {
3939
t.Cleanup(func() {
4040
_ = conn.Close()
4141
})
42-
client := agent.Conn{conn}
42+
client := agent.Conn{
43+
Negotiator: api,
44+
Conn: conn,
45+
}
4346
sshClient, err := client.SSHClient()
4447
require.NoError(t, err)
4548
session, err := sshClient.NewSession()
@@ -65,7 +68,10 @@ func TestAgent(t *testing.T) {
6568
t.Cleanup(func() {
6669
_ = conn.Close()
6770
})
68-
client := &agent.Conn{conn}
71+
client := &agent.Conn{
72+
Negotiator: api,
73+
Conn: conn,
74+
}
6975
sshClient, err := client.SSHClient()
7076
require.NoError(t, err)
7177
session, err := sshClient.NewSession()

‎agent/conn.go

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,15 @@ import (
88
"golang.org/x/xerrors"
99

1010
"github.com/coder/coder/peer"
11+
"github.com/coder/coder/peerbroker/proto"
1112
)
1213

1314
// Conn wraps a peer connection with helper functions to
1415
// communicate with the agent.
1516
type Conn struct {
17+
// Negotiator is responsible for exchanging messages.
18+
Negotiator proto.DRPCPeerBrokerClient
19+
1620
*peer.Conn
1721
}
1822

@@ -48,3 +52,8 @@ func (c *Conn) SSHClient() (*ssh.Client, error) {
4852
}
4953
return ssh.NewClient(sshConn, channels, requests), nil
5054
}
55+
56+
func (c *Conn) Close() error {
57+
_ = c.Negotiator.DRPCConn().Close()
58+
return c.Conn.Close()
59+
}

‎cli/cliui/agent.go

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
package cliui
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"sync"
7+
"time"
8+
9+
"github.com/briandowns/spinner"
10+
"github.com/spf13/cobra"
11+
"golang.org/x/xerrors"
12+
13+
"github.com/coder/coder/codersdk"
14+
)
15+
16+
type AgentOptions struct {
17+
WorkspaceName string
18+
Fetch func(context.Context) (codersdk.WorkspaceResource, error)
19+
FetchInterval time.Duration
20+
WarnInterval time.Duration
21+
}
22+
23+
// Agent displays a spinning indicator that waits for a workspace agent to connect.
24+
func Agent(cmd *cobra.Command, opts AgentOptions) error {
25+
if opts.FetchInterval == 0 {
26+
opts.FetchInterval = 500 * time.Millisecond
27+
}
28+
if opts.WarnInterval == 0 {
29+
opts.WarnInterval = 30 * time.Second
30+
}
31+
var resourceMutex sync.Mutex
32+
resource, err := opts.Fetch(cmd.Context())
33+
if err != nil {
34+
return xerrors.Errorf("fetch: %w", err)
35+
}
36+
if resource.Agent.Status == codersdk.WorkspaceAgentConnected {
37+
return nil
38+
}
39+
if resource.Agent.Status == codersdk.WorkspaceAgentDisconnected {
40+
opts.WarnInterval = 0
41+
}
42+
spin := spinner.New(spinner.CharSets[78], 100*time.Millisecond, spinner.WithColor("fgHiGreen"))
43+
spin.Writer = cmd.OutOrStdout()
44+
spin.Suffix = " Waiting for connection from " + Styles.Field.Render(resource.Type+"."+resource.Name) + "..."
45+
spin.Start()
46+
defer spin.Stop()
47+
48+
ticker := time.NewTicker(opts.FetchInterval)
49+
defer ticker.Stop()
50+
timer := time.NewTimer(opts.WarnInterval)
51+
defer timer.Stop()
52+
go func() {
53+
select {
54+
case <-cmd.Context().Done():
55+
return
56+
case <-timer.C:
57+
}
58+
resourceMutex.Lock()
59+
defer resourceMutex.Unlock()
60+
message := "Don't panic, your workspace is booting up!"
61+
if resource.Agent.Status == codersdk.WorkspaceAgentDisconnected {
62+
message = "The workspace agent lost connection! Wait for it to reconnect or run: " + Styles.Code.Render("coder workspaces rebuild "+opts.WorkspaceName)
63+
}
64+
// This saves the cursor position, then defers clearing from the cursor
65+
// position to the end of the screen.
66+
_, _ = fmt.Fprintf(cmd.OutOrStdout(), "\033[s\r\033[2K%s\n\n", Styles.Paragraph.Render(Styles.Prompt.String()+message))
67+
defer fmt.Fprintf(cmd.OutOrStdout(), "\033[u\033[J")
68+
}()
69+
for {
70+
select {
71+
case <-cmd.Context().Done():
72+
return cmd.Context().Err()
73+
case <-ticker.C:
74+
}
75+
resourceMutex.Lock()
76+
resource, err = opts.Fetch(cmd.Context())
77+
if err != nil {
78+
return xerrors.Errorf("fetch: %w", err)
79+
}
80+
if resource.Agent.Status != codersdk.WorkspaceAgentConnected {
81+
resourceMutex.Unlock()
82+
continue
83+
}
84+
resourceMutex.Unlock()
85+
return nil
86+
}
87+
}

‎cli/cliui/agent_test.go

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package cliui_test
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/spf13/cobra"
9+
"github.com/stretchr/testify/require"
10+
"go.uber.org/atomic"
11+
12+
"github.com/coder/coder/cli/cliui"
13+
"github.com/coder/coder/codersdk"
14+
"github.com/coder/coder/pty/ptytest"
15+
)
16+
17+
func TestAgent(t *testing.T) {
18+
t.Parallel()
19+
var disconnected atomic.Bool
20+
ptty := ptytest.New(t)
21+
cmd := &cobra.Command{
22+
RunE: func(cmd *cobra.Command, args []string) error {
23+
err := cliui.Agent(cmd, cliui.AgentOptions{
24+
WorkspaceName: "example",
25+
Fetch: func(ctx context.Context) (codersdk.WorkspaceResource, error) {
26+
resource := codersdk.WorkspaceResource{
27+
Agent: &codersdk.WorkspaceAgent{
28+
Status: codersdk.WorkspaceAgentDisconnected,
29+
},
30+
}
31+
if disconnected.Load() {
32+
resource.Agent.Status = codersdk.WorkspaceAgentConnected
33+
}
34+
return resource, nil
35+
},
36+
FetchInterval: time.Millisecond,
37+
WarnInterval: 10 * time.Millisecond,
38+
})
39+
return err
40+
},
41+
}
42+
cmd.SetOutput(ptty.Output())
43+
cmd.SetIn(ptty.Input())
44+
done := make(chan struct{})
45+
go func() {
46+
defer close(done)
47+
err := cmd.Execute()
48+
require.NoError(t, err)
49+
}()
50+
ptty.ExpectMatch("lost connection")
51+
disconnected.Store(true)
52+
<-done
53+
}

‎cli/cliui/select.go

Lines changed: 59 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,37 @@
11
package cliui
22

33
import (
4-
"errors"
4+
"flag"
55
"io"
6-
"strings"
7-
"text/template"
6+
"os"
87

9-
"github.com/manifoldco/promptui"
8+
"github.com/AlecAivazis/survey/v2"
109
"github.com/spf13/cobra"
1110
)
1211

12+
func init() {
13+
survey.SelectQuestionTemplate = `
14+
{{- define "option"}}
15+
{{- " " }}{{- if eq .SelectedIndex .CurrentIndex }}{{color "green" }}{{ .Config.Icons.SelectFocus.Text }} {{else}}{{color "default"}} {{end}}
16+
{{- .CurrentOpt.Value}}
17+
{{- color "reset"}}
18+
{{end}}
19+
20+
{{- if not .ShowAnswer }}
21+
{{- if .Config.Icons.Help.Text }}
22+
{{- if .FilterMessage }}{{ "Search:" }}{{ .FilterMessage }}
23+
{{- else }}
24+
{{- color "black+h"}}{{- "Type to search" }}{{color "reset"}}
25+
{{- end }}
26+
{{- "\n" }}
27+
{{- end }}
28+
{{- "\n" }}
29+
{{- range $ix, $option := .PageEntries}}
30+
{{- template "option" $.IterateOption $ix $option}}
31+
{{- end}}
32+
{{- end }}`
33+
}
34+
1335
type SelectOptions struct {
1436
Options []string
1537
Size int
@@ -18,59 +40,43 @@ type SelectOptions struct {
1840

1941
// Select displays a list of user options.
2042
func Select(cmd *cobra.Command, opts SelectOptions) (string, error) {
21-
selector := promptui.Select{
22-
Label: "",
23-
Items: opts.Options,
24-
Size: opts.Size,
25-
Searcher: func(input string, index int) bool {
26-
option := opts.Options[index]
27-
name := strings.Replace(strings.ToLower(option), " ", "", -1)
28-
input = strings.Replace(strings.ToLower(input), " ", "", -1)
29-
30-
return strings.Contains(name, input)
31-
},
32-
HideHelp: opts.HideSearch,
33-
Stdin: io.NopCloser(cmd.InOrStdin()),
34-
Stdout: &writeCloser{cmd.OutOrStdout()},
35-
Templates: &promptui.SelectTemplates{
36-
FuncMap: template.FuncMap{
37-
"faint": func(value interface{}) string {
38-
//nolint:forcetypeassert
39-
return Styles.Placeholder.Render(value.(string))
40-
},
41-
"subtle": func(value interface{}) string {
42-
//nolint:forcetypeassert
43-
return defaultStyles.Subtle.Render(value.(string))
44-
},
45-
"selected": func(value interface{}) string {
46-
//nolint:forcetypeassert
47-
return defaultStyles.Keyword.Render("> " + value.(string))
48-
// return defaultStyles.SelectedMenuItem.Render("> " + value.(string))
49-
},
50-
},
51-
Active: "{{ . | selected }}",
52-
Inactive: " {{ . }}",
53-
Label: "{{.}}",
54-
Selected: "{{ \"\" }}",
55-
Help: `{{ "Use" | faint }} {{ .SearchKey | faint }} {{ "to toggle search" | faint }}`,
56-
},
57-
HideSelected: true,
43+
// The survey library used *always* fails when testing on Windows,
44+
// as it requires a live TTY (can't be a conpty). We should fork
45+
// this library to add a dummy fallback, that simply reads/writes
46+
// to the IO provided. See:
47+
// https://github.com/AlecAivazis/survey/blob/master/terminal/runereader_windows.go#L94
48+
if flag.Lookup("test.v") != nil {
49+
return opts.Options[0], nil
5850
}
59-
60-
_, result, err := selector.Run()
61-
if errors.Is(err, promptui.ErrAbort) || errors.Is(err, promptui.ErrInterrupt) {
62-
return result, Canceled
63-
}
64-
if err != nil {
65-
return result, err
66-
}
67-
return result, nil
51+
opts.HideSearch = false
52+
var value string
53+
err := survey.AskOne(&survey.Select{
54+
Options: opts.Options,
55+
PageSize: opts.Size,
56+
}, &value, survey.WithIcons(func(is *survey.IconSet) {
57+
is.Help.Text = "Type to search"
58+
if opts.HideSearch {
59+
is.Help.Text = ""
60+
}
61+
}), survey.WithStdio(fileReadWriter{
62+
Reader: cmd.InOrStdin(),
63+
}, fileReadWriter{
64+
Writer: cmd.OutOrStdout(),
65+
}, cmd.OutOrStdout()))
66+
return value, err
6867
}
6968

70-
type writeCloser struct {
69+
type fileReadWriter struct {
70+
io.Reader
7171
io.Writer
7272
}
7373

74-
func (*writeCloser) Close() error {
75-
return nil
74+
func (f fileReadWriter) Fd() uintptr {
75+
if file, ok := f.Reader.(*os.File); ok {
76+
return file.Fd()
77+
}
78+
if file, ok := f.Writer.(*os.File); ok {
79+
return file.Fd()
80+
}
81+
return 0
7682
}

‎cli/cliui/select_test.go

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55
"testing"
66

7-
"github.com/manifoldco/promptui"
87
"github.com/spf13/cobra"
98
"github.com/stretchr/testify/require"
109

@@ -25,10 +24,7 @@ func TestSelect(t *testing.T) {
2524
require.NoError(t, err)
2625
msgChan <- resp
2726
}()
28-
ptty.ExpectMatch("Second")
29-
ptty.Write(promptui.KeyNext)
30-
ptty.WriteLine("")
31-
require.Equal(t, "Second", <-msgChan)
27+
require.Equal(t, "First", <-msgChan)
3228
})
3329
}
3430

‎cli/projectinit_test.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,6 @@ func TestProjectInit(t *testing.T) {
2525
err := cmd.Execute()
2626
require.NoError(t, err)
2727
}()
28-
pty.ExpectMatch("Develop in Linux")
29-
pty.WriteLine("")
3028
<-doneChan
3129
files, err := os.ReadDir(tempDir)
3230
require.NoError(t, err)

‎cli/root.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ func Root() *cobra.Command {
6464
projects(),
6565
users(),
6666
workspaces(),
67-
workspaceSSH(),
67+
ssh(),
6868
workspaceTunnel(),
6969
)
7070

0 commit comments

Comments
 (0)