fix(coderd): filter build instance agents in SQL (#25031)

Replaces the per-agent Go-side template-version filter in
`handleAuthInstanceID` with a purpose-built SQL query.

`GetWorkspaceBuildAgentsByInstanceID` joins `workspace_agents ->
workspace_resources -> workspace_builds -> provisioner_jobs ->
workspaces` and excludes:

- non-`workspace_build` provisioner jobs (template-version-import,
dry-run)
- deleted agents and sub-agents
- deleted workspaces

The handler:

- drops the per-candidate `GetWorkspaceResourceByID` /
`GetProvisionerJobByID` lookups
- drops the `provisioner_jobs.input` JSON parsing and the follow-up
`GetWorkspaceBuildByID` call
- compares `latestHistory.ID` against `selected.WorkspaceBuildID`
returned directly from the query
- preserves the existing recycled-instance safety check and matching
response codes

One intentional behavior tightening: agents whose workspace is deleted
now return 404 (previously they could reach the recycled-instance check
and return 400, or 200 if the stale build was still latest). This
matches the existing token-auth path, which already refuses to
authenticate against deleted workspaces.

The original `GetWorkspaceAgentsByInstanceID` query is intentionally
untouched. It remains the generic raw lookup used elsewhere in tests and
helpers.

The dbauthz wrapper for the new query uses the system-read fast path
with `fetchWithPostFilter` for non-system reads, with `RBACObject()`
delegating to the embedded `WorkspaceTable`.

Tests:

- new `TestGetWorkspaceBuildAgentsByInstanceID` covering newest-first
ordering, exclusion of deleted/sub agents, exclusion of template-import
and dry-run jobs, and exclusion of deleted workspaces
- new dbauthz mock test for `GetWorkspaceBuildAgentsByInstanceID`
- new `TestPostWorkspaceAuthAWSInstanceIdentity/RecycledInstanceID`
exercising the recycled-instance rejection branch (HTTP 400 when the
agent's build is no longer latest)
- existing `TestPostWorkspaceAuth{AWS,Azure,Google}InstanceIdentity`
continue to cover the handler end to end (including the template-version
+ workspace-build same-instance-ID scenario via
`setupInstanceIDWorkspace`)

> Mux is acting on Mike's behalf.

coderd/database/dbauthz/dbauthz.go

@@ -4873,6 +4873,14 @@ func (q *querier) GetWorkspaceAppsCreatedAfter(ctx context.Context, createdAt ti
 	return q.db.GetWorkspaceAppsCreatedAfter(ctx, createdAt)
 }
 
+func (q *querier) GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]database.GetWorkspaceBuildAgentsByInstanceIDRow, error) {
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err == nil {
+		return q.db.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID)
+	}
+
+	return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetWorkspaceBuildAgentsByInstanceID)(ctx, authInstanceID)
+}
+
 func (q *querier) GetWorkspaceBuildByID(ctx context.Context, buildID uuid.UUID) (database.WorkspaceBuild, error) {
 	build, err := q.db.GetWorkspaceBuildByID(ctx, buildID)
 	if err != nil {

coderd/database/dbauthz/dbauthz_test.go

@@ -3257,6 +3257,19 @@ func (s *MethodTestSuite) TestWorkspace() {
 			Returns([]database.WorkspaceAgent{agt}).
 			FailSystemObjectChecks()
 	}))
+	s.Run("GetWorkspaceBuildAgentsByInstanceID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		w := testutil.Fake(s.T(), faker, database.WorkspaceTable{})
+		agt := testutil.Fake(s.T(), faker, database.WorkspaceAgent{})
+		row := testutil.Fake(s.T(), faker, database.GetWorkspaceBuildAgentsByInstanceIDRow{})
+		row.WorkspaceAgent = agt
+		row.WorkspaceTable = w
+		authInstanceID := "instance-id"
+		dbm.EXPECT().GetWorkspaceBuildAgentsByInstanceID(gomock.Any(), authInstanceID).Return([]database.GetWorkspaceBuildAgentsByInstanceIDRow{row}, nil).AnyTimes()
+		check.Args(authInstanceID).
+			Asserts(rbac.ResourceSystem, policy.ActionRead, w, policy.ActionRead).
+			Returns([]database.GetWorkspaceBuildAgentsByInstanceIDRow{row}).
+			FailSystemObjectChecks()
+	}))
 	s.Run("UpdateWorkspaceAgentLifecycleStateByID", 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{})

