fix: Refactor agent to consume API client (#4715)

* fix: Refactor agent to consume API client

This simplifies a lot of code by creating an interface for
the codersdk client into the agent. It also moves agent
authentication code so instance identity will work between
restarts.

Fixes #3485 and #4082.

* Fix client reconnections

agent/agent.go

@@ -34,6 +34,7 @@ import (
 
 	"cdr.dev/slog"
 	"github.com/coder/coder/agent/usershell"
+	"github.com/coder/coder/buildinfo"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/pty"
 	"github.com/coder/coder/tailnet"
@@ -52,22 +53,20 @@ const (
 )
 
 type Options struct {
-	CoordinatorDialer           CoordinatorDialer
-	FetchMetadata               FetchMetadata
-	StatsReporter               StatsReporter
-	WorkspaceAgentApps          WorkspaceAgentApps
-	PostWorkspaceAgentAppHealth PostWorkspaceAgentAppHealth
-	ReconnectingPTYTimeout      time.Duration
-	EnvironmentVariables        map[string]string
-	Logger                      slog.Logger
+	ExchangeToken          func(ctx context.Context) error
+	Client                 Client
+	ReconnectingPTYTimeout time.Duration
+	EnvironmentVariables   map[string]string
+	Logger                 slog.Logger
 }
 
-// CoordinatorDialer is a function that constructs a new broker.
-// A dialer must be passed in to allow for reconnects.
-type CoordinatorDialer func(context.Context) (net.Conn, error)
-
-// FetchMetadata is a function to obtain metadata for the agent.
-type FetchMetadata func(context.Context) (codersdk.WorkspaceAgentMetadata, error)
+type Client interface {
+	WorkspaceAgentMetadata(ctx context.Context) (codersdk.WorkspaceAgentMetadata, error)
+	ListenWorkspaceAgent(ctx context.Context) (net.Conn, error)
+	AgentReportStats(ctx context.Context, log slog.Logger, stats func() *codersdk.AgentStats) (io.Closer, error)
+	PostWorkspaceAgentAppHealth(ctx context.Context, req codersdk.PostWorkspaceAppHealthsRequest) error
+	PostWorkspaceAgentVersion(ctx context.Context, version string) error
+}
 
 func New(options Options) io.Closer {
 	if options.ReconnectingPTYTimeout == 0 {
@@ -75,24 +74,23 @@ func New(options Options) io.Closer {
 	}
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	server := &agent{
-		reconnectingPTYTimeout:      options.ReconnectingPTYTimeout,
-		logger:                      options.Logger,
-		closeCancel:                 cancelFunc,
-		closed:                      make(chan struct{}),
-		envVars:                     options.EnvironmentVariables,
-		coordinatorDialer:           options.CoordinatorDialer,
-		fetchMetadata:               options.FetchMetadata,
-		stats:                       &Stats{},
-		statsReporter:               options.StatsReporter,
-		workspaceAgentApps:          options.WorkspaceAgentApps,
-		postWorkspaceAgentAppHealth: options.PostWorkspaceAgentAppHealth,
+		reconnectingPTYTimeout: options.ReconnectingPTYTimeout,
+		logger:                 options.Logger,
+		closeCancel:            cancelFunc,
+		closed:                 make(chan struct{}),
+		envVars:                options.EnvironmentVariables,
+		client:                 options.Client,
+		exchangeToken:          options.ExchangeToken,
+		stats:                  &Stats{},
 	}
 	server.init(ctx)
 	return server
 }
 
 type agent struct {
-	logger slog.Logger
+	logger        slog.Logger
+	client        Client
+	exchangeToken func(ctx context.Context) error
 
 	reconnectingPTYs       sync.Map
 	reconnectingPTYTimeout time.Duration
@@ -104,100 +102,130 @@ type agent struct {
 
 	envVars map[string]string
 	// metadata is atomic because values can change after reconnection.
-	metadata      atomic.Value
-	fetchMetadata FetchMetadata
-	sshServer     *ssh.Server
+	metadata  atomic.Value
+	sshServer *ssh.Server
 
-	network                     *tailnet.Conn
-	coordinatorDialer           CoordinatorDialer
-	stats                       *Stats
-	statsReporter               StatsReporter
-	workspaceAgentApps          WorkspaceAgentApps
-	postWorkspaceAgentAppHealth PostWorkspaceAgentAppHealth
+	network *tailnet.Conn
+	stats   *Stats
 }
 
-func (a *agent) run(ctx context.Context) {
-	var metadata codersdk.WorkspaceAgentMetadata
-	var err error
-	// An exponential back-off occurs when the connection is failing to dial.
-	// This is to prevent server spam in case of a coderd outage.
-	for retrier := retry.New(50*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
-		a.logger.Info(ctx, "connecting")
-		metadata, err = a.fetchMetadata(ctx)
-		if err != nil {
-			if errors.Is(err, context.Canceled) {
-				return
-			}
-			if a.isClosed() {
-				return
-			}
-			a.logger.Warn(context.Background(), "failed to dial", slog.Error(err))
+// runLoop attempts to start the agent in a retry loop.
+// Coder may be offline temporarily, a connection issue
+// may be happening, but regardless after the intermittent
+// failure, you'll want the agent to reconnect.
+func (a *agent) runLoop(ctx context.Context) {
+	for retrier := retry.New(100*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
+		a.logger.Info(ctx, "running loop")
+		err := a.run(ctx)
+		// Cancel after the run is complete to clean up any leaked resources!
+		if err == nil {
 			continue
 		}
-		a.logger.Info(context.Background(), "fetched metadata")
-		break
-	}
-	select {
-	case <-ctx.Done():
-		return
-	default:
-	}
-	a.metadata.Store(metadata)
-
-	// The startup script has not ran yet!
-	go func() {
-		err := a.runStartupScript(ctx, metadata.StartupScript)
 		if errors.Is(err, context.Canceled) {
 			return
 		}
-		if err != nil {
-			a.logger.Warn(ctx, "agent script failed", slog.Error(err))
+		if a.isClosed() {
+			return
 		}
-	}()
-
-	if metadata.DERPMap != nil {
-		go a.runTailnet(ctx, metadata.DERPMap)
-	}
-
-	if a.workspaceAgentApps != nil && a.postWorkspaceAgentAppHealth != nil {
-		go NewWorkspaceAppHealthReporter(a.logger, a.workspaceAgentApps, a.postWorkspaceAgentAppHealth)(ctx)
+		if errors.Is(err, io.EOF) {
+			a.logger.Info(ctx, "likely disconnected from coder", slog.Error(err))
+			continue
+		}
+		a.logger.Warn(ctx, "run exited with error", slog.Error(err))
 	}
 }
 
-func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
+func (a *agent) run(ctx context.Context) error {
+	// This allows the agent to refresh it's token if necessary.
+	// For instance identity this is required, since the instance
+	// may not have re-provisioned, but a new agent ID was created.
+	if a.exchangeToken != nil {
+		err := a.exchangeToken(ctx)
+		if err != nil {
+			return xerrors.Errorf("exchange token: %w", err)
+		}
+	}
+
+	err := a.client.PostWorkspaceAgentVersion(ctx, buildinfo.Version())
+	if err != nil {
+		return xerrors.Errorf("update workspace agent version: %w", err)
+	}
+
+	metadata, err := a.client.WorkspaceAgentMetadata(ctx)
+	if err != nil {
+		return xerrors.Errorf("fetch metadata: %w", err)
+	}
+	a.logger.Info(context.Background(), "fetched metadata")
+	oldMetadata := a.metadata.Swap(metadata)
+
+	// The startup script should only execute on the first run!
+	if oldMetadata == nil {
+		go func() {
+			err := a.runStartupScript(ctx, metadata.StartupScript)
+			if errors.Is(err, context.Canceled) {
+				return
+			}
+			if err != nil {
+				a.logger.Warn(ctx, "agent script failed", slog.Error(err))
+			}
+		}()
+	}
+
+	// This automatically closes when the context ends!
+	appReporterCtx, appReporterCtxCancel := context.WithCancel(ctx)
+	defer appReporterCtxCancel()
+	go NewWorkspaceAppHealthReporter(
+		a.logger, metadata.Apps, a.client.PostWorkspaceAgentAppHealth)(appReporterCtx)
+
+	a.logger.Debug(ctx, "running tailnet with derpmap", slog.F("derpmap", metadata.DERPMap))
+
 	a.closeMutex.Lock()
-	defer a.closeMutex.Unlock()
-	if a.isClosed() {
-		return
+	network := a.network
+	a.closeMutex.Unlock()
+	if a.network == nil {
+		a.logger.Debug(ctx, "creating tailnet")
+		network, err = a.createTailnet(ctx, metadata.DERPMap)
+		if err != nil {
+			return xerrors.Errorf("create tailnet: %w", err)
+		}
+		a.closeMutex.Lock()
+		a.network = network
+		a.closeMutex.Unlock()
+	} else {
+		// Update the DERP map!
+		network.SetDERPMap(metadata.DERPMap)
 	}
-	a.logger.Debug(ctx, "running tailnet with derpmap", slog.F("derpmap", derpMap))
-	if a.network != nil {
-		a.network.SetDERPMap(derpMap)
-		return
+
+	a.logger.Debug(ctx, "running coordinator")
+	err = a.runCoordinator(ctx, network)
+	if err != nil {
+		a.logger.Debug(ctx, "coordinator exited", slog.Error(err))
+		return xerrors.Errorf("run coordinator: %w", err)
 	}
-	var err error
-	a.network, err = tailnet.NewConn(&tailnet.Options{
+	return nil
+}
+
+func (a *agent) createTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) (*tailnet.Conn, error) {
+	network, err := tailnet.NewConn(&tailnet.Options{
 		Addresses: []netip.Prefix{netip.PrefixFrom(codersdk.TailnetIP, 128)},
 		DERPMap:   derpMap,
 		Logger:    a.logger.Named("tailnet"),
 	})
 	if err != nil {
-		a.logger.Critical(ctx, "create tailnet", slog.Error(err))
-		return
+		return nil, xerrors.Errorf("create tailnet: %w", err)
 	}
-	a.network.SetForwardTCPCallback(func(conn net.Conn, listenerExists bool) net.Conn {
+	a.network = network
+	network.SetForwardTCPCallback(func(conn net.Conn, listenerExists bool) net.Conn {
 		if listenerExists {
 			// If a listener already exists, we would double-wrap the conn.
 			return conn
 		}
 		return a.stats.wrapConn(conn)
 	})
-	go a.runCoordinator(ctx)
 
-	sshListener, err := a.network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetSSHPort))
+	sshListener, err := network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetSSHPort))
 	if err != nil {
-		a.logger.Critical(ctx, "listen for ssh", slog.Error(err))
-		return
+		return nil, xerrors.Errorf("listen on the ssh port: %w", err)
 	}
 	go func() {
 		for {
@@ -209,10 +237,9 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
 		}
 	}()
 
-	reconnectingPTYListener, err := a.network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetReconnectingPTYPort))
+	reconnectingPTYListener, err := network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetReconnectingPTYPort))
 	if err != nil {
-		a.logger.Critical(ctx, "listen for reconnecting pty", slog.Error(err))
-		return
+		return nil, xerrors.Errorf("listen for reconnecting pty: %w", err)
 	}
 	go func() {
 		for {
@@ -244,10 +271,9 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
 		}
 	}()
 
-	speedtestListener, err := a.network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetSpeedtestPort))
+	speedtestListener, err := network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetSpeedtestPort))
 	if err != nil {
-		a.logger.Critical(ctx, "listen for speedtest", slog.Error(err))
-		return
+		return nil, xerrors.Errorf("listen for speedtest: %w", err)
 	}
 	go func() {
 		for {
@@ -266,10 +292,9 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
 		}
 	}()
 
-	statisticsListener, err := a.network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetStatisticsPort))
+	statisticsListener, err := network.Listen("tcp", ":"+strconv.Itoa(codersdk.TailnetStatisticsPort))
 	if err != nil {
-		a.logger.Critical(ctx, "listen for statistics", slog.Error(err))
-		return
+		return nil, xerrors.Errorf("listen for statistics: %w", err)
 	}
 	go func() {
 		defer statisticsListener.Close()
@@ -290,59 +315,26 @@ func (a *agent) runTailnet(ctx context.Context, derpMap *tailcfg.DERPMap) {
 			a.logger.Critical(ctx, "serve statistics HTTP server", slog.Error(err))
 		}
 	}()
