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
2022-03-28 18:19:28 -06:00 · 2022-03-28 18:19:28 -06:00 · 82dfd6c72f
parent 620c889842
commit 82dfd6c72f
26 changed files with 539 additions and 231 deletions
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@ -8,8 +8,10 @@
    "drpcconn",
    "drpcmux",
    "drpcserver",
+    "Dsts",
    "fatih",
    "goarch",
+    "gographviz",
    "goleak",
    "gossh",
    "hashicorp",
--- a/agent/agent_test.go
+++ b/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()
--- a/agent/conn.go
+++ b/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()
+}
--- a/cli/cliui/agent.go
+++ b/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
+	}
+}
--- a/cli/cliui/agent_test.go
+++ b/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
+}
--- a/cli/cliui/select.go
+++ b/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
 }
--- a/cli/cliui/select_test.go
+++ b/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)
 	})
 }

--- a/cli/projectinit_test.go
+++ b/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)
--- a/cli/root.go
+++ b/cli/root.go
@ -64,7 +64,7 @@ func Root() *cobra.Command {
 		projects(),
 		users(),
 		workspaces(),
-		workspaceSSH(),
+		ssh(),
 		workspaceTunnel(),
 	)

--- a/cli/ssh.go
+++ b/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
--- a/cli/ssh_test.go
+++ b/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
+	})
+}
--- a/cli/start.go
+++ b/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 {
--- a/cli/workspaceagent.go
+++ b/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
 }
--- a/cli/workspacecreate.go
+++ b/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
 		},