coderd/database/dbmetrics/querymetrics.go

@@ -3249,6 +3249,14 @@ func (m queryMetricsStore) GetWorkspaceAppsCreatedAfter(ctx context.Context, cre
 	return r0, r1
 }
 
+func (m queryMetricsStore) GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]database.GetWorkspaceBuildAgentsByInstanceIDRow, error) {
+	start := time.Now()
+	r0, r1 := m.s.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID)
+	m.queryLatencies.WithLabelValues("GetWorkspaceBuildAgentsByInstanceID").Observe(time.Since(start).Seconds())
+	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceBuildAgentsByInstanceID").Inc()
+	return r0, r1
+}
+
 func (m queryMetricsStore) GetWorkspaceBuildByID(ctx context.Context, id uuid.UUID) (database.WorkspaceBuild, error) {
 	start := time.Now()
 	r0, r1 := m.s.GetWorkspaceBuildByID(ctx, id)

coderd/database/dbmock/dbmock.go

@@ -6078,6 +6078,21 @@ func (mr *MockStoreMockRecorder) GetWorkspaceAppsCreatedAfter(ctx, createdAt any
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAppsCreatedAfter", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAppsCreatedAfter), ctx, createdAt)
 }
 
+// GetWorkspaceBuildAgentsByInstanceID mocks base method.
+func (m *MockStore) GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]database.GetWorkspaceBuildAgentsByInstanceIDRow, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetWorkspaceBuildAgentsByInstanceID", ctx, authInstanceID)
+	ret0, _ := ret[0].([]database.GetWorkspaceBuildAgentsByInstanceIDRow)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetWorkspaceBuildAgentsByInstanceID indicates an expected call of GetWorkspaceBuildAgentsByInstanceID.
+func (mr *MockStoreMockRecorder) GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceBuildAgentsByInstanceID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceBuildAgentsByInstanceID), ctx, authInstanceID)
+}
+
 // GetWorkspaceBuildByID mocks base method.
 func (m *MockStore) GetWorkspaceBuildByID(ctx context.Context, id uuid.UUID) (database.WorkspaceBuild, error) {
 	m.ctrl.T.Helper()

coderd/database/modelmethods.go

@@ -932,6 +932,11 @@ func (r GetWorkspaceAgentAndWorkspaceByIDRow) RBACObject() rbac.Object {
 	return r.WorkspaceTable.RBACObject()
 }
 
+// A workspace agent belongs to the owner of the associated workspace.
+func (r GetWorkspaceBuildAgentsByInstanceIDRow) RBACObject() rbac.Object {
+	return r.WorkspaceTable.RBACObject()
+}
+
 // UpsertConnectionLogParams contains the parameters for upserting a
 // connection log entry. This struct is hand-maintained (not generated
 // by sqlc) because the single-row UpsertConnectionLog query was

coderd/database/querier.go

@@ -802,6 +802,7 @@ type sqlcQuerier interface {
 	GetWorkspaceAppsByAgentID(ctx context.Context, agentID uuid.UUID) ([]WorkspaceApp, error)
 	GetWorkspaceAppsByAgentIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceApp, error)
 	GetWorkspaceAppsCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceApp, error)
+	GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]GetWorkspaceBuildAgentsByInstanceIDRow, error)
 	GetWorkspaceBuildByID(ctx context.Context, id uuid.UUID) (WorkspaceBuild, error)
 	GetWorkspaceBuildByJobID(ctx context.Context, jobID uuid.UUID) (WorkspaceBuild, error)
 	GetWorkspaceBuildByWorkspaceIDAndBuildNumber(ctx context.Context, arg GetWorkspaceBuildByWorkspaceIDAndBuildNumberParams) (WorkspaceBuild, error)

coderd/database/querier_test.go

@@ -7187,6 +7187,103 @@ func markWorkspaceAgentDeleted(ctx context.Context, t *testing.T, sqlDB *sql.DB,
 	require.NoError(t, err)
 }
 
