coder/cli/cliui/agent.go

package cliui

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/briandowns/spinner"
	"github.com/muesli/reflow/indent"
	"github.com/muesli/reflow/wordwrap"
	"golang.org/x/xerrors"

	"github.com/coder/coder/codersdk"
)

var AgentStartError = xerrors.New("agent startup exited with non-zero exit status")

type AgentOptions struct {
	WorkspaceName string
	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.
}

// Agent displays a spinning indicator that waits for a workspace agent to connect.
func Agent(ctx context.Context, writer io.Writer, 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
	agent, err := opts.Fetch(ctx)
	if err != nil {
		return xerrors.Errorf("fetch: %w", err)
	}

	// Fast path if the agent is ready (avoid showing connecting prompt).
	// We don't take the fast path for opts.NoWait yet because we want to
	// show the message.
	if agent.Status == codersdk.WorkspaceAgentConnected &&
		(agent.LoginBeforeReady || agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady) {
		return nil
	}

	ctx, cancel := signal.NotifyContext(ctx, os.Interrupt)
	defer cancel()

	spin := spinner.New(spinner.CharSets[78], 100*time.Millisecond, spinner.WithColor("fgHiGreen"))
	spin.Writer = writer
	spin.ForceOutput = true
	spin.Suffix = waitingMessage(agent, opts).Spin

	waitMessage := &message{}
	showMessage := func() {
		resourceMutex.Lock()
		defer resourceMutex.Unlock()

		m := waitingMessage(agent, opts)
		if m.Prompt == waitMessage.Prompt {
			return
		}
		moveUp := ""
		if waitMessage.Prompt != "" {
			// If this is an update, move a line up
			// to keep it tidy and aligned.
			moveUp = "\033[1A"
		}
		waitMessage = m

		// Stop the spinner while we write our message.
		spin.Stop()
		spin.Suffix = waitMessage.Spin
		// Clear the line and (if necessary) move up a line to write our message.
		_, _ = fmt.Fprintf(writer, "\033[2K%s\n%s\n", moveUp, waitMessage.Prompt)
		select {
		case <-ctx.Done():
		default:
			// Safe to resume operation.
			if spin.Suffix != "" {
				spin.Start()
			}
		}
	}

	// Fast path for showing the error message even when using no wait,
	// we do this just before starting the spinner to avoid needless
	// spinning.
	if agent.Status == codersdk.WorkspaceAgentConnected &&
		!agent.LoginBeforeReady && opts.NoWait {
		showMessage()
		return nil
	}

	// Start spinning after fast paths are handled.
	if spin.Suffix != "" {
		spin.Start()
	}
	defer spin.Stop()

	warnAfter := time.NewTimer(opts.WarnInterval)
	defer warnAfter.Stop()
	warningShown := make(chan struct{})
	go func() {
		select {
		case <-ctx.Done():
			close(warningShown)
		case <-warnAfter.C:
			close(warningShown)
			showMessage()
		}
	}()

	fetchInterval := time.NewTicker(opts.FetchInterval)
	defer fetchInterval.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-fetchInterval.C:
		}
		resourceMutex.Lock()
		agent, err = opts.Fetch(ctx)
		if err != nil {
			resourceMutex.Unlock()
			return xerrors.Errorf("fetch: %w", err)
		}
		resourceMutex.Unlock()
		switch agent.Status {
		case codersdk.WorkspaceAgentConnected:
			// 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.LoginBeforeReady && !opts.NoWait {
				switch agent.LifecycleState {
				case codersdk.WorkspaceAgentLifecycleReady:
					return nil
				case codersdk.WorkspaceAgentLifecycleStartTimeout:
					showMessage()
				case codersdk.WorkspaceAgentLifecycleStartError:
					showMessage()
					return AgentStartError
				default:
					select {
					case <-warningShown:
						showMessage()
					default:
						// This state is normal, we don't want
						// to show a message prematurely.
					}
				}
				continue
			}
			return nil
		case codersdk.WorkspaceAgentTimeout, codersdk.WorkspaceAgentDisconnected:
			showMessage()
		}
	}
}

type message struct {
	Spin         string
	Prompt       string
	Troubleshoot bool
}

func waitingMessage(agent codersdk.WorkspaceAgent, opts AgentOptions) (m *message) {
	m = &message{
		Spin:   fmt.Sprintf("Waiting for connection from %s...", Styles.Field.Render(agent.Name)),
		Prompt: "Don't panic, your workspace is booting up!",
	}
	defer func() {
		if agent.Status == codersdk.WorkspaceAgentConnected && opts.NoWait {
			m.Spin = ""
		}
		if m.Spin != "" {
			m.Spin = " " + m.Spin
		}

		// We don't want to wrap the troubleshooting URL, so we'll handle word
		// wrapping ourselves (vs using lipgloss).
		w := wordwrap.NewWriter(Styles.Paragraph.GetWidth() - Styles.Paragraph.GetMarginLeft()*2)
		w.Breakpoints = []rune{' ', '\n'}

		_, _ = fmt.Fprint(w, m.Prompt)
		if m.Troubleshoot {
			if agent.TroubleshootingURL != "" {
				_, _ = fmt.Fprintf(w, " See troubleshooting instructions at:\n%s", agent.TroubleshootingURL)
			} else {
				_, _ = fmt.Fprint(w, " Wait for it to (re)connect or restart your workspace.")
			}
		}
		_, _ = fmt.Fprint(w, "\n")

		// We want to prefix the prompt with a caret, but we want text on the
		// following lines to align with the text on the first line (i.e. added
		// spacing).
		ind := " " + Styles.Prompt.String()
		iw := indent.NewWriter(1, func(w io.Writer) {
			_, _ = w.Write([]byte(ind))
			ind = "   " // Set indentation to space after initial prompt.
		})
		_, _ = fmt.Fprint(iw, w.String())
		m.Prompt = iw.String()
	}()

	switch agent.Status {
	case codersdk.WorkspaceAgentTimeout:
		m.Prompt = "The workspace agent is having trouble connecting."
	case codersdk.WorkspaceAgentDisconnected:
		m.Prompt = "The workspace agent lost connection!"
	case codersdk.WorkspaceAgentConnected:
		m.Spin = fmt.Sprintf("Waiting for %s to become ready...", Styles.Field.Render(agent.Name))
		m.Prompt = "Don't panic, your workspace agent has connected and the workspace is getting ready!"
		if opts.NoWait {
			m.Prompt = "Your workspace is still getting ready, it may be in an incomplete state."
		}

		switch agent.LifecycleState {
		case codersdk.WorkspaceAgentLifecycleStartTimeout:
			m.Prompt = "The workspace is taking longer than expected to get ready, the agent startup script is still executing."
		case codersdk.WorkspaceAgentLifecycleStartError:
			m.Spin = ""
			m.Prompt = "The workspace ran into a problem while getting ready, the agent startup script exited with non-zero status."
		default:
			// Not a failure state, no troubleshooting necessary.
			return m
		}
	default:
		// Not a failure state, no troubleshooting necessary.
		return m
	}
	m.Troubleshoot = true
	return m
}