--- a/cli/workspaces.go
+++ b/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
--- a/cmd/cliui/main.go
+++ b/cmd/cliui/main.go
@ -1,6 +1,7 @@
 package main

 import (
+	"context"
 	"errors"
 	"fmt"
 	"os"
@ -62,10 +63,11 @@ func main() {
 	root.AddCommand(&cobra.Command{
 		Use: "select",
 		RunE: func(cmd *cobra.Command, args []string) error {
-			_, err := cliui.Select(cmd, cliui.SelectOptions{
+			value, err := cliui.Select(cmd, cliui.SelectOptions{
 				Options: []string{"Tomato", "Banana", "Onion", "Grape", "Lemon"},
 				Size:    3,
 			})
+			fmt.Printf("Selected: %q\n", value)
 			return err
 		},
 	})
@ -156,6 +158,35 @@ func main() {
 		},
 	})

+	root.AddCommand(&cobra.Command{
+		Use: "agent",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			resource := codersdk.WorkspaceResource{
+				Type: "google_compute_instance",
+				Name: "dev",
+				Agent: &codersdk.WorkspaceAgent{
+					Status: codersdk.WorkspaceAgentDisconnected,
+				},
+			}
+			go func() {
+				time.Sleep(3 * time.Second)
+				resource.Agent.Status = codersdk.WorkspaceAgentConnected
+			}()
+			err := cliui.Agent(cmd, cliui.AgentOptions{
+				WorkspaceName: "dev",
+				Fetch: func(ctx context.Context) (codersdk.WorkspaceResource, error) {
+					return resource, nil
+				},
+				WarnInterval: 2 * time.Second,
+			})
+			if err != nil {
+				return err
+			}
+			fmt.Printf("Completed!\n")
+			return nil
+		},
+	})
+
 	err := root.Execute()
 	if err != nil {
 		_, _ = fmt.Println(err.Error())
--- a/coderd/coderd.go
+++ b/coderd/coderd.go
@ -9,13 +9,14 @@ import (
 	"github.com/go-chi/chi/v5"
 	"google.golang.org/api/idtoken"

+	chitrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/go-chi/chi.v5"
+
 	"cdr.dev/slog"
 	"github.com/coder/coder/coderd/awsidentity"
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/site"
-	chitrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/go-chi/chi.v5"
 )

 // Options are requires parameters for Coder to start.
--- a/coderd/provisionerdaemons.go
+++ b/coderd/provisionerdaemons.go
@ -273,12 +273,6 @@ func (server *provisionerdServer) UpdateJob(ctx context.Context, request *proto.
 	if job.WorkerID.UUID.String() != server.ID.String() {
 		return nil, xerrors.New("you don't own this job")
 	}
-	if job.CanceledAt.Valid {
-		// Allows for graceful cancelation on the backend!
-		return &proto.UpdateJobResponse{
-			Canceled: true,
-		}, nil
-	}
 	err = server.Database.UpdateProvisionerJobByID(ctx, database.UpdateProvisionerJobByIDParams{
 		ID:        parsedID,
 		UpdatedAt: database.Now(),
@ -401,11 +395,14 @@ func (server *provisionerdServer) UpdateJob(ctx context.Context, request *proto.
 		}

 		return &proto.UpdateJobResponse{
+			Canceled:        job.CanceledAt.Valid,
 			ParameterValues: protoParameters,
 		}, nil
 	}

-	return &proto.UpdateJobResponse{}, nil
+	return &proto.UpdateJobResponse{
+		Canceled: job.CanceledAt.Valid,
+	}, nil
 }

 func (server *provisionerdServer) FailJob(ctx context.Context, failJob *proto.FailedJob) (*proto.Empty, error) {
@ -446,7 +443,7 @@ func (server *provisionerdServer) FailJob(ctx context.Context, failJob *proto.Fa
 			return nil, xerrors.Errorf("unmarshal workspace provision input: %w", err)
 		}
 		err = server.Database.UpdateWorkspaceBuildByID(ctx, database.UpdateWorkspaceBuildByIDParams{
-			ID:               jobID,
+			ID:               input.WorkspaceBuildID,
 			UpdatedAt:        database.Now(),
 			ProvisionerState: jobType.WorkspaceBuild.State,
 		})
--- a/codersdk/workspaceresources.go
+++ b/codersdk/workspaceresources.go
@ -133,15 +133,9 @@ func (c *Client) DialWorkspaceAgent(ctx context.Context, resource uuid.UUID, ice
 	if err != nil {
 		return nil, xerrors.Errorf("dial peer: %w", err)
 	}
-	go func() {
-		// The stream is kept alive to renegotiate the RTC connection
-		// if need-be. The calling context can be canceled to end
-		// the negotiation stream, but not the peer connection.
-		<-peerConn.Closed()
-		_ = conn.Close(websocket.StatusNormalClosure, "")
-	}()
 	return &agent.Conn{
-		Conn: peerConn,
+		Negotiator: client,
+		Conn:       peerConn,
 	}, nil
 }

--- a/examples/gcp-linux/main.tf
+++ b/examples/gcp-linux/main.tf
@ -8,7 +8,7 @@ terraform {
 }

 variable "service_account" {
-  description = <<EOT
+  description = <<EOF
 Coder requires a Google Cloud Service Account to provision workspaces.

 1. Create a service account:
@ -19,7 +19,7 @@ Coder requires a Google Cloud Service Account to provision workspaces.
 3. Click on the created key, and navigate to the "Keys" tab.
 4. Click "Add key", then "Create new key".
 5. Generate a JSON private key, and paste the contents in \'\' quotes below.
-EOT
+EOF
  sensitive   = true
 }

--- a/examples/gcp-windows/main.tf
+++ b/examples/gcp-windows/main.tf
@ -1,33 +1,48 @@
 terraform {
  required_providers {
    coder = {
-      source  = "coder/coder"
+      source = "coder/coder"
+      version = "0.2.1"
    }
  }
 }

-variable "gcp_credentials" {
-  sensitive = true
+variable "service_account" {
+  description = <<EOF
+Coder requires a Google Cloud Service Account to provision workspaces.
+
+1. Create a service account:
+   https://console.cloud.google.com/projectselector/iam-admin/serviceaccounts/create
+2. Add the roles:
+   - Compute Admin
+   - Service Account User
+3. Click on the created key, and navigate to the "Keys" tab.
+4. Click "Add key", then "Create new key".
+5. Generate a JSON private key, and paste the contents in \'\' quotes below.
+EOF
+  sensitive   = true
 }

-variable "gcp_project" {
-  description = "The Google Cloud project to manage resources in."
-}
-
-variable "gcp_region" {
-  default = "us-central1"
+variable "zone" {
+  description = "What region should your workspace live in?"
+  default     = "us-central1-a"
+  validation {
+    condition     = contains(["northamerica-northeast1-a", "us-central1-a", "us-west2-c", "europe-west4-b", "southamerica-east1-a"], var.zone)
+    error_message = "Invalid zone!"
+  }
 }

 provider "google" {
-  project = var.gcp_project
-  region = var.gcp_region
-  credentials = var.gcp_credentials
+  zone        = var.zone
+  credentials = var.service_account
+  project     = jsondecode(var.service_account).project_id
 }

 data "coder_workspace" "me" {
 }

 data "coder_agent_script" "dev" {
+  auth = "google-instance-identity"
  arch = "amd64"
  os   = "windows"
 }
@ -36,15 +51,24 @@ data "google_compute_default_service_account" "default" {
 }

 resource "random_string" "random" {
-  count = data.coder_workspace.me.transition == "start" ? 1 : 0
  length  = 8
  special = false
 }

+resource "google_compute_disk" "root" {
+  name  = "coder-${lower(random_string.random.result)}"
+  type  = "pd-ssd"
+  zone  = var.zone
+  image = "projects/windows-cloud/global/images/windows-server-2022-dc-core-v20220215"
+  lifecycle {
+    ignore_changes = [image]
+  }
+}
+
 resource "google_compute_instance" "dev" {
-  zone         = "us-central1-a"
+  zone         = var.zone
  count        = data.coder_workspace.me.transition == "start" ? 1 : 0
-  name         = "coder-${lower(random_string.random[0].result)}"
+  name         = "coder-${data.coder_workspace.me.owner}-${data.coder_workspace.me.name}"
  machine_type = "e2-medium"
  network_interface {
    network = "default"
@ -53,15 +77,14 @@ resource "google_compute_instance" "dev" {
    }
  }
  boot_disk {
-    initialize_params {
-      image = "projects/windows-cloud/global/images/windows-server-2022-dc-core-v20220215"
-    }
+    auto_delete = false
+    source = google_compute_disk.root.name
  }
  service_account {
    email  = data.google_compute_default_service_account.default.email
    scopes = ["cloud-platform"]
  }
-  metadata = {
+ metadata = {
    windows-startup-script-ps1 = data.coder_agent_script.dev.value
    serial-port-enable = "TRUE"
  }
@ -69,8 +92,5 @@ resource "google_compute_instance" "dev" {

 resource "coder_agent" "dev" {
  count = length(google_compute_instance.dev)
-  auth {
-    type        = "google-instance-identity"
-    instance_id = google_compute_instance.dev[0].instance_id
-  }
+  instance_id = google_compute_instance.dev[0].instance_id
 }
--- a/go.mod
+++ b/go.mod
@ -31,6 +31,8 @@ replace github.com/golang/glog => github.com/coder/glog v1.0.1-0.20220322161911-
 require (
 	cdr.dev/slog v1.4.1
 	cloud.google.com/go/compute v1.5.0
+	github.com/AlecAivazis/survey/v2 v2.3.4
+	github.com/awalterschulze/gographviz v2.0.3+incompatible
 	github.com/bgentry/speakeasy v0.1.0
 	github.com/briandowns/spinner v1.18.1
 	github.com/charmbracelet/charm v0.10.3
@ -53,7 +55,6 @@ require (
 	github.com/hashicorp/hcl/v2 v2.11.1
 	github.com/hashicorp/terraform-config-inspect v0.0.0-20211115214459-90acf1ca460f
 	github.com/hashicorp/terraform-exec v0.15.0
-	github.com/hashicorp/terraform-json v0.13.0
 	github.com/hashicorp/yamux v0.0.0-20211028200310-0bc27b27de87
 	github.com/justinas/nosurf v1.1.1
 	github.com/kirsle/configdir v0.0.0-20170128060238-e45d2f54772f
@ -71,6 +72,7 @@ require (
 	github.com/quasilyte/go-ruleguard/dsl v0.3.19
 	github.com/rs/zerolog v1.26.1
 	github.com/spf13/cobra v1.4.0
+	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.7.1
 	github.com/tabbed/pqtype v0.1.1
 	github.com/unrolled/secure v1.10.0
@ -81,7 +83,6 @@ require (
 	golang.org/x/oauth2 v0.0.0-20220309155454-6242fa91716a
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 	golang.org/x/sys v0.0.0-20220310020820-b874c991c1a5
-	golang.org/x/term v0.0.0-20210927222741-03fcf44c2211
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1
 	google.golang.org/api v0.73.0
 	google.golang.org/protobuf v1.28.0
@ -157,11 +158,13 @@ require (
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/hashicorp/go-multierror v1.1.1 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
+	github.com/hashicorp/terraform-json v0.13.0 // indirect
 	github.com/imdario/mergo v0.3.12 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/juju/ansiterm v0.0.0-20210929141451-8b71cc96ebdc // indirect
+	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
 	github.com/klauspost/compress v1.15.0 // indirect
 	github.com/leodido/go-urn v1.2.1 // indirect
 	github.com/lucas-clemente/quic-go v0.25.1-0.20220307142123-ad1cb27c1b64 // indirect
@ -174,6 +177,7 @@ require (
 	github.com/mattn/go-colorable v0.1.12 // indirect
 	github.com/mattn/go-runewidth v0.0.13 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 // indirect
+	github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b // indirect
 	github.com/miekg/dns v1.1.46 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	github.com/mitchellh/go-wordwrap v1.0.1 // indirect
@ -219,7 +223,6 @@ require (
 	github.com/spf13/afero v1.8.1 // indirect
 	github.com/spf13/cast v1.4.1 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
-	github.com/spf13/pflag v1.0.5
 	github.com/tinylib/msgp v1.1.2 // indirect
 	github.com/xeipuuv/gojsonpointer v0.0.0-20190905194746-02993c407bfb // indirect
 	github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect
@ -229,6 +232,7 @@ require (
 	go.opencensus.io v0.23.0 // indirect
 	golang.org/x/mod v0.5.1 // indirect
 	golang.org/x/net v0.0.0-20220225172249-27dd8689420f // indirect
+	golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	golang.org/x/time v0.0.0-20211116232009-f0f3c7e86c11 // indirect
 	golang.org/x/tools v0.1.9 // indirect
--- a/go.sum
+++ b/go.sum
@ -66,6 +66,8 @@ dmitri.shuralyov.com/html/belt v0.0.0-20180602232347-f7d459c86be0/go.mod h1:JLBr
 dmitri.shuralyov.com/service/change v0.0.0-20181023043359-a85b471d5412/go.mod h1:a1inKt/atXimZ4Mv927x+r7UpyzRUf4emIoiiSC2TN4=
 dmitri.shuralyov.com/state v0.0.0-20180228185332-28bcc343414c/go.mod h1:0PRwlb0D6DFvNNtx+9ybjezNCa8XF0xaYcETyp6rHWU=
 gioui.org v0.0.0-20210308172011-57750fc8a0a6/go.mod h1:RSH6KIUZ0p2xy5zHDxgAM4zumjgTw83q2ge/PI+yyw8=
+github.com/AlecAivazis/survey/v2 v2.3.4 h1:pchTU9rsLUSvWEl2Aq9Pv3k0IE2fkqtGxazskAMd9Ng=
+github.com/AlecAivazis/survey/v2 v2.3.4/go.mod h1:hrV6Y/kQCLhIZXGcriDCUBtB3wnN7156gMXJ3+b23xM=
 github.com/Azure/azure-pipeline-go v0.2.3/go.mod h1:x841ezTBIMG6O3lAcl8ATHnsOPVl2bqk7S3ta6S6u4k=
 github.com/Azure/azure-sdk-for-go v16.2.1+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=
 github.com/Azure/azure-storage-blob-go v0.14.0/go.mod h1:SMqIBi+SuiQH32bvyjngEewEeXoPfKMgWlBDaYf6fck=
@ -134,6 +136,8 @@ github.com/Microsoft/hcsshim v0.8.23/go.mod h1:4zegtUJth7lAvFyc6cH2gGQ5B3OFQim01
 github.com/Microsoft/hcsshim/test v0.0.0-20201218223536-d3e5debf77da/go.mod h1:5hlzMzRKMLyo42nCZ9oml8AdTlq/0cvIaBv6tK1RehU=
 github.com/Microsoft/hcsshim/test v0.0.0-20210227013316-43a75bb4edd3/go.mod h1:mw7qgWloBUl75W/gVH3cQszUg1+gUITj7D6NY7ywVnY=
 github.com/NYTimes/gziphandler v0.0.0-20170623195520-56545f4a5d46/go.mod h1:3wb06e3pkSAbeQ52E9H9iFoQsEEwGN64994WTCIhntQ=
+github.com/Netflix/go-expect v0.0.0-20220104043353-73e0943537d2 h1:+vx7roKuyA63nhn5WAunQHLTznkw5W8b1Xc0dNjp83s=
+github.com/Netflix/go-expect v0.0.0-20220104043353-73e0943537d2/go.mod h1:HBCaDeC1lPdgDeDbhX8XFpy1jqjK0IBG8W5K+xYqA0w=
 github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5 h1:TngWCqHvy9oXAN6lEVMRuU21PR1EtLVZJmdB18Gu3Rw=
 github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5/go.mod h1:lmUJ/7eu/Q8D7ML55dXQrVaamCz2vxCfdQBasLZfHKk=
 github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
@ -191,6 +195,8 @@ github.com/armon/go-radix v1.0.0/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgI
 github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs=
 github.com/asaskevich/govalidator v0.0.0-20190424111038-f61b66f89f4a/go.mod h1:lB+ZfQJz7igIIfQNfa7Ml4HSf2uFQQRzpGGRXenZAgY=
 github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
+github.com/awalterschulze/gographviz v2.0.3+incompatible h1:9sVEXJBJLwGX7EQVhLm2elIKCm7P2YHFC8v6096G09E=
+github.com/awalterschulze/gographviz v2.0.3+incompatible/go.mod h1:GEV5wmg4YquNw7v1kkyoX9etIk8yVmXj+AkDHuuETHs=
 github.com/aws/aws-sdk-go v1.15.11/go.mod h1:mFuSZ37Z9YOHbQEwBWztmVzqXrEkub65tZoCYDt7FT0=
 github.com/aws/aws-sdk-go v1.17.7/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=
 github.com/aws/aws-sdk-go v1.25.37/go.mod h1:KmX6BPdI08NWTb3/sm4ZGu5ShLoqVDhKgpiN924inxo=
@ -443,6 +449,7 @@ github.com/cpuguy83/go-md2man/v2 v2.0.1/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46t
 github.com/creack/pty v1.1.7/go.mod h1:lj5s0c3V2DBrqTV7llrYr5NG6My20zk30Fl46Y7DoTY=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/creack/pty v1.1.11/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
+github.com/creack/pty v1.1.17/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
 github.com/creack/pty v1.1.18 h1:n56/Zwd5o6whRC5PMGretI4IdRLlmBXYNjScPaBgsbY=
 github.com/creack/pty v1.1.18/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
 github.com/cyphar/filepath-securejoin v0.2.2/go.mod h1:FpkQEhXnPnOthhzymB7CGsFk2G9VLXONKD9G7QGMM+4=
@ -947,6 +954,8 @@ github.com/hashicorp/yamux v0.0.0-20180604194846-3520598351bb/go.mod h1:+NfK9FKe
 github.com/hashicorp/yamux v0.0.0-20181012175058-2f1d1f20f75d/go.mod h1:+NfK9FKeTrX5uv1uIXGdwYDTeHna2qgaIlx54MXqjAM=
 github.com/hashicorp/yamux v0.0.0-20211028200310-0bc27b27de87 h1:xixZ2bWeofWV68J+x6AzmKuVM/JWCQwkWm6GW/MUR6I=
 github.com/hashicorp/yamux v0.0.0-20211028200310-0bc27b27de87/go.mod h1:CtWFDAQgb7dxtzFs4tWbplKIe2jSi3+5vKbgIO0SLnQ=
+github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec h1:qv2VnGeEQHchGaZ/u7lxST/RaJw+cv273q79D81Xbog=
+github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec/go.mod h1:Q48J4R4DvxnHolD5P8pOtXigYlRuPLGl6moFx3ulM68=
 github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
@ -1069,6 +1078,7 @@ github.com/k0kubun/pp v2.3.0+incompatible/go.mod h1:GWse8YhT0p8pT4ir3ZgBbfZild3t
 github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8=
 github.com/karrick/godirwalk v1.8.0/go.mod h1:H5KPZjojv4lE+QYImBI8xVtrBRgYrIVsaRPx4tDPEn4=
 github.com/karrick/godirwalk v1.10.3/go.mod h1:RoGL9dQei4vP9ilrpETWE8CLOZ1kiN0LhBygSwrAsHA=
+github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs=
 github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8=
 github.com/kevinburke/ssh_config v0.0.0-20201106050909-4977a11b4351 h1:DowS9hvgyYSX4TO5NpyC606/Z4SxnNYbT+WX27or6Ck=
 github.com/kevinburke/ssh_config v0.0.0-20201106050909-4977a11b4351/go.mod h1:CT57kijsi8u/K/BOFA39wgDQJ9CxiF4nAY/ojJ6r6mM=
@ -1212,6 +1222,8 @@ github.com/mattn/go-sqlite3 v1.14.12/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 h1:I0XW9+e1XWDxdcEniV4rQAIOPUGDq67JSCiRCgGCZLI=
 github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4=
+github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4=
+github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE=
 github.com/microcosm-cc/bluemonday v1.0.1/go.mod h1:hsXNsILzKxV+sX77C5b8FSuKF00vh2OMYv+xgHpAMF4=
 github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
 github.com/miekg/dns v1.1.25/go.mod h1:bPDLeHnStXmXAq1m/Ch/hvfNHr14JKNPMBo3VZKjuso=
@ -2106,6 +2118,7 @@ golang.org/x/sys v0.0.0-20220310020820-b874c991c1a5/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210422114643-f5beecf764ed/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
+golang.org/x/term v0.0.0-20210503060354-a79de5458b56/go.mod h1:tfny5GFUkzUvx4ps4ajbZsCe5lw1metzhBm9T3x7oIY=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 h1:JGgROgKl9N8DuW20oFS5gxc+lE67/N3FcwmBPMe7ArY=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/text v0.0.0-20160726164857-2910a502d2bf/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
--- a/peerbroker/dial.go
+++ b/peerbroker/dial.go
@ -1,6 +1,9 @@
 package peerbroker

 import (
+	"context"
+	"errors"
+	"io"
 	"reflect"

 	"github.com/pion/webrtc/v3"
@ -54,6 +57,11 @@ func Dial(stream proto.DRPCPeerBroker_NegotiateConnectionClient, iceServers []we
 		for {
 			serverToClientMessage, err := stream.Recv()
 			if err != nil {
+				// p2p connections should never die if this stream does due
+				// to proper closure or context cancellation!
+				if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) {
+					return
+				}
 				_ = peerConn.CloseWithError(xerrors.Errorf("recv: %w", err))
 				return
 			}
--- a/peerbroker/listen.go
+++ b/peerbroker/listen.go
@ -166,8 +166,10 @@ func (b *peerBrokerService) NegotiateConnection(stream proto.DRPCPeerBroker_Nego
 	for {
 		clientToServerMessage, err := stream.Recv()
 		if err != nil {
-			if errors.Is(err, io.EOF) {
-				break
+			// p2p connections should never die if this stream does due
+			// to proper closure or context cancellation!
+			if errors.Is(err, io.EOF) || errors.Is(err, context.Canceled) {
+				return nil
 			}
 			return peerConn.CloseWithError(xerrors.Errorf("recv: %w", err))
 		}
@ -186,6 +188,4 @@ func (b *peerBrokerService) NegotiateConnection(stream proto.DRPCPeerBroker_Nego
 			return peerConn.CloseWithError(xerrors.Errorf("unhandled message: %s", reflect.TypeOf(clientToServerMessage).String()))
 		}
 	}
-
-	return nil
 }
--- a/provisioner/terraform/provision.go
+++ b/provisioner/terraform/provision.go
@ -11,8 +11,8 @@ import (
 	"path/filepath"
 	"strings"

+	"github.com/awalterschulze/gographviz"
 	"github.com/hashicorp/terraform-exec/tfexec"
-	tfjson "github.com/hashicorp/terraform-json"
 	"github.com/mitchellh/mapstructure"
 	"golang.org/x/xerrors"

@ -80,7 +80,7 @@ func (t *terraform) Provision(stream proto.DRPCProvisioner_ProvisionStream) erro
 			_ = stream.Send(&proto.Provision_Response{
 				Type: &proto.Provision_Response_Log{
 					Log: &proto.Log{
-						Level:  proto.LogLevel_INFO,
+						Level:  proto.LogLevel_DEBUG,
 						Output: scanner.Text(),
 					},
 				},
@ -140,7 +140,6 @@ func (t *terraform) Provision(stream proto.DRPCProvisioner_ProvisionStream) erro
 			if err != nil {
 				return
 			}
-
 			logLevel, err := convertTerraformLogLevel(log.Level)
 			if err != nil {
 				// Not a big deal, but we should handle this at some point!
@ -209,7 +208,7 @@ func (t *terraform) Provision(stream proto.DRPCProvisioner_ProvisionStream) erro
 		case <-stream.Context().Done():
 			return
 		case <-shutdown.Done():
-			_ = cmd.Process.Signal(os.Kill)
+			_ = cmd.Process.Signal(os.Interrupt)
 		}
 	}()
 	cmd.Stdout = writer
@ -266,27 +265,13 @@ func parseTerraformPlan(ctx context.Context, terraform *tfexec.Terraform, planfi
 		return nil, xerrors.Errorf("show terraform plan file: %w", err)
 	}

-	// Maps resource dependencies to expression references.
-	// This is *required* for a plan, because "DependsOn"
-	// does not propagate.
-	resourceDependencies := map[string][]string{}
-	for _, resource := range plan.Config.RootModule.Resources {
-		if resource.Expressions == nil {
-			resource.Expressions = map[string]*tfjson.Expression{}
-		}
-		// Count expression is separated for logical reasons,
-		// but it's simpler syntactically for us to combine here.
-		if resource.CountExpression != nil {
-			resource.Expressions["count"] = resource.CountExpression
-		}
-		for _, expression := range resource.Expressions {
-			dependencies, exists := resourceDependencies[resource.Address]
-			if !exists {
-				dependencies = []string{}
-			}
-			dependencies = append(dependencies, expression.References...)
-			resourceDependencies[resource.Address] = dependencies
-		}
+	rawGraph, err := terraform.Graph(ctx)
+	if err != nil {
+		return nil, xerrors.Errorf("graph: %w", err)
+	}
+	resourceDependencies, err := findDirectDependencies(rawGraph)
+	if err != nil {
+		return nil, xerrors.Errorf("find dependencies: %w", err)
 	}

 	resources := make([]*proto.Resource, 0)
@ -322,31 +307,21 @@ func parseTerraformPlan(ctx context.Context, terraform *tfexec.Terraform, planfi

 		agents[resource.Address] = agent
 	}
-
 	for _, resource := range plan.PlannedValues.RootModule.Resources {
 		if resource.Type == "coder_agent" {
 			continue
 		}
-		// The resource address on planned values can include the indexed
-		// value like "[0]", but the config doesn't have these, and we don't
-		// care which index the resource is.
-		resourceAddress := fmt.Sprintf("%s.%s", resource.Type, resource.Name)
-		var agent *proto.Agent
-		// Associate resources that depend on an agent.
-		for _, dependency := range resourceDependencies[resourceAddress] {
-			var has bool
-			agent, has = agents[dependency]
-			if has {
-				break
-			}
+		resourceKey := strings.Join([]string{resource.Type, resource.Name}, ".")
+		resourceNode, exists := resourceDependencies[resourceKey]
+		if !exists {
+			continue
 		}
-		// Associate resources where the agent depends on it.
-		for agentAddress := range agents {
-			for _, depend := range resourceDependencies[agentAddress] {
-				if depend != resourceAddress {
-					continue
-				}
-				agent = agents[agentAddress]
+		// Associate resources that depend on an agent.
+		var agent *proto.Agent
+		for _, dep := range resourceNode {
+			var has bool
+			agent, has = agents[dep]
+			if has {
 				break
 			}
 		}
@ -378,6 +353,14 @@ func parseTerraformApply(ctx context.Context, terraform *tfexec.Terraform, state
 	}
 	resources := make([]*proto.Resource, 0)
 	if state.Values != nil {
+		rawGraph, err := terraform.Graph(ctx)
+		if err != nil {
+			return nil, xerrors.Errorf("graph: %w", err)
+		}
+		resourceDependencies, err := findDirectDependencies(rawGraph)
+		if err != nil {
+			return nil, xerrors.Errorf("find dependencies: %w", err)
+		}
 		type agentAttributes struct {
 			ID            string            `mapstructure:"id"`
 			Token         string            `mapstructure:"token"`
@ -386,7 +369,6 @@ func parseTerraformApply(ctx context.Context, terraform *tfexec.Terraform, state
 			StartupScript string            `mapstructure:"startup_script"`
 		}
 		agents := map[string]*proto.Agent{}
-		agentDepends := map[string][]string{}

 		// Store all agents inside the maps!
 		for _, resource := range state.Values.RootModule.Resources {
@ -413,34 +395,26 @@ func parseTerraformApply(ctx context.Context, terraform *tfexec.Terraform, state
 			}
 			resourceKey := strings.Join([]string{resource.Type, resource.Name}, ".")
 			agents[resourceKey] = agent
-			agentDepends[resourceKey] = resource.DependsOn
 		}

 		for _, resource := range state.Values.RootModule.Resources {
 			if resource.Type == "coder_agent" {
 				continue
 			}
-			var agent *proto.Agent
+			resourceKey := strings.Join([]string{resource.Type, resource.Name}, ".")
+			resourceNode, exists := resourceDependencies[resourceKey]
+			if !exists {
+				continue
+			}
 			// Associate resources that depend on an agent.
-			for _, dep := range resource.DependsOn {
+			var agent *proto.Agent
+			for _, dep := range resourceNode {
 				var has bool
 				agent, has = agents[dep]
 				if has {
 					break
 				}
 			}
-			if agent == nil {
-				// Associate resources where the agent depends on it.
-				for agentKey, dependsOn := range agentDepends {
-					for _, depend := range dependsOn {
-						if depend != strings.Join([]string{resource.Type, resource.Name}, ".") {
-							continue
-						}
-						agent = agents[agentKey]
-						break
-					}
-				}
-			}

 			resources = append(resources, &proto.Resource{
 				Name:  resource.Name,
@ -489,3 +463,46 @@ func convertTerraformLogLevel(logLevel string) (proto.LogLevel, error) {
 		return proto.LogLevel(0), xerrors.Errorf("invalid log level %q", logLevel)
 	}
 }
+
+// findDirectDependencies maps Terraform resources to their parent and
+// children nodes. This parses GraphViz output from Terraform which
+// certainly is not ideal, but seems reliable.
+func findDirectDependencies(rawGraph string) (map[string][]string, error) {
+	parsedGraph, err := gographviz.ParseString(rawGraph)
+	if err != nil {
+		return nil, xerrors.Errorf("parse graph: %w", err)
+	}
+	graph, err := gographviz.NewAnalysedGraph(parsedGraph)
+	if err != nil {
+		return nil, xerrors.Errorf("analyze graph: %w", err)
+	}
+	direct := map[string][]string{}
+	for _, node := range graph.Nodes.Nodes {
+		label, exists := node.Attrs["label"]
+		if !exists {
+			continue
+		}
+		label = strings.Trim(label, `"`)
+
+		dependencies := make([]string, 0)
+		for _, edges := range []map[string][]*gographviz.Edge{
+			graph.Edges.SrcToDsts[node.Name],
+			graph.Edges.DstToSrcs[node.Name],
+		} {
+			for destination := range edges {
+				dependencyNode, exists := graph.Nodes.Lookup[destination]
+				if !exists {
+					continue
+				}
+				label, exists := dependencyNode.Attrs["label"]
+				if !exists {
+					continue
+				}
+				label = strings.Trim(label, `"`)
+				dependencies = append(dependencies, label)
+			}
+		}
+		direct[label] = dependencies
+	}
+	return direct, nil
+}