+type workspaceBuildAgentQueryFixture struct {
+	Workspace database.WorkspaceTable
+	Build     database.WorkspaceBuild
+	Agent     database.WorkspaceAgent
+}
+
+func setupWorkspaceBuildAgentQueryWorkspace(t testing.TB, db database.Store, deleted bool) database.WorkspaceTable {
+	t.Helper()
+
+	org := dbgen.Organization(t, db, database.Organization{})
+	user := dbgen.User(t, db, database.User{})
+	template := dbgen.Template(t, db, database.Template{
+		CreatedBy:      user.ID,
+		OrganizationID: org.ID,
+	})
+	return dbgen.Workspace(t, db, database.WorkspaceTable{
+		OwnerID:        user.ID,
+		OrganizationID: org.ID,
+		TemplateID:     template.ID,
+		Deleted:        deleted,
+	})
+}
+
+func setupWorkspaceBuildAgentQueryFixture(
+	t testing.TB,
+	db database.Store,
+	authInstanceID string,
+	name string,
+	createdAt time.Time,
+	workspace database.WorkspaceTable,
+) workspaceBuildAgentQueryFixture {
+	t.Helper()
+
+	if workspace.ID == uuid.Nil {
+		workspace = setupWorkspaceBuildAgentQueryWorkspace(t, db, false)
+	}
+	templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
+		TemplateID:     uuid.NullUUID{UUID: workspace.TemplateID, Valid: true},
+		OrganizationID: workspace.OrganizationID,
+		CreatedBy:      workspace.OwnerID,
+	})
+	job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+		OrganizationID: workspace.OrganizationID,
+		Type:           database.ProvisionerJobTypeWorkspaceBuild,
+	})
+	build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+		WorkspaceID:       workspace.ID,
+		TemplateVersionID: templateVersion.ID,
+		JobID:             job.ID,
+	})
+	resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+		JobID: job.ID,
+	})
+	agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+		Name:       name,
+		ResourceID: resource.ID,
+		CreatedAt:  createdAt,
+		AuthInstanceID: sql.NullString{
+			String: authInstanceID,
+			Valid:  true,
+		},
+	})
+
+	return workspaceBuildAgentQueryFixture{
+		Workspace: workspace,
+		Build:     build,
+		Agent:     agent,
+	}
+}
+
+func setupProvisionerJobAgentQueryFixture(
+	t testing.TB,
+	db database.Store,
+	authInstanceID string,
+	name string,
+	createdAt time.Time,
+	jobType database.ProvisionerJobType,
+) database.WorkspaceAgent {
+	t.Helper()
+
+	job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+		Type: jobType,
+	})
+	resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+		JobID: job.ID,
+	})
+	return dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+		Name:       name,
+		ResourceID: resource.ID,
+		CreatedAt:  createdAt,
+		AuthInstanceID: sql.NullString{
+			String: authInstanceID,
+			Valid:  true,
+		},
+	})
+}
+
 func TestGetWorkspaceAgentsByInstanceID(t *testing.T) {
 	t.Parallel()
 
@@ -7304,6 +7401,90 @@ func TestGetWorkspaceAgentsByInstanceID(t *testing.T) {
 	})
 }
 
