feat: add agent log streaming and follow provisioner format (#8170)

codersdk/workspaceagents.go

@@ -11,6 +11,7 @@ import (
 	"net/http/cookiejar"
 	"net/netip"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/google/uuid"
@@ -64,6 +65,17 @@ func (l WorkspaceAgentLifecycle) Starting() bool {
 	}
 }
 
+// ShuttingDown returns true if the agent is in the process of shutting
+// down or has shut down.
+func (l WorkspaceAgentLifecycle) ShuttingDown() bool {
+	switch l {
+	case WorkspaceAgentLifecycleShuttingDown, WorkspaceAgentLifecycleShutdownTimeout, WorkspaceAgentLifecycleShutdownError, WorkspaceAgentLifecycleOff:
+		return true
+	default:
+		return false
+	}
+}
+
 // WorkspaceAgentLifecycleOrder is the order in which workspace agent
 // lifecycle states are expected to be reported during the lifetime of
 // the agent process. For instance, the agent can go from starting to
@@ -536,20 +548,52 @@ func (c *Client) WorkspaceAgentListeningPorts(ctx context.Context, agentID uuid.
 	return listeningPorts, json.NewDecoder(res.Body).Decode(&listeningPorts)
 }
 
-func (c *Client) WorkspaceAgentStartupLogsAfter(ctx context.Context, agentID uuid.UUID, after int64) (<-chan []WorkspaceAgentStartupLog, io.Closer, error) {
-	afterQuery := ""
+//nolint:revive // Follow is a control flag on the server as well.
+func (c *Client) WorkspaceAgentStartupLogsAfter(ctx context.Context, agentID uuid.UUID, after int64, follow bool) (<-chan []WorkspaceAgentStartupLog, io.Closer, error) {
+	var queryParams []string
 	if after != 0 {
-		afterQuery = fmt.Sprintf("&after=%d", after)
+		queryParams = append(queryParams, fmt.Sprintf("after=%d", after))
 	}
-	followURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/startup-logs?follow%s", agentID, afterQuery))
+	if follow {
+		queryParams = append(queryParams, "follow")
+	}
+	var query string
+	if len(queryParams) > 0 {
+		query = "?" + strings.Join(queryParams, "&")
+	}
+	reqURL, err := c.URL.Parse(fmt.Sprintf("/api/v2/workspaceagents/%s/startup-logs%s", agentID, query))
 	if err != nil {
 		return nil, nil, err
 	}
+
+	if !follow {
+		resp, err := c.Request(ctx, http.MethodGet, reqURL.String(), nil)
+		if err != nil {
+			return nil, nil, xerrors.Errorf("execute request: %w", err)
+		}
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusOK {
+			return nil, nil, ReadBodyAsError(resp)
+		}
+
+		var logs []WorkspaceAgentStartupLog
+		err = json.NewDecoder(resp.Body).Decode(&logs)
+		if err != nil {
+			return nil, nil, xerrors.Errorf("decode startup logs: %w", err)
+		}
+
+		ch := make(chan []WorkspaceAgentStartupLog, 1)
+		ch <- logs
+		close(ch)
+		return ch, closeFunc(func() error { return nil }), nil
+	}
+
 	jar, err := cookiejar.New(nil)
 	if err != nil {
 		return nil, nil, xerrors.Errorf("create cookie jar: %w", err)
 	}
-	jar.SetCookies(followURL, []*http.Cookie{{
+	jar.SetCookies(reqURL, []*http.Cookie{{
 		Name:  SessionTokenCookie,
 		Value: c.SessionToken(),
 	}})
@@ -557,7 +601,7 @@ func (c *Client) WorkspaceAgentStartupLogsAfter(ctx context.Context, agentID uui
 		Jar:       jar,
 		Transport: c.HTTPClient.Transport,
 	}
-	conn, res, err := websocket.Dial(ctx, followURL.String(), &websocket.DialOptions{
+	conn, res, err := websocket.Dial(ctx, reqURL.String(), &websocket.DialOptions{
 		HTTPClient:      httpClient,
 		CompressionMode: websocket.CompressionDisabled,
 	})