feat: support multiple agents with shared instance-identity auth (#24325)

> This PR was authored by Mux on behalf of Mike. ## Summary Adds support for multiple peer root workspace agents sharing the same `auth_instance_id`, so AWS, Azure, and GCP instance-identity auth can issue the correct session token for a selected agent instead of assuming a single root agent per instance. ## Problem When a Terraform template attaches two or more `coder_agent` resources (with `auth = "aws-instance-identity"`) to a single compute instance, every agent shares the same cloud instance ID. The existing singular lookup picks whichever agent was created most recently, silently ignoring the others. ## Solution Introduce an optional pre-auth agent selector (`CODER_AGENT_NAME`) and make the server-side lookup ambiguity-aware. **Database layer:** - `GetWorkspaceAgentsByInstanceID` (`:many`): returns all matching root agents for an instance ID. - `GetWorkspaceAgentByInstanceIDAndName` (`:one`): returns the named root agent for disambiguation. **SDK and CLI:** - `agent_name` field added to AWS, Azure, and GCP request structs (`omitempty` for backward compatibility). - `CODER_AGENT_NAME` env var and `--agent-name` flag wired into the agent bootstrap before instance-identity auth runs. **Server handler (`handleAuthInstanceID`):** - When `agent_name` is present: direct lookup by (instance ID, name). - When absent: legacy lookup, then resource-scoped ambiguity check. Returns 409 with available agent names if multiple root agents match. - Whitespace-only names are trimmed and treated as unspecified. - Sub-agents remain excluded (`parent_id IS NULL` filter). **Verification template:** - `examples/templates/aws-multi-agent/` provisions one EC2 instance with two agents (`main` and `dev`), both using instance-identity auth with `CODER_AGENT_NAME` set in the cloud-init user data. ## Backward compatibility Existing single-agent deployments work unchanged. The `agent_name` field is optional with `omitempty`, and the unnamed path preserves today's behavior when only one root agent matches.
2026-09-24 15:04:27 +08:00 · 2026-04-16 13:59:09 +02:00
parent 1cf0354f72
commit e5707a13d6
29 changed files with 1563 additions and 286 deletions
@@ -4,7 +4,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"sort"
+	"strings"

+	"github.com/google/uuid"
 	"github.com/mitchellh/mapstructure"

 	"github.com/coder/coder/v2/coderd/awsidentity"
@@ -26,7 +29,7 @@ import (
 // @Accept json
 // @Produce json
 // @Tags Agents
-// @Param request body agentsdk.AzureInstanceIdentityToken true "Instance identity token"
+// @Param request body agentsdk.AzureInstanceIdentityToken true "Instance identity token. The optional agent_name field disambiguates when multiple agents share the same instance ID."
 // @Success 200 {object} agentsdk.AuthenticateResponse
 // @Router /workspaceagents/azure-instance-identity [post]
 func (api *API) postWorkspaceAuthAzureInstanceIdentity(rw http.ResponseWriter, r *http.Request) {
@@ -45,7 +48,7 @@ func (api *API) postWorkspaceAuthAzureInstanceIdentity(rw http.ResponseWriter, r
 		})
 		return
 	}
-	api.handleAuthInstanceID(rw, r, instanceID)
+	api.handleAuthInstanceID(rw, r, instanceID, req.AgentName)
 }

 // AWS supports instance identity verification:
@@ -58,7 +61,7 @@ func (api *API) postWorkspaceAuthAzureInstanceIdentity(rw http.ResponseWriter, r
 // @Accept json
 // @Produce json
 // @Tags Agents
-// @Param request body agentsdk.AWSInstanceIdentityToken true "Instance identity token"
+// @Param request body agentsdk.AWSInstanceIdentityToken true "Instance identity token. The optional agent_name field disambiguates when multiple agents share the same instance ID."
 // @Success 200 {object} agentsdk.AuthenticateResponse
 // @Router /workspaceagents/aws-instance-identity [post]
 func (api *API) postWorkspaceAuthAWSInstanceIdentity(rw http.ResponseWriter, r *http.Request) {
@@ -75,7 +78,7 @@ func (api *API) postWorkspaceAuthAWSInstanceIdentity(rw http.ResponseWriter, r *
 		})
 		return
 	}
-	api.handleAuthInstanceID(rw, r, identity.InstanceID)
+	api.handleAuthInstanceID(rw, r, identity.InstanceID, req.AgentName)
 }

 // Google Compute Engine supports instance identity verification:
@@ -88,7 +91,7 @@ func (api *API) postWorkspaceAuthAWSInstanceIdentity(rw http.ResponseWriter, r *
 // @Accept json
 // @Produce json
 // @Tags Agents
-// @Param request body agentsdk.GoogleInstanceIdentityToken true "Instance identity token"
+// @Param request body agentsdk.GoogleInstanceIdentityToken true "Instance identity token. The optional agent_name field disambiguates when multiple agents share the same instance ID."
 // @Success 200 {object} agentsdk.AuthenticateResponse
 // @Router /workspaceagents/google-instance-identity [post]
 func (api *API) postWorkspaceAuthGoogleInstanceIdentity(rw http.ResponseWriter, r *http.Request) {
@@ -122,19 +125,18 @@ func (api *API) postWorkspaceAuthGoogleInstanceIdentity(rw http.ResponseWriter,
 		})
 		return
 	}