+
+	return network, nil
 }
 
-// runCoordinator listens for nodes and updates the self-node as it changes.
-func (a *agent) runCoordinator(ctx context.Context) {
-	for {
-		reconnect := a.runCoordinatorWithRetry(ctx)
-		if !reconnect {
-			return
-		}
-	}
-}
-
-func (a *agent) runCoordinatorWithRetry(ctx context.Context) (reconnect bool) {
-	var coordinator net.Conn
-	var err error
-	// An exponential back-off occurs when the connection is failing to dial.
-	// This is to prevent server spam in case of a coderd outage.
-	for retrier := retry.New(50*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
-		coordinator, err = a.coordinatorDialer(ctx)
-		if err != nil {
-			if errors.Is(err, context.Canceled) {
-				return false
-			}
-			if a.isClosed() {
-				return false
-			}
-			a.logger.Warn(context.Background(), "failed to dial", slog.Error(err))
-			continue
-		}
-		//nolint:revive // Defer is ok because we're exiting this loop.
-		defer coordinator.Close()
-		a.logger.Info(context.Background(), "connected to coordination server")
-		break
+// runCoordinator runs a coordinator and returns whether a reconnect
+// should occur.
+func (a *agent) runCoordinator(ctx context.Context, network *tailnet.Conn) error {
+	coordinator, err := a.client.ListenWorkspaceAgent(ctx)
+	if err != nil {
+		return err
 	}
+	defer coordinator.Close()
+	a.logger.Info(context.Background(), "connected to coordination server")
+	sendNodes, errChan := tailnet.ServeCoordinator(coordinator, network.UpdateNodes)
+	network.SetNodeCallback(sendNodes)
 	select {
 	case <-ctx.Done():
-		return false
-	default:
-	}
-	sendNodes, errChan := tailnet.ServeCoordinator(coordinator, a.network.UpdateNodes)
-	a.network.SetNodeCallback(sendNodes)
-	select {
-	case <-ctx.Done():
-		return false
+		return ctx.Err()
 	case err := <-errChan:
-		if a.isClosed() {
-			return false
-		}
-		if errors.Is(err, context.Canceled) {
-			return false
-		}
-		a.logger.Debug(ctx, "node broker accept exited; restarting connection", slog.Error(err))
-		return true
+		return err
 	}
 }
 
@@ -474,22 +466,20 @@ func (a *agent) init(ctx context.Context) {
 		},
 	}
 
-	go a.run(ctx)
-	if a.statsReporter != nil {
-		cl, err := a.statsReporter(ctx, a.logger, func() *codersdk.AgentStats {
-			return a.stats.Copy()
-		})
-		if err != nil {
-			a.logger.Error(ctx, "report stats", slog.Error(err))
-			return
-		}
-		a.connCloseWait.Add(1)
-		go func() {
-			defer a.connCloseWait.Done()
-			<-a.closed
-			cl.Close()
-		}()
+	go a.runLoop(ctx)
+	cl, err := a.client.AgentReportStats(ctx, a.logger, func() *codersdk.AgentStats {
+		return a.stats.Copy()
+	})
+	if err != nil {
+		a.logger.Error(ctx, "report stats", slog.Error(err))
+		return
 	}
+	a.connCloseWait.Add(1)
+	go func() {
+		defer a.connCloseWait.Done()
+		<-a.closed
+		cl.Close()
+	}()
 }
 
 // createCommand processes raw command input with OpenSSH-like behavior.
@@ -696,7 +686,7 @@ func (a *agent) handleReconnectingPTY(ctx context.Context, msg codersdk.Reconnec
 
 		ptty, process, err := pty.Start(cmd)
 		if err != nil {
-			a.logger.Error(ctx, "start reconnecting pty command", slog.F("id", msg.ID))
+			a.logger.Error(ctx, "start reconnecting pty command", slog.F("id", msg.ID), slog.Error(err))
 			return
 		}
 

agent/agent_test.go

@@ -16,6 +16,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -503,6 +504,45 @@ func TestAgent(t *testing.T) {
 		require.NoError(t, err)
 		t.Logf("%.2f MBits/s", res[len(res)-1].MBitsPerSecond())
 	})
+
+	t.Run("Reconnect", func(t *testing.T) {
+		t.Parallel()
+		// After the agent is disconnected from a coordinator, it's supposed
+		// to reconnect!
+		coordinator := tailnet.NewCoordinator()
+		agentID := uuid.New()
+		statsCh := make(chan *codersdk.AgentStats)
+		derpMap := tailnettest.RunDERPAndSTUN(t)
+		client := &client{
+			t:       t,
+			agentID: agentID,
+			metadata: codersdk.WorkspaceAgentMetadata{
+				DERPMap: derpMap,
+			},
+			statsChan:   statsCh,
+			coordinator: coordinator,
+		}
+		initialized := atomic.Int32{}
+		closer := agent.New(agent.Options{
+			ExchangeToken: func(ctx context.Context) error {
+				initialized.Add(1)
+				return nil
+			},
+			Client: client,
+			Logger: slogtest.Make(t, nil).Leveled(slog.LevelInfo),
+		})
+		t.Cleanup(func() {
+			_ = closer.Close()
+		})
+
+		require.Eventually(t, func() bool {
+			return coordinator.Node(agentID) != nil
+		}, testutil.WaitShort, testutil.IntervalFast)
+		client.lastWorkspaceAgent()
+		require.Eventually(t, func() bool {
+			return initialized.Load() == 2
+		}, testutil.WaitShort, testutil.IntervalFast)
+	})
 }
 
 func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) *exec.Cmd {
@@ -572,57 +612,15 @@ func setupAgent(t *testing.T, metadata codersdk.WorkspaceAgentMetadata, ptyTimeo
 	agentID := uuid.New()
 	statsCh := make(chan *codersdk.AgentStats)
 	closer := agent.New(agent.Options{
-		FetchMetadata: func(ctx context.Context) (codersdk.WorkspaceAgentMetadata, error) {
-			return metadata, nil
-		},
-		CoordinatorDialer: func(ctx context.Context) (net.Conn, error) {
-			clientConn, serverConn := net.Pipe()
-			closed := make(chan struct{})
-			t.Cleanup(func() {
-				_ = serverConn.Close()
-				_ = clientConn.Close()
-				<-closed
-			})
-			go func() {
-				_ = coordinator.ServeAgent(serverConn, agentID)
-				close(closed)
-			}()
-			return clientConn, nil
+		Client: &client{
+			t:           t,
+			agentID:     agentID,
+			metadata:    metadata,
+			statsChan:   statsCh,
+			coordinator: coordinator,
 		},
 		Logger:                 slogtest.Make(t, nil).Leveled(slog.LevelDebug),
 		ReconnectingPTYTimeout: ptyTimeout,
-		StatsReporter: func(ctx context.Context, log slog.Logger, statsFn func() *codersdk.AgentStats) (io.Closer, error) {
-			doneCh := make(chan struct{})
-			ctx, cancel := context.WithCancel(ctx)
-
-			go func() {
-				defer close(doneCh)
-
-				t := time.NewTicker(time.Millisecond * 100)
-				defer t.Stop()
-				for {
-					select {
-					case <-ctx.Done():
-						return
-					case <-t.C:
-					}
-					select {
-					case statsCh <- statsFn():
-					case <-ctx.Done():
-						return
-					default:
-						// We don't want to send old stats.
-						continue
-					}
-				}
-			}()
-			return closeFunc(func() error {
-				cancel()
-				<-doneCh
-				close(statsCh)
-				return nil
-			}), nil
-		},
 	})
 	t.Cleanup(func() {
 		_ = closer.Close()
@@ -679,3 +677,73 @@ func assertWritePayload(t *testing.T, w io.Writer, payload []byte) {
 	assert.NoError(t, err, "write payload")
 	assert.Equal(t, len(payload), n, "payload length does not match")
 }
+
+type client struct {
+	t                  *testing.T
+	agentID            uuid.UUID
+	metadata           codersdk.WorkspaceAgentMetadata
+	statsChan          chan *codersdk.AgentStats
+	coordinator        tailnet.Coordinator
+	lastWorkspaceAgent func()
+}
+
+func (c *client) WorkspaceAgentMetadata(_ context.Context) (codersdk.WorkspaceAgentMetadata, error) {
+	return c.metadata, nil
+}
+
+func (c *client) ListenWorkspaceAgent(_ context.Context) (net.Conn, error) {
+	clientConn, serverConn := net.Pipe()
+	closed := make(chan struct{})
+	c.lastWorkspaceAgent = func() {
+		_ = serverConn.Close()
+		_ = clientConn.Close()
+		<-closed
+	}
+	c.t.Cleanup(c.lastWorkspaceAgent)
+	go func() {
+		_ = c.coordinator.ServeAgent(serverConn, c.agentID)
+		close(closed)
+	}()
+	return clientConn, nil
+}
+
+func (c *client) AgentReportStats(ctx context.Context, _ slog.Logger, stats func() *codersdk.AgentStats) (io.Closer, error) {
+	doneCh := make(chan struct{})
+	ctx, cancel := context.WithCancel(ctx)
+
+	go func() {
+		defer close(doneCh)
+
+		t := time.NewTicker(time.Millisecond * 100)
+		defer t.Stop()
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-t.C:
+			}
+			select {
+			case c.statsChan <- stats():
+			case <-ctx.Done():
+				return
+			default:
+				// We don't want to send old stats.
+				continue
+			}
+		}
+	}()
+	return closeFunc(func() error {
+		cancel()
+		<-doneCh
+		close(c.statsChan)
+		return nil
+	}), nil
+}
+
+func (*client) PostWorkspaceAgentAppHealth(_ context.Context, _ codersdk.PostWorkspaceAppHealthsRequest) error {
+	return nil
+}
+
+func (*client) PostWorkspaceAgentVersion(_ context.Context, _ string) error {
+	return nil
+}

agent/apphealth.go

@@ -23,16 +23,8 @@ type PostWorkspaceAgentAppHealth func(context.Context, codersdk.PostWorkspaceApp
 type WorkspaceAppHealthReporter func(ctx context.Context)
 
 // NewWorkspaceAppHealthReporter creates a WorkspaceAppHealthReporter that reports app health to coderd.
-func NewWorkspaceAppHealthReporter(logger slog.Logger, workspaceAgentApps WorkspaceAgentApps, postWorkspaceAgentAppHealth PostWorkspaceAgentAppHealth) WorkspaceAppHealthReporter {
+func NewWorkspaceAppHealthReporter(logger slog.Logger, apps []codersdk.WorkspaceApp, postWorkspaceAgentAppHealth PostWorkspaceAgentAppHealth) WorkspaceAppHealthReporter {
 	runHealthcheckLoop := func(ctx context.Context) error {
-		apps, err := workspaceAgentApps(ctx)
-		if err != nil {
-			if xerrors.Is(err, context.Canceled) {
-				return nil
-			}
-			return xerrors.Errorf("getting workspace apps: %w", err)
-		}
-
 		// no need to run this loop if no apps for this workspace.
 		if len(apps) == 0 {
 			return nil

agent/apphealth_test.go

@@ -199,7 +199,7 @@ func setupAppReporter(ctx context.Context, t *testing.T, apps []codersdk.Workspa
 		return nil
 	}
 
-	go agent.NewWorkspaceAppHealthReporter(slogtest.Make(t, nil).Leveled(slog.LevelDebug), workspaceAgentApps, postWorkspaceAgentAppHealth)(ctx)
+	go agent.NewWorkspaceAppHealthReporter(slogtest.Make(t, nil).Leveled(slog.LevelDebug), apps, postWorkspaceAgentAppHealth)(ctx)
 
 	return workspaceAgentApps, func() {
 		for _, closeFn := range closers {

agent/stats.go

@@ -1,12 +1,9 @@
 package agent
 
 import (
-	"context"
-	"io"
 	"net"
 	"sync/atomic"
 
-	"cdr.dev/slog"
 	"github.com/coder/coder/codersdk"
 )
 
@@ -59,10 +56,3 @@ func (s *Stats) wrapConn(conn net.Conn) net.Conn {
 
 	return cs
 }
-
-// StatsReporter periodically accept and records agent stats.
-type StatsReporter func(
-	ctx context.Context,
-	log slog.Logger,
-	stats func() *codersdk.AgentStats,
-) (io.Closer, error)