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.

coderd/agentapi/subagent_test.go

@@ -213,8 +213,10 @@ func TestSubAgentAPI(t *testing.T) {
 
 		// Double-check: looking up by the parent's instance ID must
 		// still return the parent, not the sub-agent.
-		lookedUp, err := db.GetWorkspaceAgentByInstanceID(dbauthz.AsSystemRestricted(ctx), parentAgent.AuthInstanceID.String)
+		agents, err := db.GetWorkspaceAgentsByInstanceID(dbauthz.AsSystemRestricted(ctx), parentAgent.AuthInstanceID.String)
 		require.NoError(t, err)
+		require.Len(t, agents, 1)
+		lookedUp := agents[0]
 		assert.Equal(t, parentAgent.ID, lookedUp.ID, "instance ID lookup should still return the parent agent")
 	})
 

coderd/apidoc/docs.go

@@ -10096,7 +10096,7 @@ const docTemplate = `{
                 "operationId": "authenticate-agent-on-aws-instance",
                 "parameters": [
                     {
-                        "description": "Instance identity token",
+                        "description": "Instance identity token. The optional agent_name field disambiguates when multiple agents share the same instance ID.",
                         "name": "request",
                         "in": "body",
                         "required": true,
@@ -10135,7 +10135,7 @@ const docTemplate = `{
                 "operationId": "authenticate-agent-on-azure-instance",
                 "parameters": [
                     {
-                        "description": "Instance identity token",
+                        "description": "Instance identity token. The optional agent_name field disambiguates when multiple agents share the same instance ID.",
                         "name": "request",
                         "in": "body",
                         "required": true,
@@ -10202,7 +10202,7 @@ const docTemplate = `{
                 "operationId": "authenticate-agent-on-google-cloud-instance",
                 "parameters": [
                     {
-                        "description": "Instance identity token",
+                        "description": "Instance identity token. The optional agent_name field disambiguates when multiple agents share the same instance ID.",
                         "name": "request",
                         "in": "body",
                         "required": true,
@@ -12780,6 +12780,10 @@ const docTemplate = `{
                 "signature"
             ],
             "properties": {
+                "agent_name": {
+                    "description": "AgentName optionally selects a specific agent when multiple\nagents share the same instance identity. An empty string is\ntreated as unspecified.",
+                    "type": "string"
+                },
                 "document": {
                     "type": "string"
                 },
@@ -12803,6 +12807,10 @@ const docTemplate = `{
                 "signature"
             ],
             "properties": {
+                "agent_name": {
+                    "description": "AgentName optionally selects a specific agent when multiple\nagents share the same instance identity. An empty string is\ntreated as unspecified.",
+                    "type": "string"
+                },
                 "encoding": {
                     "type": "string"
                 },
@@ -12853,6 +12861,10 @@ const docTemplate = `{
                 "json_web_token"
             ],
             "properties": {
+                "agent_name": {
+                    "description": "AgentName optionally selects a specific agent when multiple\nagents share the same instance identity. An empty string is\ntreated as unspecified.",
+                    "type": "string"
+                },
                 "json_web_token": {
                     "type": "string"
                 }

coderd/apidoc/swagger.json

@@ -8949,7 +8949,7 @@
 				"operationId": "authenticate-agent-on-aws-instance",
 				"parameters": [
 					{
-						"description": "Instance identity token",
+						"description": "Instance identity token. The optional agent_name field disambiguates when multiple agents share the same instance ID.",
 						"name": "request",
 						"in": "body",
 						"required": true,
@@ -8982,7 +8982,7 @@
 				"operationId": "authenticate-agent-on-azure-instance",
 				"parameters": [
 					{
-						"description": "Instance identity token",
+						"description": "Instance identity token. The optional agent_name field disambiguates when multiple agents share the same instance ID.",
 						"name": "request",
 						"in": "body",
 						"required": true,
@@ -9039,7 +9039,7 @@
 				"operationId": "authenticate-agent-on-google-cloud-instance",
 				"parameters": [
 					{
-						"description": "Instance identity token",
+						"description": "Instance identity token. The optional agent_name field disambiguates when multiple agents share the same instance ID.",
 						"name": "request",
 						"in": "body",
 						"required": true,
@@ -11337,6 +11337,10 @@
 			"type": "object",
 			"required": ["document", "signature"],
 			"properties": {
+				"agent_name": {
+					"description": "AgentName optionally selects a specific agent when multiple\nagents share the same instance identity. An empty string is\ntreated as unspecified.",
+					"type": "string"
+				},
 				"document": {
 					"type": "string"
 				},
@@ -11357,6 +11361,10 @@
 			"type": "object",
 			"required": ["encoding", "signature"],
 			"properties": {
+				"agent_name": {
+					"description": "AgentName optionally selects a specific agent when multiple\nagents share the same instance identity. An empty string is\ntreated as unspecified.",
+					"type": "string"
+				},
 				"encoding": {
 					"type": "string"
 				},
@@ -11405,6 +11413,10 @@
 			"type": "object",
 			"required": ["json_web_token"],
 			"properties": {
+				"agent_name": {
+					"description": "AgentName optionally selects a specific agent when multiple\nagents share the same instance identity. An empty string is\ntreated as unspecified.",
+					"type": "string"
+				},
 				"json_web_token": {
 					"type": "string"
 				}

coderd/database/dbauthz/dbauthz.go

@@ -4422,22 +4422,6 @@ func (q *querier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (data
 	return q.db.GetWorkspaceAgentByID(ctx, id)
 }
 
-// GetWorkspaceAgentByInstanceID might want to be a system call? Unsure exactly,
-// but this will fail. Need to figure out what AuthInstanceID is, and if it
-// is essentially an auth token. But the caller using this function is not
-// an authenticated user. So this authz check will fail.
-func (q *querier) GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (database.WorkspaceAgent, error) {
-	agent, err := q.db.GetWorkspaceAgentByInstanceID(ctx, authInstanceID)
-	if err != nil {
-		return database.WorkspaceAgent{}, err
-	}
-	_, err = q.GetWorkspaceByAgentID(ctx, agent.ID)
-	if err != nil {
-		return database.WorkspaceAgent{}, err
-	}
-	return agent, nil
-}
-
 func (q *querier) GetWorkspaceAgentDevcontainersByAgentID(ctx context.Context, workspaceAgentID uuid.UUID) ([]database.WorkspaceAgentDevcontainer, error) {
 	_, err := q.GetWorkspaceAgentByID(ctx, workspaceAgentID)
 	if err != nil {
@@ -4527,6 +4511,33 @@ func (q *querier) GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, crea
 	return q.db.GetWorkspaceAgentUsageStatsAndLabels(ctx, createdAt)
 }
 
+func (q *querier) GetWorkspaceAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]database.WorkspaceAgent, error) {
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err == nil {
+		return q.db.GetWorkspaceAgentsByInstanceID(ctx, authInstanceID)
+	}
+
+	agents, err := q.db.GetWorkspaceAgentsByInstanceID(ctx, authInstanceID)
+	if err != nil {
+		return nil, err
+	}
+	// Filter to agents whose workspace is accessible. Template-version
+	// agents can share the same instance ID but do not belong to a
+	// workspace, so GetWorkspaceByAgentID returns sql.ErrNoRows for
+	// them. Exclude those agents rather than failing the entire lookup.
+	filtered := make([]database.WorkspaceAgent, 0, len(agents))
+	for _, agent := range agents {
+		_, err = q.GetWorkspaceByAgentID(ctx, agent.ID)
+		if err != nil {
+			if errors.Is(err, sql.ErrNoRows) {
+				continue
+			}
+			return nil, err
+		}
+		filtered = append(filtered, agent)
+	}
+	return filtered, nil
+}
+
 func (q *querier) GetWorkspaceAgentsByParentID(ctx context.Context, parentID uuid.UUID) ([]database.WorkspaceAgent, error) {
 	workspace, err := q.db.GetWorkspaceByAgentID(ctx, parentID)
 	if err != nil {

coderd/database/dbauthz/dbauthz_test.go

@@ -3012,13 +3012,16 @@ func (s *MethodTestSuite) TestWorkspace() {
 		dbm.EXPECT().BatchUpdateWorkspaceAgentMetadata(gomock.Any(), arg).Return(nil).AnyTimes()
 		check.Args(arg).Asserts(rbac.ResourceWorkspace.All(), policy.ActionUpdate).Returns()
 	}))
-	s.Run("GetWorkspaceAgentByInstanceID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+	s.Run("GetWorkspaceAgentsByInstanceID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
 		w := testutil.Fake(s.T(), faker, database.Workspace{})
 		agt := testutil.Fake(s.T(), faker, database.WorkspaceAgent{})
 		authInstanceID := "instance-id"
-		dbm.EXPECT().GetWorkspaceAgentByInstanceID(gomock.Any(), authInstanceID).Return(agt, nil).AnyTimes()
+		dbm.EXPECT().GetWorkspaceAgentsByInstanceID(gomock.Any(), authInstanceID).Return([]database.WorkspaceAgent{agt}, nil).AnyTimes()
 		dbm.EXPECT().GetWorkspaceByAgentID(gomock.Any(), agt.ID).Return(w, nil).AnyTimes()
-		check.Args(authInstanceID).Asserts(w, policy.ActionRead).Returns(agt)
+		check.Args(authInstanceID).
+			Asserts(rbac.ResourceSystem, policy.ActionRead, w, policy.ActionRead).
+			Returns([]database.WorkspaceAgent{agt}).
+			FailSystemObjectChecks()
 	}))
 	s.Run("UpdateWorkspaceAgentLifecycleStateByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
 		w := testutil.Fake(s.T(), faker, database.Workspace{})

coderd/database/dbmetrics/querymetrics.go

@@ -2864,14 +2864,6 @@ func (m queryMetricsStore) GetWorkspaceAgentByID(ctx context.Context, id uuid.UU
 	return r0, r1
 }
 
-func (m queryMetricsStore) GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (database.WorkspaceAgent, error) {
-	start := time.Now()
-	r0, r1 := m.s.GetWorkspaceAgentByInstanceID(ctx, authInstanceID)
-	m.queryLatencies.WithLabelValues("GetWorkspaceAgentByInstanceID").Observe(time.Since(start).Seconds())
-	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceAgentByInstanceID").Inc()
-	return r0, r1
-}
-
 func (m queryMetricsStore) GetWorkspaceAgentDevcontainersByAgentID(ctx context.Context, workspaceAgentID uuid.UUID) ([]database.WorkspaceAgentDevcontainer, error) {
 	start := time.Now()
 	r0, r1 := m.s.GetWorkspaceAgentDevcontainersByAgentID(ctx, workspaceAgentID)
@@ -2968,6 +2960,14 @@ func (m queryMetricsStore) GetWorkspaceAgentUsageStatsAndLabels(ctx context.Cont
 	return r0, r1
 }
 
+func (m queryMetricsStore) GetWorkspaceAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]database.WorkspaceAgent, error) {
+	start := time.Now()
+	r0, r1 := m.s.GetWorkspaceAgentsByInstanceID(ctx, authInstanceID)
+	m.queryLatencies.WithLabelValues("GetWorkspaceAgentsByInstanceID").Observe(time.Since(start).Seconds())
+	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceAgentsByInstanceID").Inc()
+	return r0, r1
+}
+
 func (m queryMetricsStore) GetWorkspaceAgentsByParentID(ctx context.Context, parentID uuid.UUID) ([]database.WorkspaceAgent, error) {
 	start := time.Now()
 	r0, r1 := m.s.GetWorkspaceAgentsByParentID(ctx, parentID)

coderd/database/dbmock/dbmock.go

@@ -5357,21 +5357,6 @@ func (mr *MockStoreMockRecorder) GetWorkspaceAgentByID(ctx, id any) *gomock.Call
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentByID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentByID), ctx, id)
 }
 
-// GetWorkspaceAgentByInstanceID mocks base method.
-func (m *MockStore) GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (database.WorkspaceAgent, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetWorkspaceAgentByInstanceID", ctx, authInstanceID)
-	ret0, _ := ret[0].(database.WorkspaceAgent)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// GetWorkspaceAgentByInstanceID indicates an expected call of GetWorkspaceAgentByInstanceID.
-func (mr *MockStoreMockRecorder) GetWorkspaceAgentByInstanceID(ctx, authInstanceID any) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentByInstanceID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentByInstanceID), ctx, authInstanceID)
-}
-
 // GetWorkspaceAgentDevcontainersByAgentID mocks base method.
 func (m *MockStore) GetWorkspaceAgentDevcontainersByAgentID(ctx context.Context, workspaceAgentID uuid.UUID) ([]database.WorkspaceAgentDevcontainer, error) {
 	m.ctrl.T.Helper()
@@ -5552,6 +5537,21 @@ func (mr *MockStoreMockRecorder) GetWorkspaceAgentUsageStatsAndLabels(ctx, creat
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentUsageStatsAndLabels", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentUsageStatsAndLabels), ctx, createdAt)
 }
 
+// GetWorkspaceAgentsByInstanceID mocks base method.
+func (m *MockStore) GetWorkspaceAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]database.WorkspaceAgent, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetWorkspaceAgentsByInstanceID", ctx, authInstanceID)
+	ret0, _ := ret[0].([]database.WorkspaceAgent)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetWorkspaceAgentsByInstanceID indicates an expected call of GetWorkspaceAgentsByInstanceID.
+func (mr *MockStoreMockRecorder) GetWorkspaceAgentsByInstanceID(ctx, authInstanceID any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentsByInstanceID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentsByInstanceID), ctx, authInstanceID)
+}
+
 // GetWorkspaceAgentsByParentID mocks base method.
 func (m *MockStore) GetWorkspaceAgentsByParentID(ctx context.Context, parentID uuid.UUID) ([]database.WorkspaceAgent, error) {
 	m.ctrl.T.Helper()

coderd/database/querier.go

@@ -683,7 +683,6 @@ type sqlcQuerier interface {
 	GetWorkspaceACLByID(ctx context.Context, id uuid.UUID) (GetWorkspaceACLByIDRow, error)
 	GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (GetWorkspaceAgentAndWorkspaceByIDRow, error)
 	GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (WorkspaceAgent, error)
-	GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (WorkspaceAgent, error)
 	GetWorkspaceAgentDevcontainersByAgentID(ctx context.Context, workspaceAgentID uuid.UUID) ([]WorkspaceAgentDevcontainer, error)
 	GetWorkspaceAgentLifecycleStateByID(ctx context.Context, id uuid.UUID) (GetWorkspaceAgentLifecycleStateByIDRow, error)
 	GetWorkspaceAgentLogSourcesByAgentIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceAgentLogSource, error)
@@ -697,6 +696,7 @@ type sqlcQuerier interface {
 	// `minute_buckets` could return 0 rows if there are no usage stats since `created_at`.
 	GetWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsRow, error)
 	GetWorkspaceAgentUsageStatsAndLabels(ctx context.Context, createdAt time.Time) ([]GetWorkspaceAgentUsageStatsAndLabelsRow, error)
+	GetWorkspaceAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]WorkspaceAgent, error)
 	GetWorkspaceAgentsByParentID(ctx context.Context, parentID uuid.UUID) ([]WorkspaceAgent, error)
 	GetWorkspaceAgentsByResourceIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceAgent, error)
 	GetWorkspaceAgentsByWorkspaceAndBuildNumber(ctx context.Context, arg GetWorkspaceAgentsByWorkspaceAndBuildNumberParams) ([]WorkspaceAgent, error)

coderd/database/querier_test.go

@@ -7184,38 +7184,55 @@ func TestGetWorkspaceAgentsByParentID(t *testing.T) {
 	})
 }
 
-func TestGetWorkspaceAgentByInstanceID(t *testing.T) {
+func setupWorkspaceAgentQueryResources(t *testing.T, db database.Store, count int) []database.WorkspaceResource {
+	t.Helper()
+
+	org := dbgen.Organization(t, db, database.Organization{})
+	job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+		Type:           database.ProvisionerJobTypeTemplateVersionImport,
+		OrganizationID: org.ID,
+	})
+
+	resources := make([]database.WorkspaceResource, 0, count)
+	for i := 0; i < count; i++ {
+		resources = append(resources, dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID: job.ID,
+		}))
+	}
+
+	return resources
+}
+
+func markWorkspaceAgentDeleted(ctx context.Context, t *testing.T, sqlDB *sql.DB, agentID uuid.UUID) {
+	t.Helper()
+
+	_, err := sqlDB.ExecContext(ctx, "UPDATE workspace_agents SET deleted = TRUE WHERE id = $1", agentID)
+	require.NoError(t, err)
+}
+
+func TestGetWorkspaceAgentsByInstanceID(t *testing.T) {
 	t.Parallel()
 
-	// Context: https://github.com/coder/coder/pull/22196
-	t.Run("DoesNotReturnSubAgents", func(t *testing.T) {
+	t.Run("ReturnsAllMatchingRootAgents", func(t *testing.T) {
 		t.Parallel()
 
-		// Given: A parent workspace agent with an AuthInstanceID and a
-		// sub-agent that shares the same AuthInstanceID.
 		db, _ := dbtestutil.NewDB(t)
-		org := dbgen.Organization(t, db, database.Organization{})
-		job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
-			Type:           database.ProvisionerJobTypeTemplateVersionImport,
-			OrganizationID: org.ID,
-		})
-		resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
-			JobID: job.ID,
-		})
-
+		resources := setupWorkspaceAgentQueryResources(t, db, 2)
 		authInstanceID := fmt.Sprintf("instance-%s-%d", t.Name(), time.Now().UnixNano())
-		parentAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
-			ResourceID: resource.ID,
+		olderCreatedAt := dbtime.Now().Add(-time.Hour)
+		newerCreatedAt := dbtime.Now()
+
+		olderAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: resources[0].ID,
+			CreatedAt:  olderCreatedAt,
 			AuthInstanceID: sql.NullString{
 				String: authInstanceID,
 				Valid:  true,
 			},
 		})
