fix!: remove startup logs eof for streaming (#8528)

* fix: remove startup logs eof for streaming

We have external utilities like logstream-kube that may send
logs after an agent shuts down unexpectedly to report additional
information. In a recent change we stopped accepting these logs,
which broke these utilities.

In the future we'll rename startup logs to agent logs or something
more generalized so this is less confusing in the future.

* fix(cli/cliui): handle never ending startup log stream in Agent

---------

Co-authored-by: Mathias Fredriksson <mafredri@gmail.com>

cli/cliui/agent.go

@@ -137,26 +137,44 @@ func Agent(ctx context.Context, writer io.Writer, agentID uuid.UUID, opts AgentO
 				}
 				defer logsCloser.Close()
 
+				var lastLog codersdk.WorkspaceAgentStartupLog
+				fetchedAgentWhileFollowing := fetchedAgent
+				if !follow {
+					fetchedAgentWhileFollowing = nil
+				}
 				for {
 					// This select is essentially and inline `fetch()`.
 					select {
 					case <-ctx.Done():
 						return ctx.Err()
-					case f := <-fetchedAgent:
+					case f := <-fetchedAgentWhileFollowing:
 						if f.err != nil {
 							return xerrors.Errorf("fetch: %w", f.err)
 						}
-						// We could handle changes in the agent status here, like
-						// if the agent becomes disconnected, we may want to stop.
-						// But for now, we'll just keep going, hopefully the agent
-						// will reconnect and update its status.
 						agent = f.agent
+
+						// If the agent is no longer starting, stop following
+						// logs because FetchLogs will keep streaming forever.
+						// We do one last non-follow request to ensure we have
+						// fetched all logs.
+						if !agent.LifecycleState.Starting() {
+							_ = logsCloser.Close()
+							fetchedAgentWhileFollowing = nil
+
+							logStream, logsCloser, err = opts.FetchLogs(ctx, agent.ID, lastLog.ID, false)
+							if err != nil {
+								return xerrors.Errorf("fetch workspace agent startup logs: %w", err)
+							}
+							// Logs are already primed, so we can call close.
+							_ = logsCloser.Close()
+						}
 					case logs, ok := <-logStream:
 						if !ok {
 							return nil
 						}
 						for _, log := range logs {
 							sw.Log(log.CreatedAt, log.Level, log.Output)
+							lastLog = log
 						}
 					}
 				}

cli/cliui/agent_test.go

@@ -46,7 +46,6 @@ func TestAgent(t *testing.T) {
 				func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentStartupLog) error {
 					agent.Status = codersdk.WorkspaceAgentConnected
 					agent.FirstConnectedAt = ptr.Ref(time.Now())
-					close(logs)
 					return nil
 				},
 			},
@@ -79,7 +78,6 @@ func TestAgent(t *testing.T) {
 					agent.FirstConnectedAt = ptr.Ref(time.Now())
 					agent.LifecycleState = codersdk.WorkspaceAgentLifecycleReady
 					agent.ReadyAt = ptr.Ref(time.Now())
-					close(logs)
 					return nil
 				},
 			},
@@ -113,10 +111,6 @@ func TestAgent(t *testing.T) {
 					agent.LastConnectedAt = ptr.Ref(time.Now())
 					return nil
 				},
-				func(_ context.Context, _ *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentStartupLog) error {
-					close(logs)
-					return nil
-				},
 			},
 			want: []string{
 				"⧗ The workspace agent lost connection",
@@ -154,7 +148,6 @@ func TestAgent(t *testing.T) {
 							Output:    "Bye now",
 						},
 					}
-					close(logs)
 					return nil
 				},
 			},
@@ -184,7 +177,6 @@ func TestAgent(t *testing.T) {
 							Output:    "Hello world",
 						},
 					}
-					close(logs)
 					return nil
 				},
 			},
@@ -205,7 +197,6 @@ func TestAgent(t *testing.T) {
 				func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentStartupLog) error {
 					agent.Status = codersdk.WorkspaceAgentDisconnected
 					agent.LifecycleState = codersdk.WorkspaceAgentLifecycleOff
-					close(logs)
 					return nil
 				},
 			},
@@ -234,7 +225,6 @@ func TestAgent(t *testing.T) {
 				func(_ context.Context, agent *codersdk.WorkspaceAgent, logs chan []codersdk.WorkspaceAgentStartupLog) error {
 					agent.ReadyAt = ptr.Ref(time.Now())
 					agent.LifecycleState = codersdk.WorkspaceAgentLifecycleShuttingDown
-					close(logs)
 					return nil
 				},
 			},
@@ -316,8 +306,21 @@ func TestAgent(t *testing.T) {
 						}
 						return agent, err
 					}
-					tc.opts.FetchLogs = func(_ context.Context, _ uuid.UUID, _ int64, _ bool) (<-chan []codersdk.WorkspaceAgentStartupLog, io.Closer, error) {
-						return logs, closeFunc(func() error { return nil }), nil
+					tc.opts.FetchLogs = func(ctx context.Context, _ uuid.UUID, _ int64, follow bool) (<-chan []codersdk.WorkspaceAgentStartupLog, io.Closer, error) {
+						if follow {
+							return logs, closeFunc(func() error { return nil }), nil
+						}
+
+						fetchLogs := make(chan []codersdk.WorkspaceAgentStartupLog, 1)
+						select {
+						case <-ctx.Done():
+							return nil, nil, ctx.Err()
+						case l := <-logs:
+							fetchLogs <- l
+						default:
+						}
+						close(fetchLogs)
+						return fetchLogs, closeFunc(func() error { return nil }), nil
 					}
 					err := cliui.Agent(inv.Context(), &buf, uuid.Nil, tc.opts)
 					return err