Skip to content

Commit 049f063

Browse files
mafredrimafredri
committed
feat(cli/cliui): add agent log streaming and new format
1 parent 29b7a0a commit 049f063

File tree

9 files changed

+604
-563
lines changed

9 files changed

+604
-563
lines changed

‎cli/cliui/agent.go

Lines changed: 182 additions & 206 deletions
Large diffs are not rendered by default.

‎cli/cliui/agent_test.go

Lines changed: 273 additions & 322 deletions
Large diffs are not rendered by default.

‎cli/cliui/provisionerjob.go

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -204,7 +204,7 @@ type stageWriter struct {
204204
}
205205

206206
func (s *stageWriter) Start(stage string) {
207-
_, _ = fmt.Fprintf(s.w, " => ⧗ %s\n", stage)
207+
_, _ = fmt.Fprintf(s.w, "==> ⧗ %s\n", stage)
208208
}
209209

210210
func (s *stageWriter) Complete(stage string, duration time.Duration) {
@@ -227,7 +227,7 @@ func (s *stageWriter) end(stage string, duration time.Duration, ok bool) {
227227
if duration < 0 {
228228
duration = 0
229229
}
230-
_, _ = fmt.Fprintf(s.w, " == %s %s [%dms]\n", mark, stage, duration.Milliseconds())
230+
_, _ = fmt.Fprintf(s.w, "=== %s %s [%dms]\n", mark, stage, duration.Milliseconds())
231231
}
232232

233233
func (s *stageWriter) Log(createdAt time.Time, level codersdk.LogLevel, line string) {
@@ -238,6 +238,12 @@ func (s *stageWriter) Log(createdAt time.Time, level codersdk.LogLevel, line str
238238

239239
render := func(s ...string) string { return strings.Join(s, " ") }
240240

241+
var lines []string
242+
if !createdAt.IsZero() {
243+
lines = append(lines, createdAt.Local().Format("2006-01-02 15:04:05.000Z07:00"))
244+
}
245+
lines = append(lines, line)
246+
241247
switch level {
242248
case codersdk.LogLevelTrace, codersdk.LogLevelDebug:
243249
if !s.verbose {
@@ -250,7 +256,7 @@ func (s *stageWriter) Log(createdAt time.Time, level codersdk.LogLevel, line str
250256
render = DefaultStyles.Warn.Render
251257
case codersdk.LogLevelInfo:
252258
}
253-
_, _ = fmt.Fprintf(w, " %s\n", render(fmt.Sprintf("%s %s", createdAt.Local().Format("2006-01-02 15:04:05.000Z07:00"), line)))
259+
_, _ = fmt.Fprintf(w, "%s\n", render(lines...))
254260
}
255261

256262
func (s *stageWriter) flushLogs() {

‎cli/portforward.go

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,6 @@ func (r *RootCmd) portForward() *clibase.Cmd {
9191
}
9292

9393
err = cliui.Agent(ctx, inv.Stderr, cliui.AgentOptions{
94-
WorkspaceName: workspace.Name,
9594
Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
9695
return client.WorkspaceAgent(ctx, workspaceAgent.ID)
9796
},

‎cli/speedtest.go

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,19 +35,18 @@ func (r *RootCmd) speedtest() *clibase.Cmd {
3535
ctx, cancel := context.WithCancel(inv.Context())
3636
defer cancel()
3737

38-
workspace, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, codersdk.Me, inv.Args[0])
38+
_, workspaceAgent, err := getWorkspaceAndAgent(ctx, inv, client, codersdk.Me, inv.Args[0])
3939
if err != nil {
4040
return err
4141
}
4242

4343
err = cliui.Agent(ctx, inv.Stderr, cliui.AgentOptions{
44-
WorkspaceName: workspace.Name,
4544
Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
4645
return client.WorkspaceAgent(ctx, workspaceAgent.ID)
4746
},
4847
Wait: false,
4948
})
50-
if err != nil && !xerrors.Is(err, cliui.AgentStartError) {
49+
if err != nil {
5150
return xerrors.Errorf("await agent: %w", err)
5251
}
5352

‎cli/ssh.go

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -176,23 +176,16 @@ func (r *RootCmd) ssh() *clibase.Cmd {
176176
// OpenSSH passes stderr directly to the calling TTY.
177177
// This is required in "stdio" mode so a connecting indicator can be displayed.
178178
err = cliui.Agent(ctx, inv.Stderr, cliui.AgentOptions{
179-
WorkspaceName: workspace.Name,
180179
Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
181180
return client.WorkspaceAgent(ctx, workspaceAgent.ID)
182181
},
183-
Wait: wait,
182+
FetchLogs: client.WorkspaceAgentStartupLogsAfter,
183+
Wait: wait,
184184
})
185185
if err != nil {
186186
if xerrors.Is(err, context.Canceled) {
187187
return cliui.Canceled
188188
}
189-
if !xerrors.Is(err, cliui.AgentStartError) {
190-
return xerrors.Errorf("await agent: %w", err)
191-
}
192-
193-
// We don't want to fail on a startup script error because it's
194-
// natural that the user will want to fix the script and try again.
195-
// We don't print the error because cliui.Agent does that for us.
196189
}
197190

198191
if r.disableDirect {

‎cmd/cliui/main.go

Lines changed: 85 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,14 @@ import (
55
"errors"
66
"fmt"
77
"io"
8+
"math/rand"
89
"net/url"
910
"os"
1011
"strings"
1112
"sync/atomic"
1213
"time"
1314

15+
"github.com/google/uuid"
1416
"golang.org/x/xerrors"
1517

1618
"github.com/coder/coder/cli/clibase"
@@ -164,25 +166,91 @@ func main() {
164166
root.Children = append(root.Children, &clibase.Cmd{
165167
Use: "agent",
166168
Handler: func(inv *clibase.Invocation) error {
167-
agent := codersdk.WorkspaceAgent{
168-
Status: codersdk.WorkspaceAgentDisconnected,
169-
LifecycleState: codersdk.WorkspaceAgentLifecycleReady,
169+
var agent codersdk.WorkspaceAgent
170+
var logs []codersdk.WorkspaceAgentStartupLog
171+
172+
fetchSteps := []func(){
173+
func() {
174+
createdAt := time.Now().Add(-time.Minute)
175+
agent = codersdk.WorkspaceAgent{
176+
CreatedAt: createdAt,
177+
Status: codersdk.WorkspaceAgentConnecting,
178+
LifecycleState: codersdk.WorkspaceAgentLifecycleCreated,
179+
}
180+
},
181+
func() {
182+
time.Sleep(time.Second)
183+
agent.Status = codersdk.WorkspaceAgentTimeout
184+
},
185+
func() {
186+
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleStarting
187+
startingAt := time.Now()
188+
agent.StartedAt = &startingAt
189+
for i := 0; i < 10; i++ {
190+
level := codersdk.LogLevelInfo
191+
if rand.Float64() > 0.75 { //nolint:gosec
192+
level = codersdk.LogLevelError
193+
}
194+
logs = append(logs, codersdk.WorkspaceAgentStartupLog{
195+
CreatedAt: time.Now().Add(-time.Duration(10-i) * 144 * time.Millisecond),
196+
Output: fmt.Sprintf("Some log %d", i),
197+
Level: level,
198+
})
199+
}
200+
},
201+
func() {
202+
time.Sleep(time.Second)
203+
firstConnectedAt := time.Now()
204+
agent.FirstConnectedAt = &firstConnectedAt
205+
lastConnectedAt := firstConnectedAt.Add(0)
206+
agent.LastConnectedAt = &lastConnectedAt
207+
agent.Status = codersdk.WorkspaceAgentConnected
208+
},
209+
func() {},
210+
func() {
211+
time.Sleep(5 * time.Second)
212+
agent.Status = codersdk.WorkspaceAgentConnected
213+
lastConnectedAt := time.Now()
214+
agent.LastConnectedAt = &lastConnectedAt
215+
},
170216
}
171-
go func() {
172-
time.Sleep(3 * time.Second)
173-
agent.Status = codersdk.WorkspaceAgentConnected
174-
}()
175217
err := cliui.Agent(inv.Context(), inv.Stdout, cliui.AgentOptions{
176-
WorkspaceName: "dev",
177-
Fetch: func(ctx context.Context) (codersdk.WorkspaceAgent, error) {
218+
FetchInterval: 100 * time.Millisecond,
219+
Wait: true,
220+
Fetch: func(_ context.Context) (codersdk.WorkspaceAgent, error) {
221+
if len(fetchSteps) == 0 {
222+
return agent, nil
223+
}
224+
step := fetchSteps[0]
225+
fetchSteps = fetchSteps[1:]
226+
step()
178227
return agent, nil
179228
},
180-
WarnInterval: 2 * time.Second,
229+
FetchLogs: func(_ context.Context, _ uuid.UUID, _ int64, follow bool) (<-chan []codersdk.WorkspaceAgentStartupLog, io.Closer, error) {
230+
logsC := make(chan []codersdk.WorkspaceAgentStartupLog, len(logs))
231+
if follow {
232+
go func() {
233+
defer close(logsC)
234+
for _, log := range logs {
235+
logsC <- []codersdk.WorkspaceAgentStartupLog{log}
236+
time.Sleep(144 * time.Millisecond)
237+
}
238+
agent.LifecycleState = codersdk.WorkspaceAgentLifecycleReady
239+
readyAt := database.Now()
240+
agent.ReadyAt = &readyAt
241+
}()
242+
} else {
243+
logsC <- logs
244+
close(logsC)
245+
}
246+
return logsC, closeFunc(func() error {
247+
return nil
248+
}), nil
249+
},
181250
})
182251
if err != nil {
183252
return err
184253
}
185-
_, _ = fmt.Printf("Completed!\n")
186254
return nil
187255
},
188256
})
@@ -278,3 +346,9 @@ func main() {
278346
os.Exit(1)
279347
}
280348
}
349+
350+
type closeFunc func() error
351+
352+
func (f closeFunc) Close() error {
353+
return f()
354+
}

‎coderd/workspaceagents_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ func TestWorkspaceAgentStartupLogs(t *testing.T) {
225225
})
226226
require.NoError(t, err)
227227

228-
logs, closer, err := client.WorkspaceAgentStartupLogsAfter(ctx, build.Resources[0].Agents[0].ID, 0)
228+
logs, closer, err := client.WorkspaceAgentStartupLogsAfter(ctx, build.Resources[0].Agents[0].ID, 0, true)
229229
require.NoError(t, err)
230230
defer func() {
231231
_ = closer.Close()
@@ -338,7 +338,7 @@ func TestWorkspaceAgentStartupLogs(t *testing.T) {
338338
agentClient := agentsdk.New(client.URL)
339339
agentClient.SetSessionToken(authToken)
340340

341-
logs, closer, err := client.WorkspaceAgentStartupLogsAfter(ctx, build.Resources[0].Agents[0].ID, 0)
341+
logs, closer, err := client.WorkspaceAgentStartupLogsAfter(ctx, build.Resources[0].Agents[0].ID, 0, true)
342342
require.NoError(t, err)
343343
defer func() {
344344
_ = closer.Close()

‎codersdk/workspaceagents.go

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"net/http/cookiejar"
1212
"net/netip"
1313
"strconv"
14+
"strings"
1415
"time"
1516

1617
"github.com/google/uuid"
@@ -64,6 +65,17 @@ func (l WorkspaceAgentLifecycle) Starting() bool {
6465
}
6566
}
6667

68+
// ShuttingDown returns true if the agent is in the process of shutting
69+
// down or has shut down.
70+
func (l WorkspaceAgentLifecycle) ShuttingDown() bool {
71+
switch l {
72+
case WorkspaceAgentLifecycleShuttingDown, WorkspaceAgentLifecycleShutdownTimeout, WorkspaceAgentLifecycleShutdownError, WorkspaceAgentLifecycleOff:
73+
return true
74+
default:
75+
return false
76+
}
77+
}
78+
6779
// WorkspaceAgentLifecycleOrder is the order in which workspace agent
6880
// lifecycle states are expected to be reported during the lifetime of
6981
// the agent process. For instance, the agent can go from starting to
@@ -536,28 +548,59 @@ func (c *Client) WorkspaceAgentListeningPorts(ctx context.Context, agentID uuid.
536548
return listeningPorts, json.NewDecoder(res.Body).Decode(&listeningPorts)
537549
}
538550

539-
func (c *Client) WorkspaceAgentStartupLogsAfter(ctx context.Context, agentID uuid.UUID, after int64) (<-chan []WorkspaceAgentStartupLog, io.Closer, error) {
540-
afterQuery := ""
551+
func (c *Client) WorkspaceAgentStartupLogsAfter(ctx context.Context, agentID uuid.UUID, after int64, follow bool) (<-chan []WorkspaceAgentStartupLog, io.Closer, error) {
552+
var queryParams []string
541553
if after != 0 {
542-
afterQuery = fmt.Sprintf("&after=%d", after)
554+
queryParams = append(queryParams, fmt.Sprintf("after=%d", after))
555+
}
556+
if follow {
557+
queryParams = append(queryParams, "follow")
543558
}
544-
followURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/startup-logs?follow%s", agentID, afterQuery))
559+
var query string
560+
if len(queryParams) > 0 {
561+
query = "?" + strings.Join(queryParams, "&")
562+
}
563+
reqURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/startup-logs%s", agentID, query))
545564
if err != nil {
546565
return nil, nil, err
547566
}
567+
568+
if !follow {
569+
resp, err := c.Request(ctx, http.MethodGet, reqURL.String(), nil)
570+
if err != nil {
571+
return nil, nil, xerrors.Errorf("execute request: %w", err)
572+
}
573+
defer resp.Body.Close()
574+
575+
if resp.StatusCode != http.StatusOK {
576+
return nil, nil, ReadBodyAsError(resp)
577+
}
578+
579+
var logs []WorkspaceAgentStartupLog
580+
err = json.NewDecoder(resp.Body).Decode(&logs)
581+
if err != nil {
582+
return nil, nil, xerrors.Errorf("decode startup logs: %w", err)
583+
}
584+
585+
ch := make(chan []WorkspaceAgentStartupLog, 1)
586+
ch <- logs
587+
close(ch)
588+
return ch, closeFunc(func() error { return nil }), nil
589+
}
590+
548591
jar, err := cookiejar.New(nil)
549592
if err != nil {
550593
return nil, nil, xerrors.Errorf("create cookie jar: %w", err)
551594
}
552-
jar.SetCookies(followURL, []*http.Cookie{{
595+
jar.SetCookies(reqURL, []*http.Cookie{{
553596
Name: SessionTokenCookie,
554597
Value: c.SessionToken(),
555598
}})
556599
httpClient := &http.Client{
557600
Jar: jar,
558601
Transport: c.HTTPClient.Transport,
559602
}
560-
conn, res, err := websocket.Dial(ctx, followURL.String(), &websocket.DialOptions{
603+
conn, res, err := websocket.Dial(ctx, reqURL.String(), &websocket.DialOptions{
561604
HTTPClient: httpClient,
562605
CompressionMode: websocket.CompressionDisabled,
563606
})

0 commit comments

Comments
 (0)