-		// Create a sub-agent with the same AuthInstanceID (simulating
-		// the old behavior before the fix).
-		_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
-			ParentID:   uuid.NullUUID{UUID: parentAgent.ID, Valid: true},
-			ResourceID: resource.ID,
+		newerAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: resources[1].ID,
+			CreatedAt:  newerCreatedAt,
 			AuthInstanceID: sql.NullString{
 				String: authInstanceID,
 				Valid:  true,
@@ -7224,13 +7241,89 @@ func TestGetWorkspaceAgentByInstanceID(t *testing.T) {
 
 		ctx := testutil.Context(t, testutil.WaitShort)
 
-		// When: We look up the agent by instance ID.
-		agent, err := db.GetWorkspaceAgentByInstanceID(ctx, authInstanceID)
+		agents, err := db.GetWorkspaceAgentsByInstanceID(ctx, authInstanceID)
 		require.NoError(t, err)
+		require.Len(t, agents, 2)
+		assert.Equal(t, []uuid.UUID{newerAgent.ID, olderAgent.ID}, []uuid.UUID{agents[0].ID, agents[1].ID})
+	})
 
-		// Then: The result must be the parent agent, not the sub-agent.
-		assert.Equal(t, parentAgent.ID, agent.ID, "instance ID lookup should return the parent agent, not a sub-agent")
-		assert.False(t, agent.ParentID.Valid, "returned agent should not have a parent (should be the parent itself)")
+	t.Run("ExcludesDeletedAndSubAgents", func(t *testing.T) {
+		t.Parallel()
+
+		db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
+		resources := setupWorkspaceAgentQueryResources(t, db, 2)
+		authInstanceID := fmt.Sprintf("instance-%s-%d", t.Name(), time.Now().UnixNano())
+		baseCreatedAt := dbtime.Now()
+
+		rootAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: resources[0].ID,
+			CreatedAt:  baseCreatedAt.Add(-time.Hour),
+			AuthInstanceID: sql.NullString{
+				String: authInstanceID,
+				Valid:  true,
+			},
+		})
+		_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ParentID:   uuid.NullUUID{UUID: rootAgent.ID, Valid: true},
+			ResourceID: resources[0].ID,
+			CreatedAt:  baseCreatedAt,
+			AuthInstanceID: sql.NullString{
+				String: authInstanceID,
+				Valid:  true,
+			},
+		})
+		deletedRootAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: resources[1].ID,
+			CreatedAt:  baseCreatedAt.Add(time.Minute),
+			AuthInstanceID: sql.NullString{
+				String: authInstanceID,
+				Valid:  true,
+			},
+		})
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		markWorkspaceAgentDeleted(ctx, t, sqlDB, deletedRootAgent.ID)
+
+		agents, err := db.GetWorkspaceAgentsByInstanceID(ctx, authInstanceID)
+		require.NoError(t, err)
+		require.Len(t, agents, 1)
+		assert.Equal(t, rootAgent.ID, agents[0].ID)
+		assert.False(t, agents[0].ParentID.Valid)
+	})
+
+	t.Run("OrdersNewestFirst", func(t *testing.T) {
+		t.Parallel()
+
+		db, _ := dbtestutil.NewDB(t)
+		resources := setupWorkspaceAgentQueryResources(t, db, 2)
+		authInstanceID := fmt.Sprintf("instance-%s-%d", t.Name(), time.Now().UnixNano())
+		olderCreatedAt := dbtime.Now().Add(-time.Hour)
+		newerCreatedAt := dbtime.Now()
+
+		olderAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: resources[0].ID,
+			CreatedAt:  olderCreatedAt,
+			AuthInstanceID: sql.NullString{
+				String: authInstanceID,
+				Valid:  true,
+			},
+		})
+		newerAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: resources[1].ID,
+			CreatedAt:  newerCreatedAt,
+			AuthInstanceID: sql.NullString{
+				String: authInstanceID,
+				Valid:  true,
+			},
+		})
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+
+		agents, err := db.GetWorkspaceAgentsByInstanceID(ctx, authInstanceID)
+		require.NoError(t, err)
+		require.Len(t, agents, 2)
+		assert.Equal(t, newerAgent.ID, agents[0].ID)
+		assert.Equal(t, olderAgent.ID, agents[1].ID)
 	})
 }
 

