chore: refactor instance identity to be a SessionTokenProvider (#19566)

Refactors Agent instance identity to be a SessionTokenProvider.

Refactors the CLI to create Agent clients via a centralized function, rather than add-hoc via individual command handlers and their flags.

This allows commands besides `coder agent`, but which still use the agent identity, to support instance identity authentication.

Fixes #19111 by unifying all API requests to go thru the SessionTokenProvider for auth credentials.

cli/agent.go

@@ -15,7 +15,6 @@ import (
 	"strings"
 	"time"
 
-	"cloud.google.com/go/compute/metadata"
 	"golang.org/x/xerrors"
 	"gopkg.in/natefinch/lumberjack.v2"
 
@@ -38,9 +37,8 @@ import (
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 )
 
-func (r *RootCmd) workspaceAgent() *serpent.Command {
+func workspaceAgent() *serpent.Command {
 	var (
-		auth                           string
 		logDir                         string
 		scriptDataDir                  string
 		pprofAddress                   string
@@ -59,6 +57,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 		devcontainerProjectDiscovery   bool
 		devcontainerDiscoveryAutostart bool
 	)
+	agentAuth := &AgentAuth{}
 	cmd := &serpent.Command{
 		Use:   "agent",
 		Short: `Starts the Coder workspace agent.`,
@@ -176,12 +175,14 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 
 			version := buildinfo.Version()
 			logger.Info(ctx, "agent is starting now",
-				slog.F("url", r.agentURL),
-				slog.F("auth", auth),
+				slog.F("url", agentAuth.agentURL),
+				slog.F("auth", agentAuth.agentAuth),
 				slog.F("version", version),
 			)
-
-			client := agentsdk.New(r.agentURL)
+			client, err := agentAuth.CreateClient(ctx)
+			if err != nil {
+				return xerrors.Errorf("create agent client: %w", err)
+			}
 			client.SDK.SetLogger(logger)
 			// Set a reasonable timeout so requests can't hang forever!
 			// The timeout needs to be reasonably long, because requests
@@ -190,7 +191,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 			client.SDK.HTTPClient.Timeout = 30 * time.Second
 			// Attach header transport so we process --agent-header and
 			// --agent-header-command flags
-			headerTransport, err := headerTransport(ctx, r.agentURL, agentHeader, agentHeaderCommand)
+			headerTransport, err := headerTransport(ctx, &agentAuth.agentURL, agentHeader, agentHeaderCommand)
 			if err != nil {
 				return xerrors.Errorf("configure header transport: %w", err)
 			}
@@ -214,68 +215,6 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 				ignorePorts[port] = "debug"
 			}
 
-			// exchangeToken returns a session token.
-			// This is abstracted to allow for the same looping condition
-			// regardless of instance identity auth type.
-			var exchangeToken func(context.Context) (agentsdk.AuthenticateResponse, error)
-			switch auth {
-			case "token":
-				token, _ := inv.ParsedFlags().GetString(varAgentToken)
-				if token == "" {
-					tokenFile, _ := inv.ParsedFlags().GetString(varAgentTokenFile)
-					if tokenFile != "" {
-						tokenBytes, err := os.ReadFile(tokenFile)
-						if err != nil {
-							return xerrors.Errorf("read token file %q: %w", tokenFile, err)
-						}
-						token = strings.TrimSpace(string(tokenBytes))
-					}
-				}
-				if token == "" {
-					return xerrors.Errorf("CODER_AGENT_TOKEN or CODER_AGENT_TOKEN_FILE must be set for token auth")
-				}
-				client.SetSessionToken(token)
-			case "google-instance-identity":
-				// This is *only* done for testing to mock client authentication.
-				// This will never be set in a production scenario.
-				var gcpClient *metadata.Client
-				gcpClientRaw := ctx.Value("gcp-client")
-				if gcpClientRaw != nil {
-					gcpClient, _ = gcpClientRaw.(*metadata.Client)
-				}
-				exchangeToken = func(ctx context.Context) (agentsdk.AuthenticateResponse, error) {
-					return client.AuthGoogleInstanceIdentity(ctx, "", gcpClient)
-				}
-			case "aws-instance-identity":
-				// This is *only* done for testing to mock client authentication.
-				// This will never be set in a production scenario.
-				var awsClient *http.Client
-				awsClientRaw := ctx.Value("aws-client")
-				if awsClientRaw != nil {
-					awsClient, _ = awsClientRaw.(*http.Client)
-					if awsClient != nil {
-						client.SDK.HTTPClient = awsClient
-					}
-				}
-				exchangeToken = func(ctx context.Context) (agentsdk.AuthenticateResponse, error) {
-					return client.AuthAWSInstanceIdentity(ctx)
-				}
-			case "azure-instance-identity":
-				// This is *only* done for testing to mock client authentication.
-				// This will never be set in a production scenario.
-				var azureClient *http.Client
-				azureClientRaw := ctx.Value("azure-client")
-				if azureClientRaw != nil {
-					azureClient, _ = azureClientRaw.(*http.Client)
-					if azureClient != nil {
-						client.SDK.HTTPClient = azureClient
-					}
-				}
-				exchangeToken = func(ctx context.Context) (agentsdk.AuthenticateResponse, error) {
-					return client.AuthAzureInstanceIdentity(ctx)
-				}
-			}
-
 			executablePath, err := os.Executable()
 			if err != nil {
 				return xerrors.Errorf("getting os executable: %w", err)
@@ -343,18 +282,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 					LogDir:        logDir,
 					ScriptDataDir: scriptDataDir,
 					// #nosec G115 - Safe conversion as tailnet listen port is within uint16 range (0-65535)
-					TailnetListenPort: uint16(tailnetListenPort),
-					ExchangeToken: func(ctx context.Context) (string, error) {
-						if exchangeToken == nil {
-							return client.SDK.SessionToken(), nil
-						}
-						resp, err := exchangeToken(ctx)
-						if err != nil {
-							return "", err
-						}
-						client.SetSessionToken(resp.SessionToken)
-						return resp.SessionToken, nil
-					},
+					TailnetListenPort:    uint16(tailnetListenPort),
 					EnvironmentVariables: environmentVariables,
 					IgnorePorts:          ignorePorts,
 					SSHMaxTimeout:        sshMaxTimeout,
@@ -365,7 +293,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 					Execer:             execer,
 					Devcontainers:      devcontainers,
 					DevcontainerAPIOptions: []agentcontainers.Option{
-						agentcontainers.WithSubAgentURL(r.agentURL.String()),
+						agentcontainers.WithSubAgentURL(agentAuth.agentURL.String()),
 						agentcontainers.WithProjectDiscovery(devcontainerProjectDiscovery),
 						agentcontainers.WithDiscoveryAutostart(devcontainerDiscoveryAutostart),
 					},
@@ -400,13 +328,6 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 	}
 
 	cmd.Options = serpent.OptionSet{
-		{
-			Flag:        "auth",
-			Default:     "token",
-			Description: "Specify the authentication type to use for the agent.",
-			Env:         "CODER_AGENT_AUTH",
-			Value:       serpent.StringOf(&auth),
-		},
 		{
 			Flag:        "log-dir",
 			Default:     os.TempDir(),
@@ -529,7 +450,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 			Value:       serpent.BoolOf(&devcontainerDiscoveryAutostart),
 		},
 	}
-
+	agentAuth.AttachOptions(cmd, false)
 	return cmd
 }
 

cli/exp_mcp.go

@@ -56,7 +56,7 @@ func (r *RootCmd) mcpConfigure() *serpent.Command {
 		},
 		Children: []*serpent.Command{
 			r.mcpConfigureClaudeDesktop(),
-			r.mcpConfigureClaudeCode(),
+			mcpConfigureClaudeCode(),
 			r.mcpConfigureCursor(),
 		},
 	}
@@ -117,7 +117,7 @@ func (*RootCmd) mcpConfigureClaudeDesktop() *serpent.Command {
 	return cmd
 }
 
-func (r *RootCmd) mcpConfigureClaudeCode() *serpent.Command {
+func mcpConfigureClaudeCode() *serpent.Command {
 	var (
 		claudeAPIKey     string
 		claudeConfigPath string
@@ -131,6 +131,7 @@ func (r *RootCmd) mcpConfigureClaudeCode() *serpent.Command {
 
 		deprecatedCoderMCPClaudeAPIKey string
 	)
+	agentAuth := &AgentAuth{}
 	cmd := &serpent.Command{
 		Use:   "claude-code <project-directory>",
 		Short: "Configure the Claude Code server. You will need to run this command for each project you want to use. Specify the project directory as the first argument.",
@@ -148,7 +149,7 @@ func (r *RootCmd) mcpConfigureClaudeCode() *serpent.Command {
 				binPath = testBinaryName
 			}
 			configureClaudeEnv := map[string]string{}
-			agentClient, err := r.createAgentClient()
+			agentClient, err := agentAuth.CreateClient(inv.Context())
 			if err != nil {
 				cliui.Warnf(inv.Stderr, "failed to create agent client: %s", err)
 			} else {
@@ -292,6 +293,7 @@ func (r *RootCmd) mcpConfigureClaudeCode() *serpent.Command {
 			},
 		},
 	}
+	agentAuth.AttachOptions(cmd, false)
 	return cmd
 }
 
@@ -403,7 +405,8 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 		appStatusSlug string
 		aiAgentAPIURL url.URL
 	)
-	return &serpent.Command{
+	agentAuth := &AgentAuth{}
+	cmd := &serpent.Command{
 		Use: "server",
 		Handler: func(inv *serpent.Invocation) error {
 			var lastReport taskReport
@@ -494,7 +497,7 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 			}
 
 			// Try to create an agent client for status reporting.  Not validated.
-			agentClient, err := r.createAgentClient()
+			agentClient, err := agentAuth.CreateClient(inv.Context())
 			if err == nil {
 				cliui.Infof(inv.Stderr, "Agent URL      : %s", agentClient.SDK.URL.String())
 				srv.agentClient = agentClient
@@ -579,6 +582,8 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 			},
 		},
 	}
+	agentAuth.AttachOptions(cmd, false)
+	return cmd
 }
 
 func (s *mcpServer) startReporter(ctx context.Context, inv *serpent.Invocation) {

cli/externalauth.go

@@ -2,19 +2,16 @@ package cli
 
 import (
 	"encoding/json"
-	"fmt"
-
-	"golang.org/x/xerrors"
 
 	"github.com/tidwall/gjson"
+	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/cli/cliui"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
-	"github.com/coder/pretty"
 	"github.com/coder/serpent"
 )
 
-func (r *RootCmd) externalAuth() *serpent.Command {
+func externalAuth() *serpent.Command {
 	return &serpent.Command{
 		Use:   "external-auth",
 		Short: "Manage external authentication",
@@ -23,14 +20,15 @@ func (r *RootCmd) externalAuth() *serpent.Command {
 			return i.Command.HelpHandler(i)
 		},
 		Children: []*serpent.Command{
-			r.externalAuthAccessToken(),
+			externalAuthAccessToken(),
 		},
 	}
 }
 
-func (r *RootCmd) externalAuthAccessToken() *serpent.Command {
+func externalAuthAccessToken() *serpent.Command {
 	var extra string
-	return &serpent.Command{
+	agentAuth := &AgentAuth{}
+	cmd := &serpent.Command{
 		Use:   "access-token <provider>",
 		Short: "Print auth for an external provider",
 		Long: "Print an access-token for an external auth provider. " +
@@ -70,12 +68,7 @@ fi
 			ctx, stop := inv.SignalNotifyContext(ctx, StopSignals...)
 			defer stop()
 
-			if r.agentToken == "" {
-				_, _ = fmt.Fprint(inv.Stderr, pretty.Sprintf(headLineStyle(), "No agent token found, this command must be run from inside a running workspace.\n"))
-				return xerrors.Errorf("agent token not found")
-			}
-
-			client, err := r.tryCreateAgentClient()
+			client, err := agentAuth.CreateClient(ctx)
 			if err != nil {
 				return xerrors.Errorf("create agent client: %w", err)
 			}
@@ -115,4 +108,6 @@ fi
 			return nil
 		},
 	}
+	agentAuth.AttachOptions(cmd, false)
+	return cmd
 }

cli/gitaskpass.go

@@ -18,8 +18,8 @@ import (
 
 // gitAskpass is used by the Coder agent to automatically authenticate
 // with Git providers based on a hostname.
-func (r *RootCmd) gitAskpass() *serpent.Command {
-	return &serpent.Command{
+func gitAskpass(agentAuth *AgentAuth) *serpent.Command {
+	cmd := &serpent.Command{
 		Use:    "gitaskpass",
 		Hidden: true,
 		Handler: func(inv *serpent.Invocation) error {
@@ -33,7 +33,7 @@ func (r *RootCmd) gitAskpass() *serpent.Command {
 				return xerrors.Errorf("parse host: %w", err)
 			}
 
-			client, err := r.tryCreateAgentClient()
+			client, err := agentAuth.CreateClient(ctx)
 			if err != nil {
 				return xerrors.Errorf("create agent client: %w", err)
 			}
@@ -90,4 +90,6 @@ func (r *RootCmd) gitAskpass() *serpent.Command {
 			return nil
 		},
 	}
+	agentAuth.AttachOptions(cmd, false)
+	return cmd
 }

cli/gitaskpass_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/pty/ptytest"
+	"github.com/coder/coder/v2/testutil"
 )
 
 func TestGitAskpass(t *testing.T) {
@@ -32,6 +33,7 @@ func TestGitAskpass(t *testing.T) {
 		url := srv.URL
 		inv, _ := clitest.New(t, "--agent-url", url, "Username for 'https://github.com':")
 		inv.Environ.Set("GIT_PREFIX", "/")
+		inv.Environ.Set("CODER_AGENT_TOKEN", "fake-token")
 		pty := ptytest.New(t)
 		inv.Stdout = pty.Output()
 		clitest.Start(t, inv)
@@ -39,6 +41,7 @@ func TestGitAskpass(t *testing.T) {
 
 		inv, _ = clitest.New(t, "--agent-url", url, "Password for 'https://potato@github.com':")
 		inv.Environ.Set("GIT_PREFIX", "/")
+		inv.Environ.Set("CODER_AGENT_TOKEN", "fake-token")
 		pty = ptytest.New(t)
 		inv.Stdout = pty.Output()
 		clitest.Start(t, inv)
@@ -56,6 +59,7 @@ func TestGitAskpass(t *testing.T) {
 		url := srv.URL
 		inv, _ := clitest.New(t, "--agent-url", url, "--no-open", "Username for 'https://github.com':")
 		inv.Environ.Set("GIT_PREFIX", "/")
+		inv.Environ.Set("CODER_AGENT_TOKEN", "fake-token")
 		pty := ptytest.New(t)
 		inv.Stderr = pty.Output()
 		err := inv.Run()
@@ -65,6 +69,7 @@ func TestGitAskpass(t *testing.T) {
 
 	t.Run("Poll", func(t *testing.T) {
 		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitShort)
 		resp := atomic.Pointer[agentsdk.ExternalAuthResponse]{}
 		resp.Store(&agentsdk.ExternalAuthResponse{
 			URL: "https://something.org",
@@ -86,6 +91,7 @@ func TestGitAskpass(t *testing.T) {
 
 		inv, _ := clitest.New(t, "--agent-url", url, "--no-open", "Username for 'https://github.com':")
 		inv.Environ.Set("GIT_PREFIX", "/")
+		inv.Environ.Set("CODER_AGENT_TOKEN", "fake-token")
 		stdout := ptytest.New(t)
 		inv.Stdout = stdout.Output()
 		stderr := ptytest.New(t)
@@ -94,7 +100,7 @@ func TestGitAskpass(t *testing.T) {
 			err := inv.Run()
 			assert.NoError(t, err)
 		}()
-		<-poll
+		testutil.RequireReceive(ctx, t, poll)
 		stderr.ExpectMatch("Open the following URL to authenticate")
 		resp.Store(&agentsdk.ExternalAuthResponse{
 			Username: "username",

cli/gitssh.go

@@ -18,7 +18,8 @@ import (
 	"github.com/coder/serpent"
 )
 
-func (r *RootCmd) gitssh() *serpent.Command {
+func gitssh() *serpent.Command {
+	agentAuth := &AgentAuth{}
 	cmd := &serpent.Command{
 		Use:    "gitssh",
 		Hidden: true,
@@ -38,7 +39,7 @@ func (r *RootCmd) gitssh() *serpent.Command {
 				return err
 			}
 
-			client, err := r.tryCreateAgentClient()
+			client, err := agentAuth.CreateClient(ctx)
 			if err != nil {
 				return xerrors.Errorf("create agent client: %w", err)
 			}
@@ -108,7 +109,7 @@ func (r *RootCmd) gitssh() *serpent.Command {
 			return nil
 		},
 	}
-
+	agentAuth.AttachOptions(cmd, false)
 	return cmd
 }
 

cli/gitssh_test.go

@@ -54,8 +54,7 @@ func prepareTestGitSSH(ctx context.Context, t *testing.T) (*agentsdk.Client, str
 	}).WithAgent().Do()
 
 	// start workspace agent
-	agentClient := agentsdk.New(client.URL)
-	agentClient.SetSessionToken(r.AgentToken)
+	agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(r.AgentToken))
 	_ = agenttest.New(t, client.URL, r.AgentToken, func(o *agent.Options) {
 		o.Client = agentClient
 	})

cli/root.go

@@ -24,6 +24,7 @@ import (
 	"text/tabwriter"
 	"time"
 
+	"cloud.google.com/go/compute/metadata"
 	"github.com/mattn/go-isatty"
 	"github.com/mitchellh/go-wordwrap"
 	"golang.org/x/mod/semver"
@@ -59,9 +60,6 @@ var (
 const (
 	varURL                     = "url"
 	varToken                   = "token"
-	varAgentToken              = "agent-token"
-	varAgentTokenFile          = "agent-token-file"
-	varAgentURL                = "agent-url"
 	varHeader                  = "header"
 	varHeaderCommand           = "header-command"
 	varNoOpen                  = "no-open"
@@ -82,6 +80,7 @@ const (
 	//nolint:gosec
 	envAgentTokenFile = "CODER_AGENT_TOKEN_FILE"
 	envAgentURL       = "CODER_AGENT_URL"
+	envAgentAuth      = "CODER_AGENT_AUTH"
 	envURL            = "CODER_URL"
 )
 
@@ -90,7 +89,7 @@ func (r *RootCmd) CoreSubcommands() []*serpent.Command {
 	return []*serpent.Command{
 		r.completion(),
 		r.dotfiles(),
-		r.externalAuth(),
+		externalAuth(),
 		r.login(),
 		r.logout(),
 		r.netcheck(),
@@ -130,11 +129,11 @@ func (r *RootCmd) CoreSubcommands() []*serpent.Command {
 		// Hidden
 		r.connectCmd(),
 		r.expCmd(),
-		r.gitssh(),
+		gitssh(),
 		r.support(),
 		r.vpnDaemon(),
 		r.vscodeSSH(),
-		r.workspaceAgent(),
+		workspaceAgent(),
 	}
 }
 
@@ -198,6 +197,7 @@ func (r *RootCmd) RunWithSubcommands(subcommands []*serpent.Command) {
 func (r *RootCmd) Command(subcommands []*serpent.Command) (*serpent.Command, error) {
 	fmtLong := `Coder %s — A tool for provisioning self-hosted development environments with Terraform.
 `
+	hiddenAgentAuth := &AgentAuth{}
 	cmd := &serpent.Command{
 		Use: "coder [global-flags] <subcommand>",
 		Long: fmt.Sprintf(fmtLong, buildinfo.Version()) + FormatExamples(
@@ -220,7 +220,7 @@ func (r *RootCmd) Command(subcommands []*serpent.Command) (*serpent.Command, err
 			// with a `gitaskpass` subcommand, we override the entrypoint
 			// to check if the command was invoked.
 			if gitauth.CheckCommand(i.Args, i.Environ.ToOS()) {
-				return r.gitAskpass().Handler(i)
+				return gitAskpass(hiddenAgentAuth).Handler(i)
 			}
 			return i.Command.HelpHandler(i)
 		},
@@ -349,9 +349,6 @@ func (r *RootCmd) Command(subcommands []*serpent.Command) (*serpent.Command, err
 		}
 	})
 
-	if r.agentURL == nil {
-		r.agentURL = new(url.URL)
-	}
 	if r.clientURL == nil {
 		r.clientURL = new(url.URL)
 	}
@@ -381,30 +378,6 @@ func (r *RootCmd) Command(subcommands []*serpent.Command) (*serpent.Command, err
 			Value:       serpent.StringOf(&r.token),
 			Group:       globalGroup,
 		},
-		{
-			Flag:        varAgentToken,
-			Env:         envAgentToken,
-			Description: "An agent authentication token.",
-			Value:       serpent.StringOf(&r.agentToken),
-			Hidden:      true,
-			Group:       globalGroup,
-		},
-		{
-			Flag:        varAgentTokenFile,
-			Env:         envAgentTokenFile,
-			Description: "A file containing an agent authentication token.",
-			Value:       serpent.StringOf(&r.agentTokenFile),
-			Hidden:      true,
-			Group:       globalGroup,
-		},
-		{
-			Flag:        varAgentURL,
-			Env:         envAgentURL,
-			Description: "URL for an agent to access your deployment.",
-			Value:       serpent.URLOf(r.agentURL),
-			Hidden:      true,
-			Group:       globalGroup,
-		},
 		{
 			Flag:        varNoVersionCheck,
 			Env:         envNoVersionCheck,
@@ -496,26 +469,25 @@ func (r *RootCmd) Command(subcommands []*serpent.Command) (*serpent.Command, err
 			Hidden:      true,
 		},
 	}
+	hiddenAgentAuth.AttachOptions(cmd, true)
 
 	return cmd, nil
 }
 
 // RootCmd contains parameters and helpers useful to all commands.
 type RootCmd struct {
-	clientURL      *url.URL
-	token          string
-	globalConfig   string
-	header         []string
-	headerCommand  string
-	agentToken     string
-	agentTokenFile string
-	agentURL       *url.URL
-	forceTTY       bool
-	noOpen         bool
-	verbose        bool
-	versionFlag    bool
-	disableDirect  bool
-	debugHTTP      bool
+	clientURL     *url.URL
+	token         string
+	globalConfig  string
+	header        []string
+	headerCommand string
+
+	forceTTY      bool
+	noOpen        bool
+	verbose       bool
+	versionFlag   bool
+	disableDirect bool
+	debugHTTP     bool
 
 	disableNetworkTelemetry bool
 	noVersionCheck          bool
@@ -672,38 +644,111 @@ func (r *RootCmd) createUnauthenticatedClient(ctx context.Context, serverURL *ur
 	return &client, err
 }
 
-// createAgentClient returns a new client from the command context.  It works
-// just like InitClient, but uses the agent token and URL instead.
-func (r *RootCmd) createAgentClient() (*agentsdk.Client, error) {
-	agentURL := r.agentURL
-	if agentURL == nil || agentURL.String() == "" {
-		return nil, xerrors.Errorf("%s must be set", envAgentURL)
-	}
-	token := r.agentToken
-	if token == "" {
-		if r.agentTokenFile == "" {
-			return nil, xerrors.Errorf("Either %s or %s must be set", envAgentToken, envAgentTokenFile)
-		}
-		tokenBytes, err := os.ReadFile(r.agentTokenFile)
-		if err != nil {
-			return nil, xerrors.Errorf("read token file %q: %w", r.agentTokenFile, err)
-		}
-		token = strings.TrimSpace(string(tokenBytes))
-	}
-	client := agentsdk.New(agentURL)
-	client.SetSessionToken(token)
-	return client, nil
+type AgentAuth struct {
+	// Agent Client config
+	agentToken     string
+	agentTokenFile string
+	agentURL       url.URL
+	agentAuth      string
 }
 
-// tryCreateAgentClient returns a new client from the command context.  It works
-// just like tryCreateAgentClient, but does not error.
-func (r *RootCmd) tryCreateAgentClient() (*agentsdk.Client, error) {
-	// TODO: Why does this not actually return any errors despite the function
-	// signature?  Could we just use createAgentClient instead, or is it expected
-	// that we return a client in some cases even without a valid URL or token?
-	client := agentsdk.New(r.agentURL)
-	client.SetSessionToken(r.agentToken)
-	return client, nil
+func (a *AgentAuth) AttachOptions(cmd *serpent.Command, hidden bool) {
+	cmd.Options = append(cmd.Options, serpent.Option{
+		Name:        "Agent Token",
+		Description: "An agent authentication token.",
+		Flag:        "agent-token",
+		Env:         envAgentToken,
+		Value:       serpent.StringOf(&a.agentToken),
+		Hidden:      hidden,
+	}, serpent.Option{
+		Name:        "Agent Token File",
+		Description: "A file containing an agent authentication token.",
+		Flag:        "agent-token-file",
+		Env:         envAgentTokenFile,
+		Value:       serpent.StringOf(&a.agentTokenFile),
+		Hidden:      hidden,
+	}, serpent.Option{
+		Name:        "Agent URL",
+		Description: "URL for an agent to access your deployment.",
+		Flag:        "agent-url",
+		Env:         envAgentURL,
+		Value:       serpent.URLOf(&a.agentURL),
+		Hidden:      hidden,
+	}, serpent.Option{
+		Name:        "Agent Auth",
+		Description: "Specify the authentication type to use for the agent.",
+		Flag:        "auth",
+		Env:         envAgentAuth,
+		Default:     "token",
+		Value:       serpent.StringOf(&a.agentAuth),
+		Hidden:      hidden,
+	})
+}
+
+// CreateClient returns a new agent client from the command context.  It works
+// just like InitClient, but uses the agent token and URL instead.
+func (a *AgentAuth) CreateClient(ctx context.Context) (*agentsdk.Client, error) {
+	agentURL := a.agentURL
+	if agentURL.String() == "" {
+		return nil, xerrors.Errorf("%s must be set", envAgentURL)
+	}
+
+	switch a.agentAuth {
+	case "token":
+		token := a.agentToken
+		if token == "" {
+			if a.agentTokenFile == "" {
+				return nil, xerrors.Errorf("Either %s or %s must be set", envAgentToken, envAgentTokenFile)
+			}
+			tokenBytes, err := os.ReadFile(a.agentTokenFile)
+			if err != nil {
+				return nil, xerrors.Errorf("read token file %q: %w", a.agentTokenFile, err)
+			}
+			token = strings.TrimSpace(string(tokenBytes))
+		}
+		if token == "" {
+			return nil, xerrors.Errorf("CODER_AGENT_TOKEN or CODER_AGENT_TOKEN_FILE must be set for token auth")
+		}
+		return agentsdk.New(&a.agentURL, agentsdk.WithFixedToken(token)), nil
+	case "google-instance-identity":
+
+		// This is *only* done for testing to mock client authentication.
+		// This will never be set in a production scenario.
+		var gcpClient *metadata.Client
+		gcpClientRaw := ctx.Value("gcp-client")
+		if gcpClientRaw != nil {
+			gcpClient, _ = gcpClientRaw.(*metadata.Client)
+		}
+		return agentsdk.New(&a.agentURL, agentsdk.WithGoogleInstanceIdentity("", gcpClient)), nil
+	case "aws-instance-identity":
+		client := agentsdk.New(&a.agentURL, agentsdk.WithAWSInstanceIdentity())
+		// This is *only* done for testing to mock client authentication.
+		// This will never be set in a production scenario.
+		var awsClient *http.Client
+		awsClientRaw := ctx.Value("aws-client")
+		if awsClientRaw != nil {
+			awsClient, _ = awsClientRaw.(*http.Client)
+			if awsClient != nil {
+				client.SDK.HTTPClient = awsClient
+			}
+		}
+		return client, nil
+	case "azure-instance-identity":
+		client := agentsdk.New(&a.agentURL, agentsdk.WithAzureInstanceIdentity())
+		// This is *only* done for testing to mock client authentication.
+		// This will never be set in a production scenario.
+		var azureClient *http.Client
+		azureClientRaw := ctx.Value("azure-client")
+		if azureClientRaw != nil {
+			azureClient, _ = azureClientRaw.(*http.Client)
+			if azureClient != nil {
+				client.SDK.HTTPClient = azureClient
+			}
+		}
+		return client, nil
+	default:
+		return nil, xerrors.Errorf("unknown agent auth type: %s", a.agentAuth)
+	}
 }
 
 type OrganizationContext struct {

cli/testdata/coder_agent_--help.golden

@@ -6,6 +6,18 @@ USAGE:
   Starts the Coder workspace agent.
 
 OPTIONS:
+      --auth string, $CODER_AGENT_AUTH (default: token)
+          Specify the authentication type to use for the agent.
+
+      --agent-token string, $CODER_AGENT_TOKEN
+          An agent authentication token.
+
+      --agent-token-file string, $CODER_AGENT_TOKEN_FILE
+          A file containing an agent authentication token.
+
+      --agent-url url, $CODER_AGENT_URL
+          URL for an agent to access your deployment.
+
       --log-human string, $CODER_AGENT_LOGGING_HUMAN (default: /dev/stderr)
           Output human-readable logs to a given file.
 
@@ -24,9 +36,6 @@ OPTIONS:
           requests. The command must output each header as `key=value` on its
           own line.
 
-      --auth string, $CODER_AGENT_AUTH (default: token)
-          Specify the authentication type to use for the agent.
-
       --block-file-transfer bool, $CODER_AGENT_BLOCK_FILE_TRANSFER (default: false)
           Block file transfer using known applications: nc,rsync,scp,sftp.
 

cli/testdata/coder_external-auth_access-token_--help.golden

@@ -25,6 +25,18 @@ USAGE:
        $ coder external-auth access-token slack --extra "authed_user.id"
 
 OPTIONS:
+      --auth string, $CODER_AGENT_AUTH (default: token)
+          Specify the authentication type to use for the agent.
+
+      --agent-token string, $CODER_AGENT_TOKEN
+          An agent authentication token.
+
+      --agent-token-file string, $CODER_AGENT_TOKEN_FILE
+          A file containing an agent authentication token.
+
+      --agent-url url, $CODER_AGENT_URL
+          URL for an agent to access your deployment.
+
       --extra string
           Extract a field from the "extra" properties of the OAuth token.