+func TestGetWorkspaceBuildAgentsByInstanceID(t *testing.T) {
+	t.Parallel()
+
+	t.Run("ReturnsWorkspaceBuildRootAgentsNewestFirst", func(t *testing.T) {
+		t.Parallel()
+
+		db, _ := dbtestutil.NewDB(t)
+		authInstanceID := fmt.Sprintf("instance-%s-%d", t.Name(), time.Now().UnixNano())
+		olderCreatedAt := dbtime.Now().Add(-time.Hour)
+		newerCreatedAt := dbtime.Now()
+
+		older := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "older", olderCreatedAt, database.WorkspaceTable{})
+		newer := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "newer", newerCreatedAt, database.WorkspaceTable{})
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+
+		agents, err := db.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID)
+		require.NoError(t, err)
+		require.Len(t, agents, 2)
+		assert.Equal(t, []uuid.UUID{newer.Agent.ID, older.Agent.ID}, []uuid.UUID{agents[0].WorkspaceAgent.ID, agents[1].WorkspaceAgent.ID})
+		assert.Equal(t, []uuid.UUID{newer.Build.ID, older.Build.ID}, []uuid.UUID{agents[0].WorkspaceBuildID, agents[1].WorkspaceBuildID})
+		assert.Equal(t, newer.Workspace.ID, agents[0].WorkspaceTable.ID)
+		assert.Equal(t, older.Workspace.ID, agents[1].WorkspaceTable.ID)
+		assert.Equal(t, newer.Workspace.OwnerID, agents[0].WorkspaceTable.OwnerID)
+		assert.Equal(t, older.Workspace.OwnerID, agents[1].WorkspaceTable.OwnerID)
+		assert.Equal(t, newer.Workspace.OrganizationID, agents[0].WorkspaceTable.OrganizationID)
+		assert.Equal(t, older.Workspace.OrganizationID, agents[1].WorkspaceTable.OrganizationID)
+		assert.False(t, agents[0].WorkspaceTable.Deleted)
+		assert.False(t, agents[1].WorkspaceTable.Deleted)
+	})
+
+	t.Run("ExcludesDeletedAgentsSubAgentsAndNonWorkspaceBuildJobs", func(t *testing.T) {
+		t.Parallel()
+
+		db, _, sqlDB := dbtestutil.NewDBWithSQLDB(t)
+		authInstanceID := fmt.Sprintf("instance-%s-%d", t.Name(), time.Now().UnixNano())
+		baseCreatedAt := dbtime.Now()
+
+		root := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "root", baseCreatedAt.Add(-time.Hour), database.WorkspaceTable{})
+		_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ParentID:   uuid.NullUUID{UUID: root.Agent.ID, Valid: true},
+			Name:       "sub",
+			ResourceID: root.Agent.ResourceID,
+			CreatedAt:  baseCreatedAt.Add(time.Minute),
+			AuthInstanceID: sql.NullString{
+				String: authInstanceID,
+				Valid:  true,
+			},
+		})
+		deletedAgent := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "deleted", baseCreatedAt.Add(2*time.Minute), database.WorkspaceTable{})
+		_ = setupProvisionerJobAgentQueryFixture(t, db, authInstanceID, "template-import", baseCreatedAt.Add(3*time.Minute), database.ProvisionerJobTypeTemplateVersionImport)
+		_ = setupProvisionerJobAgentQueryFixture(t, db, authInstanceID, "dry-run", baseCreatedAt.Add(4*time.Minute), database.ProvisionerJobTypeTemplateVersionDryRun)
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+		markWorkspaceAgentDeleted(ctx, t, sqlDB, deletedAgent.Agent.ID)
+
+		agents, err := db.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID)
+		require.NoError(t, err)
+		require.Len(t, agents, 1)
+		assert.Equal(t, root.Agent.ID, agents[0].WorkspaceAgent.ID)
+		assert.False(t, agents[0].WorkspaceAgent.ParentID.Valid)
+		assert.Equal(t, root.Build.ID, agents[0].WorkspaceBuildID)
+	})
+
+	t.Run("ExcludesDeletedWorkspaces", func(t *testing.T) {
+		t.Parallel()
+
+		db, _ := dbtestutil.NewDB(t)
+		authInstanceID := fmt.Sprintf("instance-%s-%d", t.Name(), time.Now().UnixNano())
+		baseCreatedAt := dbtime.Now()
+		active := setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "active", baseCreatedAt, database.WorkspaceTable{})
+		deletedWorkspace := setupWorkspaceBuildAgentQueryWorkspace(t, db, true)
+		_ = setupWorkspaceBuildAgentQueryFixture(t, db, authInstanceID, "deleted-workspace", baseCreatedAt.Add(time.Minute), deletedWorkspace)
+
+		ctx := testutil.Context(t, testutil.WaitShort)
+
+		agents, err := db.GetWorkspaceBuildAgentsByInstanceID(ctx, authInstanceID)
+		require.NoError(t, err)
+		require.Len(t, agents, 1)
+		assert.Equal(t, active.Agent.ID, agents[0].WorkspaceAgent.ID)
+		assert.Equal(t, active.Workspace.ID, agents[0].WorkspaceTable.ID)
+	})
+}
+
 func requireUsersMatch(t testing.TB, expected []database.User, found []database.GetUsersRow, msg string) {
 	t.Helper()
 	require.ElementsMatch(t, expected, database.ConvertUserRows(found), msg)

coderd/database/queries.sql.go

@@ -28766,6 +28766,122 @@ func (q *sqlQuerier) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Co
 	return items, nil
 }
 