coderd/database/queries.sql.go

@@ -26566,63 +26566,6 @@ func (q *sqlQuerier) GetWorkspaceAgentByID(ctx context.Context, id uuid.UUID) (W
 	return i, err
 }
 
-const getWorkspaceAgentByInstanceID = `-- name: GetWorkspaceAgentByInstanceID :one
-SELECT
-	id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, instance_metadata, resource_metadata, directory, version, last_connected_replica_id, connection_timeout_seconds, troubleshooting_url, motd_file, lifecycle_state, expanded_directory, logs_length, logs_overflowed, started_at, ready_at, subsystems, display_apps, api_version, display_order, parent_id, api_key_scope, deleted
-FROM
-	workspace_agents
-WHERE
-	auth_instance_id = $1 :: TEXT
-	-- Filter out deleted sub agents.
-	AND deleted = FALSE
-	-- Filter out sub agents, they do not authenticate with auth_instance_id.
-	AND parent_id IS NULL
-ORDER BY
-	created_at DESC
-`
-
-func (q *sqlQuerier) GetWorkspaceAgentByInstanceID(ctx context.Context, authInstanceID string) (WorkspaceAgent, error) {
-	row := q.db.QueryRowContext(ctx, getWorkspaceAgentByInstanceID, authInstanceID)
-	var i WorkspaceAgent
-	err := row.Scan(
-		&i.ID,
-		&i.CreatedAt,
-		&i.UpdatedAt,
-		&i.Name,
-		&i.FirstConnectedAt,
-		&i.LastConnectedAt,
-		&i.DisconnectedAt,
-		&i.ResourceID,
-		&i.AuthToken,
-		&i.AuthInstanceID,
-		&i.Architecture,
-		&i.EnvironmentVariables,
-		&i.OperatingSystem,
-		&i.InstanceMetadata,
-		&i.ResourceMetadata,
-		&i.Directory,
-		&i.Version,
-		&i.LastConnectedReplicaID,
-		&i.ConnectionTimeoutSeconds,
-		&i.TroubleshootingURL,
-		&i.MOTDFile,
-		&i.LifecycleState,
-		&i.ExpandedDirectory,
-		&i.LogsLength,
-		&i.LogsOverflowed,
-		&i.StartedAt,
-		&i.ReadyAt,
-		pq.Array(&i.Subsystems),
-		pq.Array(&i.DisplayApps),
-		&i.APIVersion,
-		&i.DisplayOrder,
-		&i.ParentID,
-		&i.APIKeyScope,
-		&i.Deleted,
-	)
-	return i, err
-}
-
 const getWorkspaceAgentLifecycleStateByID = `-- name: GetWorkspaceAgentLifecycleStateByID :one
 SELECT
 	lifecycle_state,
@@ -26836,6 +26779,79 @@ func (q *sqlQuerier) GetWorkspaceAgentScriptTimingsByBuildID(ctx context.Context
 	return items, nil
 }
 
+const getWorkspaceAgentsByInstanceID = `-- name: GetWorkspaceAgentsByInstanceID :many
+SELECT
+	id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, instance_metadata, resource_metadata, directory, version, last_connected_replica_id, connection_timeout_seconds, troubleshooting_url, motd_file, lifecycle_state, expanded_directory, logs_length, logs_overflowed, started_at, ready_at, subsystems, display_apps, api_version, display_order, parent_id, api_key_scope, deleted
+FROM
+	workspace_agents
+WHERE
+	auth_instance_id = $1 :: TEXT
+	-- Filter out deleted agents.
+	AND deleted = FALSE
+	-- Filter out sub agents, they do not authenticate with auth_instance_id.
+	AND parent_id IS NULL
+ORDER BY
+	created_at DESC
+`
+
+func (q *sqlQuerier) GetWorkspaceAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]WorkspaceAgent, error) {
+	rows, err := q.db.QueryContext(ctx, getWorkspaceAgentsByInstanceID, authInstanceID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items []WorkspaceAgent
+	for rows.Next() {
+		var i WorkspaceAgent
+		if err := rows.Scan(
+			&i.ID,
+			&i.CreatedAt,
+			&i.UpdatedAt,
+			&i.Name,
+			&i.FirstConnectedAt,
+			&i.LastConnectedAt,
+			&i.DisconnectedAt,
+			&i.ResourceID,
+			&i.AuthToken,
+			&i.AuthInstanceID,
+			&i.Architecture,
+			&i.EnvironmentVariables,
+			&i.OperatingSystem,
+			&i.InstanceMetadata,
+			&i.ResourceMetadata,
+			&i.Directory,
+			&i.Version,
+			&i.LastConnectedReplicaID,
+			&i.ConnectionTimeoutSeconds,
+			&i.TroubleshootingURL,
+			&i.MOTDFile,
+			&i.LifecycleState,
+			&i.ExpandedDirectory,
+			&i.LogsLength,
+			&i.LogsOverflowed,
+			&i.StartedAt,
+			&i.ReadyAt,
+			pq.Array(&i.Subsystems),
+			pq.Array(&i.DisplayApps),
+			&i.APIVersion,
+			&i.DisplayOrder,
+			&i.ParentID,
+			&i.APIKeyScope,
+			&i.Deleted,
+		); err != nil {
+			return nil, err
+		}
+		items = append(items, i)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
 const getWorkspaceAgentsByParentID = `-- name: GetWorkspaceAgentsByParentID :many
 SELECT
 	id, created_at, updated_at, name, first_connected_at, last_connected_at, disconnected_at, resource_id, auth_token, auth_instance_id, architecture, environment_variables, operating_system, instance_metadata, resource_metadata, directory, version, last_connected_replica_id, connection_timeout_seconds, troubleshooting_url, motd_file, lifecycle_state, expanded_directory, logs_length, logs_overflowed, started_at, ready_at, subsystems, display_apps, api_version, display_order, parent_id, api_key_scope, deleted

