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

cli/agent.go

@@ -23,7 +23,6 @@ import (
 	"github.com/coder/coder/buildinfo"
 	"github.com/coder/coder/cli/cliflag"
 	"github.com/coder/coder/codersdk"
-	"github.com/coder/retry"
 )
 
 func workspaceAgent() *cobra.Command {
@@ -80,6 +79,8 @@ func workspaceAgent() *cobra.Command {
 				slog.F("version", version),
 			)
 			client := codersdk.New(coderURL)
+			// Set a reasonable timeout so requests can't hang forever!
+			client.HTTPClient.Timeout = 10 * time.Second
 
 			if pprofEnabled {
 				srvClose := serveHandler(cmd.Context(), logger, nil, pprofAddress, "pprof")
@@ -143,43 +144,6 @@ func workspaceAgent() *cobra.Command {
 				}
 			}
 
-			if exchangeToken != nil {
-				logger.Info(cmd.Context(), "exchanging identity token")
-				// Agent's can start before resources are returned from the provisioner
-				// daemon. If there are many resources being provisioned, this time
-				// could be significant. This is arbitrarily set at an hour to prevent
-				// tons of idle agents from pinging coderd.
-				ctx, cancelFunc := context.WithTimeout(cmd.Context(), time.Hour)
-				defer cancelFunc()
-				for retry.New(100*time.Millisecond, 5*time.Second).Wait(ctx) {
-					var response codersdk.WorkspaceAgentAuthenticateResponse
-
-					response, err = exchangeToken(ctx)
-					if err != nil {
-						logger.Warn(ctx, "authenticate workspace", slog.F("method", auth), slog.Error(err))
-						continue
-					}
-					client.SessionToken = response.SessionToken
-					logger.Info(ctx, "authenticated", slog.F("method", auth))
-					break
-				}
-				if err != nil {
-					return xerrors.Errorf("agent failed to authenticate in time: %w", err)
-				}
-			}
-
-			retryCtx, cancelRetry := context.WithTimeout(cmd.Context(), time.Hour)
-			defer cancelRetry()
-			for retrier := retry.New(100*time.Millisecond, 5*time.Second); retrier.Wait(retryCtx); {
-				err := client.PostWorkspaceAgentVersion(retryCtx, version)
-				if err != nil {
-					logger.Warn(retryCtx, "post agent version: %w", slog.Error(err), slog.F("version", version))
-					continue
-				}
-				logger.Info(retryCtx, "updated agent version", slog.F("version", version))
-				break
-			}
-
 			executablePath, err := os.Executable()
 			if err != nil {
 				return xerrors.Errorf("getting os executable: %w", err)
@@ -190,17 +154,24 @@ func workspaceAgent() *cobra.Command {
 			}
 
 			closer := agent.New(agent.Options{
-				FetchMetadata: client.WorkspaceAgentMetadata,
-				Logger:        logger,
+				Client: client,
+				Logger: logger,
+				ExchangeToken: func(ctx context.Context) error {
+					if exchangeToken == nil {
+						return nil
+					}
+					resp, err := exchangeToken(ctx)
+					if err != nil {
+						return err
+					}
+					client.SessionToken = resp.SessionToken
+					return nil
+				},
 				EnvironmentVariables: map[string]string{
 					// Override the "CODER_AGENT_TOKEN" variable in all
 					// shells so "gitssh" works!
 					"CODER_AGENT_TOKEN": client.SessionToken,
 				},
-				CoordinatorDialer:           client.ListenWorkspaceAgentTailnet,
-				StatsReporter:               client.AgentReportStats,
-				WorkspaceAgentApps:          client.WorkspaceAgentApps,
-				PostWorkspaceAgentAppHealth: client.PostWorkspaceAgentAppHealth,
 			})
 			<-cmd.Context().Done()
 			return closer.Close()

cli/configssh_test.go

@@ -106,9 +106,8 @@ func TestConfigSSH(t *testing.T) {
 	agentClient := codersdk.New(client.URL)
 	agentClient.SessionToken = authToken
 	agentCloser := agent.New(agent.Options{
-		FetchMetadata:     agentClient.WorkspaceAgentMetadata,
-		CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
-		Logger:            slogtest.Make(t, nil).Named("agent"),
+		Client: agentClient,
+		Logger: slogtest.Make(t, nil).Named("agent"),
 	})
 	defer func() {
 		_ = agentCloser.Close()

cli/speedtest_test.go

@@ -24,9 +24,8 @@ func TestSpeedtest(t *testing.T) {
 	agentClient := codersdk.New(client.URL)
 	agentClient.SessionToken = agentToken
 	agentCloser := agent.New(agent.Options{
-		FetchMetadata:     agentClient.WorkspaceAgentMetadata,
-		CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
-		Logger:            slogtest.Make(t, nil).Named("agent"),
+		Client: agentClient,
+		Logger: slogtest.Make(t, nil).Named("agent"),
 	})
 	defer agentCloser.Close()
 	coderdtest.AwaitWorkspaceAgents(t, client, workspace.ID)

cli/ssh_test.go

@@ -89,9 +89,8 @@ func TestSSH(t *testing.T) {
 		agentClient := codersdk.New(client.URL)
 		agentClient.SessionToken = agentToken
 		agentCloser := agent.New(agent.Options{
-			FetchMetadata:     agentClient.WorkspaceAgentMetadata,
-			CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
-			Logger:            slogtest.Make(t, nil).Named("agent"),
+			Client: agentClient,
+			Logger: slogtest.Make(t, nil).Named("agent"),
 		})
 		defer func() {
 			_ = agentCloser.Close()
@@ -110,9 +109,8 @@ func TestSSH(t *testing.T) {
 			agentClient := codersdk.New(client.URL)
 			agentClient.SessionToken = agentToken
 			agentCloser := agent.New(agent.Options{
-				FetchMetadata:     agentClient.WorkspaceAgentMetadata,
-				CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
-				Logger:            slogtest.Make(t, nil).Named("agent"),
+				Client: agentClient,
+				Logger: slogtest.Make(t, nil).Named("agent"),
 			})
 			<-ctx.Done()
 			_ = agentCloser.Close()
@@ -178,9 +176,8 @@ func TestSSH(t *testing.T) {
 		agentClient := codersdk.New(client.URL)
 		agentClient.SessionToken = agentToken
 		agentCloser := agent.New(agent.Options{
-			FetchMetadata:     agentClient.WorkspaceAgentMetadata,
-			CoordinatorDialer: agentClient.ListenWorkspaceAgentTailnet,
-			Logger:            slogtest.Make(t, nil).Named("agent"),
+			Client: agentClient,
+			Logger: slogtest.Make(t, nil).Named("agent"),
 		})
 		defer agentCloser.Close()