Skip to content

Commit ec14004

Browse files
mafredrimafredri
committed
feat(cli/cliui): add agent log streaming and new format
1 parent 117215f commit ec14004

File tree

8 files changed

+257
-220
lines changed

8 files changed

+257
-220
lines changed

‎cli/cliui/agent.go

Lines changed: 119 additions & 196 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,10 @@ package cliui
22

33
import (
44
"context"
5-
"fmt"
65
"io"
7-
"os"
8-
"os/signal"
9-
"sync"
106
"time"
117

12-
"github.com/briandowns/spinner"
13-
"github.com/muesli/reflow/indent"
14-
"github.com/muesli/reflow/wordwrap"
8+
"github.com/google/uuid"
159
"golang.org/x/xerrors"
1610

1711
"github.com/coder/coder/codersdk"
@@ -25,6 +19,7 @@ var (
2519
type AgentOptions struct {
2620
WorkspaceName string
2721
Fetch func(context.Context) (codersdk.WorkspaceAgent, error)
22+
FetchLogs func(ctx context.Context, agentID uuid.UUID, after int64, follow bool) (<-chan []codersdk.WorkspaceAgentStartupLog, io.Closer, error)
2823
FetchInterval time.Duration
2924
WarnInterval time.Duration
3025
Wait bool // If true, wait for the agent to be ready (startup script).
@@ -38,227 +33,155 @@ func Agent(ctx context.Context, writer io.Writer, opts AgentOptions) error {
3833
if opts.WarnInterval == 0 {
3934
opts.WarnInterval = 30 * time.Second
4035
}
41-
var resourceMutex sync.Mutex
36+
4237
agent, err := opts.Fetch(ctx)
4338
if err != nil {
4439
return xerrors.Errorf("fetch: %w", err)
4540
}
4641

47-
// Fast path if the agent is ready (avoid showing connecting prompt).
48-
// We don't take the fast path for opts.NoWait yet because we want to
49-
// show the message.
50-
if agent.Status == codersdk.WorkspaceAgentConnected &&
51-
(agent.StartupScriptBehavior == codersdk.WorkspaceAgentStartupScriptBehaviorNonBlocking || agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady) {
52-
return nil
42+
// TODO(mafredri): Rewrite this, we don't want to spam a fetch after every break statement.
43+
fetch := func(err *error) bool {
44+
agent, *err = opts.Fetch(ctx)
45+
return *err == nil
5346
}
5447

55-
ctx, cancel := signal.NotifyContext(ctx, os.Interrupt)
56-
defer cancel()
48+
sw := &stageWriter{w: writer}
5749

58-
spin := spinner.New(spinner.CharSets[78], 100*time.Millisecond, spinner.WithColor("fgHiGreen"))
59-
spin.Writer = writer
60-
spin.ForceOutput = true
61-
spin.Suffix = waitingMessage(agent, opts).Spin
50+
showInitialConnection := true
51+
showStartupLogs := false
6252

63-
waitMessage := &message{}
64-
showMessage := func() {
65-
resourceMutex.Lock()
66-
defer resourceMutex.Unlock()
53+
printInitialConnection := func() error {
54+
showInitialConnection = false
6755

68-
m := waitingMessage(agent, opts)
69-
if m.Prompt == waitMessage.Prompt {
70-
return
71-
}
72-
moveUp := ""
73-
if waitMessage.Prompt != "" {
74-
// If this is an update, move a line up
75-
// to keep it tidy and aligned.
76-
moveUp = "\033[1A"
77-
}
78-
waitMessage = m
56+
// Since we were waiting for the agent to connect, also show
57+
// startup logs.
58+
showStartupLogs = true
7959

80-
// Stop the spinner while we write our message.
81-
spin.Stop()
82-
spin.Suffix = waitMessage.Spin
83-
// Clear the line and (if necessary) move up a line to write our message.
84-
_, _ = fmt.Fprintf(writer, "\033[2K%s\n%s\n", moveUp, waitMessage.Prompt)
85-
select {
86-
case <-ctx.Done():
87-
default:
88-
// Safe to resume operation.
89-
if spin.Suffix != "" {
90-
spin.Start()
60+
stage := "Waiting for initial connection from the workspace agent"
61+
sw.Start(stage)
62+
if agent.Status == codersdk.WorkspaceAgentConnecting {
63+
for fetch(&err) {
64+
if agent.Status != codersdk.WorkspaceAgentConnecting {
65+
break
66+
}
67+
time.Sleep(opts.FetchInterval)
68+
}
69+
if err != nil {
70+
return xerrors.Errorf("fetch: %w", err)
9171
}
9272
}
93-
}
94-
95-
// Fast path for showing the error message even when using no wait,
96-
// we do this just before starting the spinner to avoid needless
97-
// spinning.
98-
if agent.Status == codersdk.WorkspaceAgentConnected &&
99-
agent.StartupScriptBehavior == codersdk.WorkspaceAgentStartupScriptBehaviorBlocking && !opts.Wait {
100-
showMessage()
73+
if agent.Status == codersdk.WorkspaceAgentTimeout {
74+
now := time.Now()
75+
sw.Log(now, codersdk.LogLevelInfo, "The workspace agent is having trouble connecting, we will keep trying to reach it")
76+
sw.Log(now, codersdk.LogLevelInfo, "For more information and troubleshooting, see https://coder.com/docs/v2/latest/templates#agent-connection-issues")
77+
for fetch(&err) {
78+
if agent.Status != codersdk.WorkspaceAgentConnecting && agent.Status != codersdk.WorkspaceAgentTimeout {
79+
break
80+
}
81+
time.Sleep(opts.FetchInterval)
82+
}
83+
if err != nil {
84+
return xerrors.Errorf("fetch: %w", err)
85+
}
86+
}
87+
sw.Complete(stage, agent.FirstConnectedAt.Sub(agent.CreatedAt))
10188
return nil
10289
}
10390

104-
// Start spinning after fast paths are handled.
105-
if spin.Suffix != "" {
106-
spin.Start()
107-
}
108-
defer spin.Stop()
91+
printLogs := func(follow bool) error {
92+
logStream, logsCloser, err := opts.FetchLogs(ctx, agent.ID, 0, follow)
93+
if err != nil {
94+
return xerrors.Errorf("fetch logs: %w", err)
95+
}
96+
defer logsCloser.Close()
10997

110-
warnAfter := time.NewTimer(opts.WarnInterval)
111-
defer warnAfter.Stop()
112-
warningShown := make(chan struct{})
113-
go func() {
114-
select {
115-
case <-ctx.Done():
116-
close(warningShown)
117-
case <-warnAfter.C:
118-
close(warningShown)
119-
showMessage()
98+
for logs := range logStream {
99+
for _, log := range logs {
100+
sw.Log(log.CreatedAt, log.Level, log.Output)
101+
}
120102
}
121-
}()
122103

123-
fetchInterval := time.NewTicker(opts.FetchInterval)
124-
defer fetchInterval.Stop()
104+
return nil
105+
}
106+
125107
for {
126-
select {
127-
case <-ctx.Done():
128-
return ctx.Err()
129-
case <-fetchInterval.C:
130-
}
131-
resourceMutex.Lock()
132-
agent, err = opts.Fetch(ctx)
133-
if err != nil {
134-
resourceMutex.Unlock()
135-
return xerrors.Errorf("fetch: %w", err)
136-
}
137-
resourceMutex.Unlock()
108+
// TODO(mafredri): Handle shutting down lifecycle states.
109+
138110
switch agent.Status {
111+
case codersdk.WorkspaceAgentConnecting, codersdk.WorkspaceAgentTimeout:
112+
err := printInitialConnection()
113+
if err != nil {
114+
return xerrors.Errorf("initial connection: %w", err)
115+
}
116+
139117
case codersdk.WorkspaceAgentConnected:
140-
// NOTE(mafredri): Once we have access to the workspace agent's
141-
// startup script logs, we can show them here.
142-
// https://github.com/coder/coder/issues/2957
143-
if agent.StartupScriptBehavior == codersdk.WorkspaceAgentStartupScriptBehaviorBlocking && opts.Wait {
144-
switch agent.LifecycleState {
145-
case codersdk.WorkspaceAgentLifecycleReady:
146-
return nil
147-
case codersdk.WorkspaceAgentLifecycleStartTimeout:
148-
showMessage()
149-
case codersdk.WorkspaceAgentLifecycleStartError:
150-
showMessage()
151-
return AgentStartError
152-
case codersdk.WorkspaceAgentLifecycleShuttingDown, codersdk.WorkspaceAgentLifecycleShutdownTimeout,
153-
codersdk.WorkspaceAgentLifecycleShutdownError, codersdk.WorkspaceAgentLifecycleOff:
154-
showMessage()
155-
return AgentShuttingDown
156-
default:
157-
select {
158-
case <-warningShown:
159-
showMessage()
160-
default:
161-
// This state is normal, we don't want
162-
// to show a message prematurely.
163-
}
118+
if !showStartupLogs && agent.LifecycleState == codersdk.WorkspaceAgentLifecycleReady {
119+
// The workspace is ready, there's nothing to do but connect.
120+
return nil
121+
}
122+
if showInitialConnection {
123+
err := printInitialConnection()
124+
if err != nil {
125+
return xerrors.Errorf("initial connection: %w", err)
164126
}
165-
continue
166127
}
167-
return nil
168-
case codersdk.WorkspaceAgentTimeout, codersdk.WorkspaceAgentDisconnected:
169-
showMessage()
170-
}
171-
}
172-
}
173128

174-
type message struct {
175-
Spin string
176-
Prompt string
177-
Troubleshoot bool
178-
}
179-
180-
func waitingMessage(agent codersdk.WorkspaceAgent, opts AgentOptions) (m *message) {
181-
m = &message{
182-
Spin: fmt.Sprintf("Waiting for connection from %s...", DefaultStyles.Field.Render(agent.Name)),
183-
Prompt: "Don't panic, your workspace is booting up!",
184-
}
185-
defer func() {
186-
if agent.Status == codersdk.WorkspaceAgentConnected && !opts.Wait {
187-
m.Spin = ""
188-
}
189-
if m.Spin != "" {
190-
m.Spin = " " + m.Spin
191-
}
129+
stage := "Running workspace agent startup script"
130+
follow := opts.Wait
131+
if !follow {
132+
stage += " (non-blocking)"
133+
}
134+
sw.Start(stage)
192135

193-
// We don't want to wrap the troubleshooting URL, so we'll handle word
194-
// wrapping ourselves (vs using lipgloss).
195-
w := wordwrap.NewWriter(DefaultStyles.Paragraph.GetWidth() - DefaultStyles.Paragraph.GetMarginLeft()*2)
196-
w.Breakpoints = []rune{' ', '\n'}
136+
err = printLogs(follow)
137+
if err != nil {
138+
return xerrors.Errorf("print logs: %w", err)
139+
}
197140

198-
_, _ = fmt.Fprint(w, m.Prompt)
199-
if m.Troubleshoot {
200-
if agent.TroubleshootingURL != "" {
201-
_, _ = fmt.Fprintf(w, " See troubleshooting instructions at:\n%s", agent.TroubleshootingURL)
202-
} else {
203-
_, _ = fmt.Fprint(w, " Wait for it to (re)connect or restart your workspace.")
141+
for fetch(&err) {
142+
if !follow || !agent.LifecycleState.Starting() {
143+
break
144+
}
145+
time.Sleep(opts.FetchInterval)
146+
}
147+
if err != nil {
148+
return xerrors.Errorf("fetch: %w", err)
204149
}
205-
}
206-
_, _ = fmt.Fprint(w, "\n")
207150

208-
// We want to prefix the prompt with a caret, but we want text on the
209-
// following lines to align with the text on the first line (i.e. added
210-
// spacing).
211-
ind := " " + DefaultStyles.Prompt.String()
212-
iw := indent.NewWriter(1, func(w io.Writer) {
213-
_, _ = w.Write([]byte(ind))
214-
ind = " " // Set indentation to space after initial prompt.
215-
})
216-
_, _ = fmt.Fprint(iw, w.String())
217-
m.Prompt = iw.String()
218-
}()
151+
switch agent.LifecycleState {
152+
case codersdk.WorkspaceAgentLifecycleReady:
153+
sw.Complete(stage, agent.ReadyAt.Sub(*agent.StartedAt))
154+
case codersdk.WorkspaceAgentLifecycleStartError:
155+
sw.Log(time.Time{}, codersdk.LogLevelWarn, "Warning: The startup script exited with an error and your workspace may be incomplete.")
156+
sw.Log(time.Time{}, codersdk.LogLevelWarn, "For more information and troubleshooting, see https://coder.com/docs/v2/latest/templates#startup-script-exited-with-an-error")
157+
sw.Fail(stage, agent.ReadyAt.Sub(*agent.StartedAt))
158+
default:
159+
sw.Log(time.Time{}, codersdk.LogLevelWarn, "Notice: The startup script is still running and your workspace may be incomplete.")
160+
sw.Log(time.Time{}, codersdk.LogLevelWarn, "For more information and troubleshooting, see https://coder.com/docs/v2/latest/templates#your-workspace-may-be-incomplete")
161+
// Note: We don't complete or fail the stage here, it's
162+
// intentionally left open to indicate this stage never
163+
// completed.
164+
}
219165

220-
switch agent.Status {
221-
case codersdk.WorkspaceAgentTimeout:
222-
m.Prompt = "The workspace agent is having trouble connecting."
223-
case codersdk.WorkspaceAgentDisconnected:
224-
m.Prompt = "The workspace agent lost connection!"
225-
case codersdk.WorkspaceAgentConnected:
226-
m.Spin = fmt.Sprintf("Waiting for %s to become ready...", DefaultStyles.Field.Render(agent.Name))
227-
m.Prompt = "Don't panic, your workspace agent has connected and the workspace is getting ready!"
228-
if !opts.Wait {
229-
m.Prompt = "Your workspace is still getting ready, it may be in an incomplete state."
230-
}
166+
return nil
231167

232-
switch agent.LifecycleState {
233-
case codersdk.WorkspaceAgentLifecycleStartTimeout:
234-
m.Prompt = "The workspace is taking longer than expected to get ready, the agent startup script is still executing."
235-
case codersdk.WorkspaceAgentLifecycleStartError:
236-
m.Spin = ""
237-
m.Prompt = "The workspace ran into a problem while getting ready, the agent startup script exited with non-zero status."
238-
default:
239-
switch agent.LifecycleState {
240-
case codersdk.WorkspaceAgentLifecycleShutdownTimeout:
241-
m.Spin = ""
242-
m.Prompt = "The workspace is shutting down, but is taking longer than expected to shut down and the agent shutdown script is still executing."
243-
m.Troubleshoot = true
244-
case codersdk.WorkspaceAgentLifecycleShutdownError:
245-
m.Spin = ""
246-
m.Prompt = "The workspace ran into a problem while shutting down, the agent shutdown script exited with non-zero status."
247-
m.Troubleshoot = true
248-
case codersdk.WorkspaceAgentLifecycleShuttingDown:
249-
m.Spin = ""
250-
m.Prompt = "The workspace is shutting down."
251-
case codersdk.WorkspaceAgentLifecycleOff:
252-
m.Spin = ""
253-
m.Prompt = "The workspace is not running."
168+
case codersdk.WorkspaceAgentDisconnected:
169+
// Since we were waiting for the agent to reconnect, also show
170+
// startup logs.
171+
showStartupLogs = true
172+
showInitialConnection = false
173+
174+
stage := "The workspace agent lost connection, waiting for it to reconnect"
175+
sw.Start(stage)
176+
for fetch(&err) {
177+
if agent.Status != codersdk.WorkspaceAgentDisconnected {
178+
break
179+
}
180+
}
181+
if err != nil {
182+
return xerrors.Errorf("fetch: %w", err)
254183
}
255-
// Not a failure state, no troubleshooting necessary.
256-
return m
184+
sw.Complete(stage, agent.LastConnectedAt.Sub(*agent.DisconnectedAt))
257185
}
258-
default:
259-
// Not a failure state, no troubleshooting necessary.
260-
return m
261186
}
262-
m.Troubleshoot = true
263-
return m
264187
}

0 commit comments

Comments
 (0)