coderd/database/queries/workspaceagents.sql

@@ -8,14 +8,14 @@ WHERE
 	-- Filter out deleted sub agents.
 	AND deleted = FALSE;
 
--- name: GetWorkspaceAgentByInstanceID :one
+-- name: GetWorkspaceAgentsByInstanceID :many
 SELECT
 	*
 FROM
 	workspace_agents
 WHERE
 	auth_instance_id = @auth_instance_id :: TEXT
-	-- Filter out deleted sub agents.
+	-- Filter out deleted agents.
 	AND deleted = FALSE
 	-- Filter out sub agents, they do not authenticate with auth_instance_id.
 	AND parent_id IS NULL

coderd/provisionerdserver/provisionerdserver_test.go

@@ -4286,8 +4286,10 @@ func TestInsertWorkspaceResource(t *testing.T) {
 
 					// Looking up by the parent's instance ID must still
 					// return the parent, not the sub-agent.
-					lookedUp, err := db.GetWorkspaceAgentByInstanceID(ctx, parentAgent.AuthInstanceID.String)
+					agents, err := db.GetWorkspaceAgentsByInstanceID(ctx, parentAgent.AuthInstanceID.String)
 					require.NoError(t, err)
+					require.Len(t, agents, 1)
+					lookedUp := agents[0]
 					assert.Equal(t, parentAgent.ID, lookedUp.ID, "instance ID lookup should still return the parent agent")
 				},
 			},