+const getWorkspaceBuildAgentsByInstanceID = `-- name: GetWorkspaceBuildAgentsByInstanceID :many
+SELECT
+	workspace_agents.id, workspace_agents.created_at, workspace_agents.updated_at, workspace_agents.name, workspace_agents.first_connected_at, workspace_agents.last_connected_at, workspace_agents.disconnected_at, workspace_agents.resource_id, workspace_agents.auth_token, workspace_agents.auth_instance_id, workspace_agents.architecture, workspace_agents.environment_variables, workspace_agents.operating_system, workspace_agents.instance_metadata, workspace_agents.resource_metadata, workspace_agents.directory, workspace_agents.version, workspace_agents.last_connected_replica_id, workspace_agents.connection_timeout_seconds, workspace_agents.troubleshooting_url, workspace_agents.motd_file, workspace_agents.lifecycle_state, workspace_agents.expanded_directory, workspace_agents.logs_length, workspace_agents.logs_overflowed, workspace_agents.started_at, workspace_agents.ready_at, workspace_agents.subsystems, workspace_agents.display_apps, workspace_agents.api_version, workspace_agents.display_order, workspace_agents.parent_id, workspace_agents.api_key_scope, workspace_agents.deleted,
+	workspace_builds.id AS workspace_build_id,
+	workspaces.id, workspaces.created_at, workspaces.updated_at, workspaces.owner_id, workspaces.organization_id, workspaces.template_id, workspaces.deleted, workspaces.name, workspaces.autostart_schedule, workspaces.ttl, workspaces.last_used_at, workspaces.dormant_at, workspaces.deleting_at, workspaces.automatic_updates, workspaces.favorite, workspaces.next_start_at, workspaces.group_acl, workspaces.user_acl
+FROM
+	workspace_agents
+JOIN
+	workspace_resources
+ON
+	workspace_resources.id = workspace_agents.resource_id
+JOIN
+	workspace_builds
+ON
+	workspace_builds.job_id = workspace_resources.job_id
+JOIN
+	provisioner_jobs
+ON
+	provisioner_jobs.id = workspace_builds.job_id
+JOIN
+	workspaces
+ON
+	workspaces.id = workspace_builds.workspace_id
+WHERE
+	workspace_agents.auth_instance_id = $1 :: TEXT
+	AND workspace_agents.deleted = FALSE
+	AND workspace_agents.parent_id IS NULL
+	AND provisioner_jobs.type = 'workspace_build'::provisioner_job_type
+	AND workspaces.deleted = FALSE
+ORDER BY
+	workspace_agents.created_at DESC
+`
+
+type GetWorkspaceBuildAgentsByInstanceIDRow struct {
+	WorkspaceAgent   WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"`
+	WorkspaceBuildID uuid.UUID      `db:"workspace_build_id" json:"workspace_build_id"`
+	WorkspaceTable   WorkspaceTable `db:"workspace_table" json:"workspace_table"`
+}
+
+func (q *sqlQuerier) GetWorkspaceBuildAgentsByInstanceID(ctx context.Context, authInstanceID string) ([]GetWorkspaceBuildAgentsByInstanceIDRow, error) {
+	rows, err := q.db.QueryContext(ctx, getWorkspaceBuildAgentsByInstanceID, authInstanceID)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items []GetWorkspaceBuildAgentsByInstanceIDRow
+	for rows.Next() {
+		var i GetWorkspaceBuildAgentsByInstanceIDRow
+		if err := rows.Scan(
+			&i.WorkspaceAgent.ID,
+			&i.WorkspaceAgent.CreatedAt,
+			&i.WorkspaceAgent.UpdatedAt,
+			&i.WorkspaceAgent.Name,
+			&i.WorkspaceAgent.FirstConnectedAt,
+			&i.WorkspaceAgent.LastConnectedAt,
+			&i.WorkspaceAgent.DisconnectedAt,
+			&i.WorkspaceAgent.ResourceID,
+			&i.WorkspaceAgent.AuthToken,
+			&i.WorkspaceAgent.AuthInstanceID,
+			&i.WorkspaceAgent.Architecture,
+			&i.WorkspaceAgent.EnvironmentVariables,
+			&i.WorkspaceAgent.OperatingSystem,
+			&i.WorkspaceAgent.InstanceMetadata,
+			&i.WorkspaceAgent.ResourceMetadata,
+			&i.WorkspaceAgent.Directory,
+			&i.WorkspaceAgent.Version,
+			&i.WorkspaceAgent.LastConnectedReplicaID,
+			&i.WorkspaceAgent.ConnectionTimeoutSeconds,
+			&i.WorkspaceAgent.TroubleshootingURL,
+			&i.WorkspaceAgent.MOTDFile,
+			&i.WorkspaceAgent.LifecycleState,
+			&i.WorkspaceAgent.ExpandedDirectory,
+			&i.WorkspaceAgent.LogsLength,
+			&i.WorkspaceAgent.LogsOverflowed,
+			&i.WorkspaceAgent.StartedAt,
+			&i.WorkspaceAgent.ReadyAt,
+			pq.Array(&i.WorkspaceAgent.Subsystems),
+			pq.Array(&i.WorkspaceAgent.DisplayApps),
+			&i.WorkspaceAgent.APIVersion,
+			&i.WorkspaceAgent.DisplayOrder,
+			&i.WorkspaceAgent.ParentID,
+			&i.WorkspaceAgent.APIKeyScope,
+			&i.WorkspaceAgent.Deleted,
+			&i.WorkspaceBuildID,
+			&i.WorkspaceTable.ID,
+			&i.WorkspaceTable.CreatedAt,
+			&i.WorkspaceTable.UpdatedAt,
+			&i.WorkspaceTable.OwnerID,
+			&i.WorkspaceTable.OrganizationID,
+			&i.WorkspaceTable.TemplateID,
+			&i.WorkspaceTable.Deleted,
+			&i.WorkspaceTable.Name,
+			&i.WorkspaceTable.AutostartSchedule,
+			&i.WorkspaceTable.Ttl,
+			&i.WorkspaceTable.LastUsedAt,
+			&i.WorkspaceTable.DormantAt,
+			&i.WorkspaceTable.DeletingAt,
+			&i.WorkspaceTable.AutomaticUpdates,
+			&i.WorkspaceTable.Favorite,
+			&i.WorkspaceTable.NextStartAt,
+			&i.WorkspaceTable.GroupACL,
+			&i.WorkspaceTable.UserACL,
+		); 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 insertWorkspaceAgent = `-- name: InsertWorkspaceAgent :one
 INSERT INTO
 	workspace_agents (

coderd/database/queries/workspaceagents.sql

@@ -22,6 +22,38 @@ WHERE
 ORDER BY
 	created_at DESC;
 
+-- name: GetWorkspaceBuildAgentsByInstanceID :many
+SELECT
+	sqlc.embed(workspace_agents),
+	workspace_builds.id AS workspace_build_id,
+	sqlc.embed(workspaces)
+FROM
+	workspace_agents
+JOIN
+	workspace_resources
+ON
+	workspace_resources.id = workspace_agents.resource_id
+JOIN
+	workspace_builds
+ON
+	workspace_builds.job_id = workspace_resources.job_id
+JOIN
+	provisioner_jobs
+ON
+	provisioner_jobs.id = workspace_builds.job_id
+JOIN
+	workspaces
+ON
+	workspaces.id = workspace_builds.workspace_id
+WHERE
+	workspace_agents.auth_instance_id = @auth_instance_id :: TEXT
+	AND workspace_agents.deleted = FALSE
+	AND workspace_agents.parent_id IS NULL
+	AND provisioner_jobs.type = 'workspace_build'::provisioner_job_type
+	AND workspaces.deleted = FALSE
+ORDER BY
+	workspace_agents.created_at DESC;
+
 -- name: GetWorkspaceAgentsByResourceIDs :many
 SELECT
 	*

coderd/workspaceresourceauth.go

@@ -1,21 +1,17 @@
 package coderd
 
 import (
-	"encoding/json"
 	"fmt"
 	"net/http"
 	"sort"
 	"strings"
 
-	"github.com/google/uuid"
 	"github.com/mitchellh/mapstructure"
 
 	"github.com/coder/coder/v2/coderd/awsidentity"
 	"github.com/coder/coder/v2/coderd/azureidentity"
-	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/httpapi"
-	"github.com/coder/coder/v2/coderd/provisionerdserver"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 )
@@ -136,120 +132,61 @@ func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, in
 	systemCtx := dbauthz.AsSystemRestricted(ctx)
 	agentName = strings.TrimSpace(agentName)
 
-	agents, err := api.Database.GetWorkspaceAgentsByInstanceID(systemCtx, instanceID)
+	agents, err := api.Database.GetWorkspaceBuildAgentsByInstanceID(systemCtx, instanceID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Internal error fetching provisioner job agent.",
+			Message: "Internal error fetching workspace agent.",
 			Detail:  err.Error(),
 		})
 		return
 	}
-
-	// 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.
-	//
-	// We attach the provisioner job to each candidate during the filter
-	// loop so the post-selection code below can read it directly from the
-	// chosen candidate instead of re-querying. The previous code re-fetched
-	// the resource and job for the surviving agent, firing the
-	// resource->job->build->workspace dbauthz cascade twice and saturating
-	// the pgx pool under load.
-	type instanceCandidate struct {
-		agent database.WorkspaceAgent
-		job   database.ProvisionerJob
-	}
-	buildCandidates := make([]instanceCandidate, 0, len(agents))
-	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 {
-			buildCandidates = append(buildCandidates, instanceCandidate{
-				agent: candidate,
-				job:   job,
-			})
-		}
-	}
-	if len(buildCandidates) == 0 {
+	if len(agents) == 0 {
 		httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
 			Message: fmt.Sprintf("Instance with id %q not found.", instanceID),
 		})
 		return
 	}
 
-	var selected instanceCandidate
+	selected := agents[0]
 	if agentName != "" {
-		for _, candidate := range buildCandidates {
-			if candidate.agent.Name == agentName {
+		found := false
+		for _, candidate := range agents {
+			if candidate.WorkspaceAgent.Name == agentName {
 				selected = candidate
+				found = true
 				break
 			}
 		}
-		if selected.agent.ID == uuid.Nil {
+		if !found {
 			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(buildCandidates) != 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(buildCandidates))
-			for i, candidate := range buildCandidates {
-				names[i] = candidate.agent.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
+	} 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.WorkspaceAgent.Name
 		}
-		selected = buildCandidates[0]
-	}
-	agent := selected.agent
-	job := selected.job
-	var jobData provisionerdserver.WorkspaceProvisionJob
-	err = json.Unmarshal(job.Input, &jobData)
-	if err != nil {
-		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Internal error extracting job data.",
-			Detail:  err.Error(),
-		})
-		return
-	}
-	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.",
-			Detail:  err.Error(),
+		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 := selected.WorkspaceAgent
 	// 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.
-	latestHistory, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(systemCtx, resourceHistory.WorkspaceID)
+	latestHistory, err := api.Database.GetLatestWorkspaceBuildByWorkspaceID(systemCtx, selected.WorkspaceTable.ID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Internal error fetching the latest workspace build.",
@@ -257,7 +194,7 @@ func (api *API) handleAuthInstanceID(rw http.ResponseWriter, r *http.Request, in
 		})
 		return
 	}
-	if latestHistory.ID != resourceHistory.ID {
+	if latestHistory.ID != selected.WorkspaceBuildID {
 		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
 			Message: fmt.Sprintf("Resource found for id %q, but isn't registered on the latest history.", instanceID),
 		})

coderd/workspaceresourceauth_test.go

@@ -91,6 +91,50 @@ func TestPostWorkspaceAuthAWSInstanceIdentity(t *testing.T) {
 		require.NoError(t, err)
 	})
 
+	t.Run("RecycledInstanceID", func(t *testing.T) {
+		t.Parallel()
+
+		instanceID := newTestInstanceID(t)
+		certificates, metadataClient := coderdtest.NewAWSInstanceIdentity(t, instanceID)
+		setup := setupInstanceIDWorkspaceWithResources(t, &coderdtest.Options{
+			AWSCertificates: certificates,
+		}, workspaceAgentsForInstanceID(instanceID, "dev"))
+
+		successorVersion := coderdtest.CreateTemplateVersion(t, setup.client, setup.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: workspaceAgentsForInstanceID(newTestInstanceID(t), "dev"),
+						}},
+					},
+				},
+			}},
+		}, func(req *codersdk.CreateTemplateVersionRequest) {
+			req.TemplateID = setup.template.ID
+		})
+		coderdtest.AwaitTemplateVersionJobCompleted(t, setup.client, successorVersion.ID)
+		build := coderdtest.CreateWorkspaceBuild(t, setup.client, setup.workspace, database.WorkspaceTransitionStart, func(req *codersdk.CreateWorkspaceBuildRequest) {
+			req.TemplateVersionID = successorVersion.ID
+		})
+		coderdtest.AwaitWorkspaceBuildJobCompleted(t, setup.client, build.ID)
+
+		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+		defer cancel()
+
+		agentClient := agentsdk.New(setup.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.StatusBadRequest, apiErr.StatusCode())
+		require.Contains(t, apiErr.Message, "isn't registered on the latest history")
+	})
+
 	t.Run("Ambiguous/MultipleAgentsNoSelector", func(t *testing.T) {
 		t.Parallel()
 
@@ -126,35 +170,11 @@ func TestPostWorkspaceAuthAWSInstanceIdentity(t *testing.T) {
 		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)
+		res := postAWSInstanceIdentity(ctx, t, client, metadataClient, "")
 		defer res.Body.Close()
 
 		require.Equal(t, http.StatusConflict, res.StatusCode)
-		err = codersdk.ReadBodyAsError(res)
+		err := codersdk.ReadBodyAsError(res)
 		var apiErr *codersdk.Error
 		require.ErrorAs(t, err, &apiErr)
 		require.Equal(t, http.StatusConflict, apiErr.StatusCode())
@@ -174,35 +194,11 @@ func TestPostWorkspaceAuthAWSInstanceIdentity(t *testing.T) {
 		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)
+		res := postAWSInstanceIdentity(ctx, t, client, metadataClient, "   ")
 		defer res.Body.Close()
 
 		require.Equal(t, http.StatusConflict, res.StatusCode)
-		err = codersdk.ReadBodyAsError(res)
+		err := codersdk.ReadBodyAsError(res)
 		var apiErr *codersdk.Error
 		require.ErrorAs(t, err, &apiErr)
 		require.Equal(t, http.StatusConflict, apiErr.StatusCode())
@@ -368,9 +364,28 @@ func TestPostWorkspaceAuthGoogleInstanceIdentity(t *testing.T) {
 	})
 }
 
+type instanceIDWorkspaceSetup struct {
+	client    *codersdk.Client
+	store     database.Store
+	user      codersdk.CreateFirstUserResponse
+	template  codersdk.Template
+	workspace codersdk.Workspace
+}
+
 func setupInstanceIDWorkspace(t *testing.T, opts *coderdtest.Options, agents []*proto.Agent) (*codersdk.Client, database.Store) {
 	t.Helper()
 
+	setup := setupInstanceIDWorkspaceWithResources(t, opts, agents)
+	return setup.client, setup.store
+}
+
+func setupInstanceIDWorkspaceWithResources(
+	t *testing.T,
+	opts *coderdtest.Options,
+	agents []*proto.Agent,
+) instanceIDWorkspaceSetup {
+	t.Helper()
+
 	actualOpts := &coderdtest.Options{}
 	if opts != nil {
 		*actualOpts = *opts
@@ -398,7 +413,13 @@ func setupInstanceIDWorkspace(t *testing.T, opts *coderdtest.Options, agents []*
 	workspace := coderdtest.CreateWorkspace(t, client, template.ID)
 	coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
 
-	return client, store
+	return instanceIDWorkspaceSetup{
+		client:    client,
+		store:     store,
+		user:      user,
+		template:  template,
+		workspace: workspace,
+	}
 }
 
 func workspaceAgentsForInstanceID(instanceID string, names ...string) []*proto.Agent {
@@ -427,6 +448,59 @@ func requireWorkspaceAgentByInstanceIDAndName(t testing.TB, store database.Store
 	return database.WorkspaceAgent{}
 }
 
+const awsInstanceIdentityMetadataURL = "http://169.254.169.254/latest/dynamic/instance-identity"
+
+func postAWSInstanceIdentity(
+	ctx context.Context,
+	t testing.TB,
+	client *codersdk.Client,
+	metadataClient *http.Client,
+	agentName string,
+) *http.Response {
+	t.Helper()
+
+	signature := readAWSInstanceMetadata(ctx, t, metadataClient, "signature")
+	document := readAWSInstanceMetadata(ctx, t, metadataClient, "document")
+	reqBody, err := json.Marshal(map[string]string{
+		"signature":  signature,
+		"document":   document,
+		"agent_name": agentName,
+	})
+	require.NoError(t, err)
+
+	res, err := client.RequestWithoutSessionToken(
+		ctx,
+		http.MethodPost,
+		"/api/v2/workspaceagents/aws-instance-identity",
+		reqBody,
+	)
+	require.NoError(t, err)
+	return res
+}
+
+func readAWSInstanceMetadata(
+	ctx context.Context,
+	t testing.TB,
+	metadataClient *http.Client,
+	path string,
+) string {
+	t.Helper()
+
+	req, err := http.NewRequestWithContext(
+		ctx,
+		http.MethodGet,
+		awsInstanceIdentityMetadataURL+"/"+path,
+		nil,
+	)
+	require.NoError(t, err)
+	res, err := metadataClient.Do(req)
+	require.NoError(t, err)
+	defer res.Body.Close()
+	body, err := io.ReadAll(res.Body)
+	require.NoError(t, err)
+	return string(body)
+}
+
 func newTestInstanceID(t testing.TB) string {
 	t.Helper()
 	return fmt.Sprintf("instance-%d", time.Now().UnixNano())