-	api.handleAuthInstanceID(rw, r, claims.Google.ComputeEngine.InstanceID)
+	api.handleAuthInstanceID(rw, r, claims.Google.ComputeEngine.InstanceID, req.AgentName)
 }

-func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, instanceID string) {
+func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, instanceID string, agentName string) {
 	ctx := r.Context()
-	//nolint:gocritic // needed for auth instance id
-	agent, err := api.Database.GetWorkspaceAgentByInstanceID(dbauthz.AsSystemRestricted(ctx), instanceID)
-	if httpapi.Is404Error(err) {
-		httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
-			Message: fmt.Sprintf("Instance with id %q not found.", instanceID),
-		})
-		return
-	}
+	// Instance identity auth happens before the agent has a session token, so
+	// these lookups must use a restricted system context.
+	//nolint:gocritic // Instance identity auth happens before agent auth.
+	systemCtx := dbauthz.AsSystemRestricted(ctx)
+	agentName = strings.TrimSpace(agentName)
+
+	agents, err := api.Database.GetWorkspaceAgentsByInstanceID(systemCtx, instanceID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error fetching provisioner job agent.",
@@ -142,8 +144,77 @@ func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, in
 		})
 		return
 	}
-	//nolint:gocritic // needed for auth instance id
-	resource, err := api.Database.GetWorkspaceResourceByID(dbauthz.AsSystemRestricted(ctx), agent.ResourceID)
+
+	// Template version agents can share an instance ID with workspace build
+	// agents. Keep only workspace build agents before resolving ambiguity so
+	// template version agents do not force CODER_AGENT_NAME.
+	buildAgents := agents[:0]
+	for _, candidate := range agents {
+		resource, err := api.Database.GetWorkspaceResourceByID(systemCtx, candidate.ResourceID)
+		if err != nil {
+			httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+				Message: "Internal error fetching provisioner job resource.",
+				Detail:  err.Error(),
+			})
+			return
+		}
+		job, err := api.Database.GetProvisionerJobByID(systemCtx, resource.JobID)
+		if err != nil {
+			httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+				Message: "Internal error fetching provisioner job.",
+				Detail:  err.Error(),
+			})
+			return
+		}
+		if job.Type == database.ProvisionerJobTypeWorkspaceBuild {
+			buildAgents = append(buildAgents, candidate)
+		}
+	}
+	agents = buildAgents
+	if len(agents) == 0 {
+		httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
+			Message: fmt.Sprintf("Instance with id %q not found.", instanceID),
+		})
+		return
+	}
+
+	var agent database.WorkspaceAgent
+	if agentName != "" {
+		for _, candidate := range agents {
+			if candidate.Name == agentName {
+				agent = candidate
+				break
+			}
+		}
+		if agent.ID == uuid.Nil {
+			httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
+				Message: fmt.Sprintf("No agent found with instance ID %q and name %q.", instanceID, agentName),
+			})
+			return
+		}
+	} else {
+		if len(agents) != 1 {
+			// Include agent names in the error message to help operators
+			// configure CODER_AGENT_NAME. The caller has already proven
+			// cloud instance identity, so agent names are not sensitive
+			// here.
+			names := make([]string, len(agents))
+			for i, candidate := range agents {
+				names[i] = candidate.Name
+			}
+			sort.Strings(names)
+			httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{
+				Message: fmt.Sprintf(
+					"Multiple agents found with instance ID %q. Set CODER_AGENT_NAME to one of: %s",
+					instanceID,
+					strings.Join(names, ", "),
+				),
+			})
+			return
+		}
+		agent = agents[0]
+	}
+	resource, err := api.Database.GetWorkspaceResourceByID(systemCtx, agent.ResourceID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error fetching provisioner job resource.",
@@ -151,8 +222,7 @@ func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, in
 		})
 		return
 	}
-	//nolint:gocritic // needed for auth instance id
-	job, err := api.Database.GetProvisionerJobByID(dbauthz.AsSystemRestricted(ctx), resource.JobID)
+	job, err := api.Database.GetProvisionerJobByID(systemCtx, resource.JobID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error fetching provisioner job.",
@@ -175,8 +245,7 @@ func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, in
 		})
 		return
 	}
-	//nolint:gocritic // needed for auth instance id
-	resourceHistory, err := api.Database.GetWorkspaceBuildByID(dbauthz.AsSystemRestricted(ctx), jobData.WorkspaceBuildID)
+	resourceHistory, err := api.Database.GetWorkspaceBuildByID(systemCtx, jobData.WorkspaceBuildID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error fetching workspace build.",
@@ -187,8 +256,7 @@ func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, in
 	// This token should only be exchanged if the instance ID is valid
 	// for the latest history. If an instance ID is recycled by a cloud,
 	// we'd hate to leak access to a user's workspace.
-	//nolint:gocritic // needed for auth instance id
-	latestHistory, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(dbauthz.AsSystemRestricted(ctx), resourceHistory.WorkspaceID)
+	latestHistory, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(systemCtx, resourceHistory.WorkspaceID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error fetching the latest workspace build.",