coderd/workspaceresourceauth.go

@@ -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.",

coderd/workspaceresourceauth_test.go

@@ -2,12 +2,20 @@ package coderd_test
 
 import (
 	"context"
+	"database/sql"
+	"encoding/json"
+	"fmt"
+	"io"
 	"net/http"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/v2/coderd/coderdtest"
+	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/dbauthz"
+	"github.com/coder/coder/v2/coderd/database/dbgen"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/provisioner/echo"
@@ -17,96 +25,274 @@ import (
 
 func TestPostWorkspaceAuthAzureInstanceIdentity(t *testing.T) {
 	t.Parallel()
-	instanceID := "instanceidentifier"
-	certificates, metadataClient := coderdtest.NewAzureInstanceIdentity(t, instanceID)
-	client := coderdtest.New(t, &coderdtest.Options{
-		AzureCertificates:        certificates,
-		IncludeProvisionerDaemon: true,
-	})
-	user := coderdtest.CreateFirstUser(t, client)
-	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
-		Parse: echo.ParseComplete,
-		ProvisionGraph: []*proto.Response{{
-			Type: &proto.Response_Graph{
-				Graph: &proto.GraphComplete{
-					Resources: []*proto.Resource{{
-						Name: "somename",
-						Type: "someinstance",
-						Agents: []*proto.Agent{{
-							Name: "dev",
-							Auth: &proto.Agent_InstanceId{
-								InstanceId: instanceID,
-							},
-						}},
-					}},
-				},
-			},
-		}},
-	})
-	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
-	coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
-	workspace := coderdtest.CreateWorkspace(t, client, template.ID)
-	coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
 
-	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
-	defer cancel()
+	t.Run("Success", func(t *testing.T) {
+		t.Parallel()
 
-	agentClient := agentsdk.New(client.URL, agentsdk.WithAzureInstanceIdentity())
-	agentClient.SDK.HTTPClient = metadataClient
-	err := agentClient.RefreshToken(ctx)
-	require.NoError(t, err)
+		instanceID := newTestInstanceID(t)
+		certificates, metadataClient := coderdtest.NewAzureInstanceIdentity(t, instanceID)
+		client, _ := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			AzureCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "dev"))
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		agentClient := agentsdk.New(client.URL, agentsdk.WithAzureInstanceIdentity())
+		agentClient.SDK.HTTPClient = metadataClient
+
+		err := agentClient.RefreshToken(ctx)
+		require.NoError(t, err)
+	})
+
+	t.Run("Ambiguous/AzureWithSelector", func(t *testing.T) {
+		t.Parallel()
+
+		instanceID := newTestInstanceID(t)
+		certificates, metadataClient := coderdtest.NewAzureInstanceIdentity(t, instanceID)
+		client, store := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			AzureCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "alpha", "beta"))
+
+		expectedAgent := requireWorkspaceAgentByInstanceIDAndName(t, store, instanceID, "alpha")
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		agentClient := agentsdk.New(client.URL, agentsdk.WithAzureInstanceIdentity(
+			agentsdk.WithInstanceIdentityAgentName("alpha"),
+		))
+		agentClient.SDK.HTTPClient = metadataClient
+
+		err := agentClient.RefreshToken(ctx)
+		require.NoError(t, err)
+		require.Equal(t, expectedAgent.AuthToken.String(), agentClient.SDK.SessionToken())
+	})
 }
 
 func TestPostWorkspaceAuthAWSInstanceIdentity(t *testing.T) {
 	t.Parallel()
-	t.Run("Success", func(t *testing.T) {
+
+	t.Run("Ambiguous/SingleAgent", func(t *testing.T) {
 		t.Parallel()
-		instanceID := "instanceidentifier"
+
+		instanceID := newTestInstanceID(t)
 		certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
-		client := coderdtest.New(t, &coderdtest.Options{
-			AWSCertificates:          certificates,
-			IncludeProvisionerDaemon: true,
-		})
-		user := coderdtest.CreateFirstUser(t, client)
-		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
-			Parse: echo.ParseComplete,
-			ProvisionGraph: []*proto.Response{{
-				Type: &proto.Response_Graph{
-					Graph: &proto.GraphComplete{
-						Resources: []*proto.Resource{{
-							Name: "somename",
-							Type: "someinstance",
-							Agents: []*proto.Agent{{
-								Name: "dev",
-								Auth: &proto.Agent_InstanceId{
-									InstanceId: instanceID,
-								},
-							}},
-						}},
-					},
-				},
-			}},
-		})
-		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
-		coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
-		workspace := coderdtest.CreateWorkspace(t, client, template.ID)
-		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
+		client, _ := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			AWSCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "dev"))
 
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 		defer cancel()
 
 		agentClient := agentsdk.New(client.URL, agentsdk.WithAWSInstanceIdentity())
 		agentClient.SDK.HTTPClient = metadataClient
+
 		err := agentClient.RefreshToken(ctx)
 		require.NoError(t, err)
 	})
+
+	t.Run("Ambiguous/MultipleAgentsNoSelector", func(t *testing.T) {
+		t.Parallel()
+
+		instanceID := newTestInstanceID(t)
+		certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
+		client, _ := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			AWSCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "alpha", "beta"))
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		agentClient := agentsdk.New(client.URL, agentsdk.WithAWSInstanceIdentity())
+		agentClient.SDK.HTTPClient = metadataClient
+
+		err := agentClient.RefreshToken(ctx)
+		var apiErr *codersdk.Error
+		require.ErrorAs(t, err, &apiErr)
+		require.Equal(t, http.StatusConflict, apiErr.StatusCode())
+		require.Contains(t, apiErr.Message, "CODER_AGENT_NAME")
+		require.Contains(t, apiErr.Message, "alpha, beta")
+	})
+
+	t.Run("Ambiguous/EmptyAgentNameTreatedAsUnset", func(t *testing.T) {
+		t.Parallel()
+
+		instanceID := newTestInstanceID(t)
+		certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
+		client, _ := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			AWSCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "alpha", "beta"))
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		signatureReq, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://169.254.169.254/latest/dynamic/instance-identity/signature", nil)
+		require.NoError(t, err)
+		signatureRes, err := metadataClient.Do(signatureReq)
+		require.NoError(t, err)
+		defer signatureRes.Body.Close()
+		signature, err := io.ReadAll(signatureRes.Body)
+		require.NoError(t, err)
+
+		documentReq, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://169.254.169.254/latest/dynamic/instance-identity/document", nil)
+		require.NoError(t, err)
+		documentRes, err := metadataClient.Do(documentReq)
+		require.NoError(t, err)
+		defer documentRes.Body.Close()
+		document, err := io.ReadAll(documentRes.Body)
+		require.NoError(t, err)
+
+		reqBody, err := json.Marshal(map[string]string{
+			"signature":  string(signature),
+			"document":   string(document),
+			"agent_name": "",
+		})
+		require.NoError(t, err)
+
+		res, err := client.RequestWithoutSessionToken(ctx, http.MethodPost, "/api/v2/workspaceagents/aws-instance-identity", reqBody)
+		require.NoError(t, err)
+		defer res.Body.Close()
+
+		require.Equal(t, http.StatusConflict, res.StatusCode)
+		err = codersdk.ReadBodyAsError(res)
+		var apiErr *codersdk.Error
+		require.ErrorAs(t, err, &apiErr)
+		require.Equal(t, http.StatusConflict, apiErr.StatusCode())
+		require.Contains(t, apiErr.Message, "CODER_AGENT_NAME")
+		require.Contains(t, apiErr.Message, "alpha, beta")
+	})
+
+	t.Run("Ambiguous/WhitespaceAgentNameTreatedAsUnset", func(t *testing.T) {
+		t.Parallel()
+
+		instanceID := newTestInstanceID(t)
+		certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
+		client, _ := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			AWSCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "alpha", "beta"))
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		signatureReq, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://169.254.169.254/latest/dynamic/instance-identity/signature", nil)
+		require.NoError(t, err)
+		signatureRes, err := metadataClient.Do(signatureReq)
+		require.NoError(t, err)
+		defer signatureRes.Body.Close()
+		signature, err := io.ReadAll(signatureRes.Body)
+		require.NoError(t, err)
+
+		documentReq, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://169.254.169.254/latest/dynamic/instance-identity/document", nil)
+		require.NoError(t, err)
+		documentRes, err := metadataClient.Do(documentReq)
+		require.NoError(t, err)
+		defer documentRes.Body.Close()
+		document, err := io.ReadAll(documentRes.Body)
+		require.NoError(t, err)
+
+		reqBody, err := json.Marshal(map[string]string{
+			"signature":  string(signature),
+			"document":   string(document),
+			"agent_name": "   ",
+		})
+		require.NoError(t, err)
+
+		res, err := client.RequestWithoutSessionToken(ctx, http.MethodPost, "/api/v2/workspaceagents/aws-instance-identity", reqBody)
+		require.NoError(t, err)
+		defer res.Body.Close()
+
+		require.Equal(t, http.StatusConflict, res.StatusCode)
+		err = codersdk.ReadBodyAsError(res)
+		var apiErr *codersdk.Error
+		require.ErrorAs(t, err, &apiErr)
+		require.Equal(t, http.StatusConflict, apiErr.StatusCode())
+		require.Contains(t, apiErr.Message, "CODER_AGENT_NAME")
+		require.Contains(t, apiErr.Message, "alpha, beta")
+	})
+
+	t.Run("Ambiguous/MultipleAgentsWithSelector", func(t *testing.T) {
+		t.Parallel()
+
+		instanceID := newTestInstanceID(t)
+		certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
+		client, store := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			AWSCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "alpha", "beta"))
+
+		expectedAgent := requireWorkspaceAgentByInstanceIDAndName(t, store, instanceID, "alpha")
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		agentClient := agentsdk.New(client.URL, agentsdk.WithAWSInstanceIdentity(
+			agentsdk.WithInstanceIdentityAgentName("alpha"),
+		))
+		agentClient.SDK.HTTPClient = metadataClient
+
+		err := agentClient.RefreshToken(ctx)
+		require.NoError(t, err)
+		require.Equal(t, expectedAgent.AuthToken.String(), agentClient.SDK.SessionToken())
+	})
+
+	t.Run("Ambiguous/MultipleAgentsUnknownSelector", func(t *testing.T) {
+		t.Parallel()
+
+		instanceID := newTestInstanceID(t)
+		certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
+		client, _ := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			AWSCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "alpha", "beta"))
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		agentClient := agentsdk.New(client.URL, agentsdk.WithAWSInstanceIdentity(
+			agentsdk.WithInstanceIdentityAgentName("nonexistent"),
+		))
+		agentClient.SDK.HTTPClient = metadataClient
+
+		err := agentClient.RefreshToken(ctx)
+		var apiErr *codersdk.Error
+		require.ErrorAs(t, err, &apiErr)
+		require.Equal(t, http.StatusNotFound, apiErr.StatusCode())
+	})
+
+	t.Run("Ambiguous/SubAgentExcluded", func(t *testing.T) {
+		t.Parallel()
+
+		instanceID := newTestInstanceID(t)
+		certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
+		client, store := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			AWSCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "dev"))
+
+		rootAgent := requireWorkspaceAgentByInstanceIDAndName(t, store, instanceID, "dev")
+		_ = dbgen.WorkspaceSubAgent(t, store, rootAgent, database.WorkspaceAgent{
+			Name: "sub",
+			AuthInstanceID: sql.NullString{
+				String: instanceID,
+				Valid:  true,
+			},
+		})
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		agentClient := agentsdk.New(client.URL, agentsdk.WithAWSInstanceIdentity())
+		agentClient.SDK.HTTPClient = metadataClient
+
+		err := agentClient.RefreshToken(ctx)
+		require.NoError(t, err)
+		require.Equal(t, rootAgent.AuthToken.String(), agentClient.SDK.SessionToken())
+	})
 }
 
 func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) {
 	t.Parallel()
+
 	t.Run("Expired", func(t *testing.T) {
 		t.Parallel()
-		instanceID := "instanceidentifier"
+
+		instanceID := newTestInstanceID(t)
 		validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, true)
 		client := coderdtest.New(t, &coderdtest.Options{
 			GoogleTokenValidator: validator,
@@ -124,7 +310,8 @@ func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) {
 
 	t.Run("InstanceNotFound", func(t *testing.T) {
 		t.Parallel()
-		instanceID := "instanceidentifier"
+
+		instanceID := newTestInstanceID(t)
 		validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
 		client := coderdtest.New(t, &coderdtest.Options{
 			GoogleTokenValidator: validator,
@@ -142,36 +329,12 @@ func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) {
 
 	t.Run("Success", func(t *testing.T) {
 		t.Parallel()
-		instanceID := "instanceidentifier"
+
+		instanceID := newTestInstanceID(t)
 		validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
-		client := coderdtest.New(t, &coderdtest.Options{
-			GoogleTokenValidator:     validator,
-			IncludeProvisionerDaemon: true,
-		})
-		user := coderdtest.CreateFirstUser(t, client)
-		version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
-			Parse: echo.ParseComplete,
-			ProvisionGraph: []*proto.Response{{
-				Type: &proto.Response_Graph{
-					Graph: &proto.GraphComplete{
-						Resources: []*proto.Resource{{
-							Name: "somename",
-							Type: "someinstance",
-							Agents: []*proto.Agent{{
-								Name: "dev",
-								Auth: &proto.Agent_InstanceId{
-									InstanceId: instanceID,
-								},
-							}},
-						}},
-					},
-				},
-			}},
-		})
-		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
-		coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
-		workspace := coderdtest.CreateWorkspace(t, client, template.ID)
-		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
+		client, _ := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			GoogleTokenValidator: validator,
+		}, workspaceAgentsForInstanceID(instanceID, "dev"))
 
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
 		defer cancel()
@@ -180,4 +343,91 @@ func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) {
 		err := agentClient.RefreshToken(ctx)
 		require.NoError(t, err)
 	})
+
+	t.Run("Ambiguous/GoogleWithSelector", func(t *testing.T) {
+		t.Parallel()
+
+		instanceID := newTestInstanceID(t)
+		validator, metadata := coderdtest.NewGoogleInstanceIdentity(t, instanceID, false)
+		client, store := setupInstanceIDWorkspace(t, &coderdtest.Options{
+			GoogleTokenValidator: validator,
+		}, workspaceAgentsForInstanceID(instanceID, "alpha", "beta"))
+
+		expectedAgent := requireWorkspaceAgentByInstanceIDAndName(t, store, instanceID, "alpha")
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		agentClient := agentsdk.New(client.URL, agentsdk.WithGoogleInstanceIdentity(
+			"",
+			metadata,
+			agentsdk.WithInstanceIdentityAgentName("alpha"),
+		))
+		err := agentClient.RefreshToken(ctx)
+		require.NoError(t, err)
+		require.Equal(t, expectedAgent.AuthToken.String(), agentClient.SDK.SessionToken())
+	})
+}
+
+func setupInstanceIDWorkspace(t *testing.T, opts *coderdtest.Options, agents []*proto.Agent) (*codersdk.Client, database.Store) {
+	t.Helper()
+
+	actualOpts := &coderdtest.Options{}
+	if opts != nil {
+		*actualOpts = *opts
+	}
+	actualOpts.IncludeProvisionerDaemon = true
+
+	client, store := coderdtest.NewWithDatabase(t, actualOpts)
+	user := coderdtest.CreateFirstUser(t, client)
+	version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
+		Parse: echo.ParseComplete,
+		ProvisionGraph: []*proto.Response{{
+			Type: &proto.Response_Graph{
+				Graph: &proto.GraphComplete{
+					Resources: []*proto.Resource{{
+						Name:   "resource",
+						Type:   "instance",
+						Agents: agents,
+					}},
+				},
+			},
+		}},
+	})
+	template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+	coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+	workspace := coderdtest.CreateWorkspace(t, client, template.ID)
+	coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
+
+	return client, store
+}
+
+func workspaceAgentsForInstanceID(instanceID string, names ...string) []*proto.Agent {
+	agents := make([]*proto.Agent, 0, len(names))
+	for _, name := range names {
+		agents = append(agents, &proto.Agent{
+			Name: name,
+			Auth: &proto.Agent_InstanceId{InstanceId: instanceID},
+		})
+	}
+	return agents
+}
+
+func requireWorkspaceAgentByInstanceIDAndName(t testing.TB, store database.Store, instanceID string, name string) database.WorkspaceAgent {
+	t.Helper()
+
+	ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong))
+	agents, err := store.GetWorkspaceAgentsByInstanceID(ctx, instanceID)
+	require.NoError(t, err)
+	for _, agent := range agents {
+		if agent.Name == name {
+			return agent
+		}
+	}
+	require.FailNow(t, "workspace agent not found", "instance ID %q, name %q", instanceID, name)
+	return database.WorkspaceAgent{}
+}
+
+func newTestInstanceID(t testing.TB) string {
+	t.Helper()
+	return fmt.Sprintf("instance-%d", time.Now().UnixNano())
 }