fix(coderd): allow agent auth during workspace shutdown (#21538)

Agents were losing authentication during workspace shutdown, causing
shutdown scripts to fail. The auth query required agents to belong to
the latest build, but during shutdown a `stop` build becomes latest while
the `start` build's agents are still running.

Modified the auth query to allow `start` build agents to authenticate
temporarily during `stop` execution. The query allows auth when:

- Agent's `start` build job succeeded
- Latest build is `stop` with `pending`/`running` job status
- Builds are adjacent (`stop` is `build_number + 1`)
- Template versions match

Auth closes once `stop` completes.

Renamed `GetWorkspaceAgentAndLatestBuildByAuthToken` to
`GetAuthenticatedWorkspaceAgentAndBuildByAuthToken` since it returns the
agent's build (not always latest) during shutdown.

Closes coder/internal#1249
Fixes #19467

coderd/database/dbauthz/dbauthz.go

@@ -2244,6 +2244,14 @@ func (q *querier) GetAuditLogsOffset(ctx context.Context, arg database.GetAuditL
 	return q.db.GetAuthorizedAuditLogsOffset(ctx, arg, prep)
 }
 
+func (q *querier) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) {
+	// This is a system function
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
+		return database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow{}, err
+	}
+	return q.db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, authToken)
+}
+
 func (q *querier) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (database.GetAuthorizationUserRolesRow, error) {
 	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
 		return database.GetAuthorizationUserRolesRow{}, err
@@ -3606,14 +3614,6 @@ func (q *querier) GetWorkspaceACLByID(ctx context.Context, id uuid.UUID) (databa
 	return q.db.GetWorkspaceACLByID(ctx, id)
 }
 
-func (q *querier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) {
-	// This is a system function
-	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
-		return database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, err
-	}
-	return q.db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken)
-}
-
 func (q *querier) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) {
 	return fetch(q.log, q.auth, q.db.GetWorkspaceAgentAndWorkspaceByID)(ctx, id)
 }

coderd/database/dbauthz/dbauthz_test.go

@@ -3472,9 +3472,9 @@ func (s *MethodTestSuite) TestSystemFunctions() {
 		dbm.EXPECT().GetReplicaByID(gomock.Any(), id).Return(database.Replica{}, sql.ErrNoRows).AnyTimes()
 		check.Args(id).Asserts(rbac.ResourceSystem, policy.ActionRead).Errors(sql.ErrNoRows)
 	}))
-	s.Run("GetWorkspaceAgentAndLatestBuildByAuthToken", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
+	s.Run("GetAuthenticatedWorkspaceAgentAndBuildByAuthToken", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
 		tok := uuid.New()
-		dbm.EXPECT().GetWorkspaceAgentAndLatestBuildByAuthToken(gomock.Any(), tok).Return(database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow{}, sql.ErrNoRows).AnyTimes()
+		dbm.EXPECT().GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(gomock.Any(), tok).Return(database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow{}, sql.ErrNoRows).AnyTimes()
 		check.Args(tok).Asserts(rbac.ResourceSystem, policy.ActionRead).Errors(sql.ErrNoRows)
 	}))
 	s.Run("GetUserLinksByUserID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {

coderd/database/dbmetrics/querymetrics.go

@@ -919,6 +919,14 @@ func (m queryMetricsStore) GetAuditLogsOffset(ctx context.Context, arg database.
 	return r0, r1
 }
 
+func (m queryMetricsStore) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) {
+	start := time.Now()
+	r0, r1 := m.s.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, authToken)
+	m.queryLatencies.WithLabelValues("GetAuthenticatedWorkspaceAgentAndBuildByAuthToken").Observe(time.Since(start).Seconds())
+	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAuthenticatedWorkspaceAgentAndBuildByAuthToken").Inc()
+	return r0, r1
+}
+
 func (m queryMetricsStore) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (database.GetAuthorizationUserRolesRow, error) {
 	start := time.Now()
 	r0, r1 := m.s.GetAuthorizationUserRoles(ctx, userID)
@@ -2207,14 +2215,6 @@ func (m queryMetricsStore) GetWorkspaceACLByID(ctx context.Context, id uuid.UUID
 	return r0, r1
 }
 
-func (m queryMetricsStore) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) {
-	start := time.Now()
-	r0, r1 := m.s.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken)
-	m.queryLatencies.WithLabelValues("GetWorkspaceAgentAndLatestBuildByAuthToken").Observe(time.Since(start).Seconds())
-	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceAgentAndLatestBuildByAuthToken").Inc()
-	return r0, r1
-}
-
 func (m queryMetricsStore) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) {
 	start := time.Now()
 	r0, r1 := m.s.GetWorkspaceAgentAndWorkspaceByID(ctx, id)

coderd/database/dbmock/dbmock.go

@@ -1572,6 +1572,21 @@ func (mr *MockStoreMockRecorder) GetAuditLogsOffset(ctx, arg any) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAuditLogsOffset", reflect.TypeOf((*MockStore)(nil).GetAuditLogsOffset), ctx, arg)
 }
 
+// GetAuthenticatedWorkspaceAgentAndBuildByAuthToken mocks base method.
+func (m *MockStore) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetAuthenticatedWorkspaceAgentAndBuildByAuthToken", ctx, authToken)
+	ret0, _ := ret[0].(database.GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetAuthenticatedWorkspaceAgentAndBuildByAuthToken indicates an expected call of GetAuthenticatedWorkspaceAgentAndBuildByAuthToken.
+func (mr *MockStoreMockRecorder) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, authToken any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAuthenticatedWorkspaceAgentAndBuildByAuthToken", reflect.TypeOf((*MockStore)(nil).GetAuthenticatedWorkspaceAgentAndBuildByAuthToken), ctx, authToken)
+}
+
 // GetAuthorizationUserRoles mocks base method.
 func (m *MockStore) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (database.GetAuthorizationUserRolesRow, error) {
 	m.ctrl.T.Helper()
@@ -4122,21 +4137,6 @@ func (mr *MockStoreMockRecorder) GetWorkspaceACLByID(ctx, id any) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceACLByID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceACLByID), ctx, id)
 }
 
-// GetWorkspaceAgentAndLatestBuildByAuthToken mocks base method.
-func (m *MockStore) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetWorkspaceAgentAndLatestBuildByAuthToken", ctx, authToken)
-	ret0, _ := ret[0].(database.GetWorkspaceAgentAndLatestBuildByAuthTokenRow)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// GetWorkspaceAgentAndLatestBuildByAuthToken indicates an expected call of GetWorkspaceAgentAndLatestBuildByAuthToken.
-func (mr *MockStoreMockRecorder) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, authToken any) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentAndLatestBuildByAuthToken", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentAndLatestBuildByAuthToken), ctx, authToken)
-}
-
 // GetWorkspaceAgentAndWorkspaceByID mocks base method.
 func (m *MockStore) GetWorkspaceAgentAndWorkspaceByID(ctx context.Context, id uuid.UUID) (database.GetWorkspaceAgentAndWorkspaceByIDRow, error) {
 	m.ctrl.T.Helper()

coderd/database/querier.go

@@ -190,6 +190,12 @@ type sqlcQuerier interface {
 	// GetAuditLogsBefore retrieves `row_limit` number of audit logs before the provided
 	// ID.
 	GetAuditLogsOffset(ctx context.Context, arg GetAuditLogsOffsetParams) ([]GetAuditLogsOffsetRow, error)
+	// GetAuthenticatedWorkspaceAgentAndBuildByAuthToken returns an authenticated
+	// workspace agent and its associated build. During normal operation, this is
+	// the latest build. During shutdown, this may be the previous START build while
+	// the STOP build is executing, allowing shutdown scripts to authenticate (see
+	// issue #19467).
+	GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error)
 	// This function returns roles for authorization purposes. Implied member roles
 	// are included.
 	GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (GetAuthorizationUserRolesRow, error)
@@ -466,7 +472,6 @@ type sqlcQuerier interface {
 	GetWebpushSubscriptionsByUserID(ctx context.Context, userID uuid.UUID) ([]WebpushSubscription, error)
 	GetWebpushVAPIDKeys(ctx context.Context) (GetWebpushVAPIDKeysRow, error)
 	GetWorkspaceACLByID(ctx context.Context, id uuid.UUID) (GetWorkspaceACLByIDRow, error)
-	GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndLatestBuildByAuthTokenRow, 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)

coderd/database/querier_test.go

@@ -8077,3 +8077,408 @@ func TestDeleteExpiredAPIKeys(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, remaining, len(unexpiredTimes))
 }
+
+func TestGetAuthenticatedWorkspaceAgentAndBuildByAuthToken_ShutdownScripts(t *testing.T) {
+	t.Parallel()
+	if testing.Short() {
+		t.SkipNow()
+	}
+
+	sqlDB := testSQLDB(t)
+	err := migrations.Up(sqlDB)
+	require.NoError(t, err)
+	db := database.New(sqlDB)
+
+	org := dbgen.Organization(t, db, database.Organization{})
+	owner := dbgen.User(t, db, database.User{})
+	tpl := dbgen.Template(t, db, database.Template{
+		OrganizationID: org.ID,
+		CreatedBy:      owner.ID,
+	})
+	ver := dbgen.TemplateVersion(t, db, database.TemplateVersion{
+		TemplateID: uuid.NullUUID{
+			UUID:  tpl.ID,
+			Valid: true,
+		},
+		OrganizationID: tpl.OrganizationID,
+		CreatedBy:      owner.ID,
+	})
+
+	t.Run("DuringStopBuild", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+			OwnerID:        owner.ID,
+			OrganizationID: org.ID,
+			TemplateID:     tpl.ID,
+		})
+
+		// Create start build with succeeded job (already completed).
+		startJob := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob)
+		startJob = dbgen.ProvisionerJob(t, db, nil, startJob)
+		startResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID:      startJob.ID,
+			Transition: database.WorkspaceTransitionStart,
+		})
+		startBuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       1,
+			Transition:        database.WorkspaceTransitionStart,
+			InitiatorID:       owner.ID,
+			JobID:             startJob.ID,
+		})
+		agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: startResource.ID,
+		})
+
+		// Create stop build (becomes latest).
+		stopJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+			JobStatus:      database.ProvisionerJobStatusRunning,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       2,
+			Transition:        database.WorkspaceTransitionStop,
+			InitiatorID:       owner.ID,
+			JobID:             stopJob.ID,
+		})
+
+		// Agent should still authenticate during stop build execution.
+		row, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent.AuthToken)
+		require.NoError(t, err, "agent should authenticate during stop build execution")
+		require.Equal(t, agent.ID, row.WorkspaceAgent.ID)
+		require.Equal(t, startBuild.ID, row.WorkspaceBuild.ID, "should return start build, not stop build")
+	})
+
+	t.Run("AfterStopJobCompletes", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+			OwnerID:        owner.ID,
+			OrganizationID: org.ID,
+			TemplateID:     tpl.ID,
+		})
+
+		// Create start build with completed job.
+		startJob := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob)
+		startJob = dbgen.ProvisionerJob(t, db, nil, startJob)
+
+		startResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID:      startJob.ID,
+			Transition: database.WorkspaceTransitionStart,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       1,
+			Transition:        database.WorkspaceTransitionStart,
+			InitiatorID:       owner.ID,
+			JobID:             startJob.ID,
+		})
+		agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: startResource.ID,
+		})
+
+		// Create stop build (becomes latest) with completed job.
+		stopJob := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusSucceeded, &stopJob)
+		stopJob = dbgen.ProvisionerJob(t, db, nil, stopJob)
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       2,
+			Transition:        database.WorkspaceTransitionStop,
+			InitiatorID:       owner.ID,
+			JobID:             stopJob.ID,
+		})
+
+		// Agent should NOT authenticate after stop job completes.
+		_, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent.AuthToken)
+		require.ErrorIs(t, err, sql.ErrNoRows, "agent should not authenticate after stop job completes")
+	})
+
+	t.Run("FailedStartBuild", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+			OwnerID:        owner.ID,
+			OrganizationID: org.ID,
+			TemplateID:     tpl.ID,
+		})
+
+		// Create START build with FAILED job.
+		startJob := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusFailed, &startJob)
+		startJob = dbgen.ProvisionerJob(t, db, nil, startJob)
+		startResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID:      startJob.ID,
+			Transition: database.WorkspaceTransitionStart,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       1,
+			Transition:        database.WorkspaceTransitionStart,
+			InitiatorID:       owner.ID,
+			JobID:             startJob.ID,
+		})
+		agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: startResource.ID,
+		})
+
+		// Create STOP build with running job.
+		stopJob := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+			JobStatus:      database.ProvisionerJobStatusRunning,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       2,
+			Transition:        database.WorkspaceTransitionStop,
+			InitiatorID:       owner.ID,
+			JobID:             stopJob.ID,
+		})
+
+		// Agent should NOT authenticate (start build failed).
+		_, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent.AuthToken)
+		require.ErrorIs(t, err, sql.ErrNoRows, "agent from failed start build should not authenticate")
+	})
+
+	t.Run("PendingStopBuild", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+			OwnerID:        owner.ID,
+			OrganizationID: org.ID,
+			TemplateID:     tpl.ID,
+		})
+
+		// Create start build with succeeded job.
+		startJob := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob)
+		startJob = dbgen.ProvisionerJob(t, db, nil, startJob)
+		startResource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID:      startJob.ID,
+			Transition: database.WorkspaceTransitionStart,
+		})
+		startBuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       1,
+			Transition:        database.WorkspaceTransitionStart,
+			InitiatorID:       owner.ID,
+			JobID:             startJob.ID,
+		})
+		agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: startResource.ID,
+		})
+
+		// Create stop build with pending job (not started yet).
+		stopJob := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusPending, &stopJob)
+		stopJob = dbgen.ProvisionerJob(t, db, nil, stopJob)
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       2,
+			Transition:        database.WorkspaceTransitionStop,
+			InitiatorID:       owner.ID,
+			JobID:             stopJob.ID,
+		})
+
+		// Agent should authenticate during pending stop build.
+		row, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent.AuthToken)
+		require.NoError(t, err, "agent should authenticate during pending stop build")
+		require.Equal(t, agent.ID, row.WorkspaceAgent.ID)
+		require.Equal(t, startBuild.ID, row.WorkspaceBuild.ID, "should return start build")
+	})
+
+	t.Run("MultipleStartStopCycles", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+			OwnerID:        owner.ID,
+			OrganizationID: org.ID,
+			TemplateID:     tpl.ID,
+		})
+
+		// Build 1: START (succeeded).
+		startJob1 := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob1)
+		startJob1 = dbgen.ProvisionerJob(t, db, nil, startJob1)
+		startResource1 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID:      startJob1.ID,
+			Transition: database.WorkspaceTransitionStart,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       1,
+			Transition:        database.WorkspaceTransitionStart,
+			InitiatorID:       owner.ID,
+			JobID:             startJob1.ID,
+		})
+		agent1 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: startResource1.ID,
+		})
+
+		// Build 2: STOP (succeeded).
+		stopJob1 := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusSucceeded, &stopJob1)
+		stopJob1 = dbgen.ProvisionerJob(t, db, nil, stopJob1)
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       2,
+			Transition:        database.WorkspaceTransitionStop,
+			InitiatorID:       owner.ID,
+			JobID:             stopJob1.ID,
+		})
+
+		// Build 3: START (succeeded).
+		startJob2 := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob2)
+		startJob2 = dbgen.ProvisionerJob(t, db, nil, startJob2)
+		startResource2 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID:      startJob2.ID,
+			Transition: database.WorkspaceTransitionStart,
+		})
+		startBuild2 := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       3,
+			Transition:        database.WorkspaceTransitionStart,
+			InitiatorID:       owner.ID,
+			JobID:             startJob2.ID,
+		})
+		agent2 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: startResource2.ID,
+		})
+
+		// Build 4: STOP (running).
+		stopJob2 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+			JobStatus:      database.ProvisionerJobStatusRunning,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       4,
+			Transition:        database.WorkspaceTransitionStop,
+			InitiatorID:       owner.ID,
+			JobID:             stopJob2.ID,
+		})
+
+		// Agent from build 3 should authenticate.
+		row, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent2.AuthToken)
+		require.NoError(t, err, "agent from most recent start should authenticate during stop")
+		require.Equal(t, agent2.ID, row.WorkspaceAgent.ID)
+		require.Equal(t, startBuild2.ID, row.WorkspaceBuild.ID)
+
+		// Agent from build 1 should NOT authenticate.
+		_, err = db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent1.AuthToken)
+		require.ErrorIs(t, err, sql.ErrNoRows, "agent from old cycle should not authenticate")
+	})
+
+	t.Run("WrongTransitionType", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+			OwnerID:        owner.ID,
+			OrganizationID: org.ID,
+			TemplateID:     tpl.ID,
+		})
+
+		// Create first start build.
+		startJob1 := database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+		}
+		setJobStatus(t, database.ProvisionerJobStatusSucceeded, &startJob1)
+		startJob1 = dbgen.ProvisionerJob(t, db, nil, startJob1)
+		startResource1 := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID:      startJob1.ID,
+			Transition: database.WorkspaceTransitionStart,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       1,
+			Transition:        database.WorkspaceTransitionStart,
+			InitiatorID:       owner.ID,
+			JobID:             startJob1.ID,
+		})
+		agent1 := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: startResource1.ID,
+		})
+
+		// Create another START build as latest (not STOP).
+		startJob2 := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
+			Type:           database.ProvisionerJobTypeWorkspaceBuild,
+			InitiatorID:    owner.ID,
+			OrganizationID: org.ID,
+			JobStatus:      database.ProvisionerJobStatusRunning,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			TemplateVersionID: ver.ID,
+			BuildNumber:       2,
+			Transition:        database.WorkspaceTransitionStart,
+			InitiatorID:       owner.ID,
+			JobID:             startJob2.ID,
+		})
+
+		// Agent from build 1 should NOT authenticate (latest is not STOP).
+		_, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agent1.AuthToken)
+		require.ErrorIs(t, err, sql.ErrNoRows, "agent should not authenticate when latest build is not STOP")
+	})
+}

coderd/database/queries.sql.go

@@ -18024,7 +18024,7 @@ func (q *sqlQuerier) DeleteWorkspaceSubAgentByID(ctx context.Context, id uuid.UU
 	return err
 }
 
-const getWorkspaceAgentAndLatestBuildByAuthToken = `-- name: GetWorkspaceAgentAndLatestBuildByAuthToken :one
+const getAuthenticatedWorkspaceAgentAndBuildByAuthToken = `-- name: GetAuthenticatedWorkspaceAgentAndBuildByAuthToken :one
 SELECT
 	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,
 	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,
@@ -18054,29 +18054,57 @@ WHERE
 	AND workspaces.deleted = FALSE
 	-- Filter out deleted sub agents.
 	AND workspace_agents.deleted = FALSE
-	-- Filter out builds that are not the latest.
-	AND workspace_build_with_user.build_number = (
-		-- Select from workspace_builds as it's one less join compared
-		-- to workspace_build_with_user.
-		SELECT
-			MAX(build_number)
-		FROM
-			workspace_builds
-		WHERE
-			workspace_id = workspace_build_with_user.workspace_id
-	)
+	-- Filter out builds that are not the latest, with exception for shutdown case.
+	-- Use CASE for short-circuiting: check normal case first (most common), then shutdown case.
+	AND CASE
+		-- Normal case: Agent's build is the latest build.
+		WHEN workspace_build_with_user.build_number = (
+			SELECT
+				MAX(build_number)
+			FROM
+				workspace_builds
+			WHERE
+				workspace_id = workspace_build_with_user.workspace_id
+		) THEN TRUE
+		-- Shutdown case: Agent from previous START build during STOP build execution.
+		WHEN workspace_build_with_user.transition = 'start'
+			-- Agent's START build job succeeded.
+			AND (SELECT job_status FROM provisioner_jobs WHERE id = workspace_build_with_user.job_id) = 'succeeded'
+			-- Latest build is a STOP build whose job is still active,
+			-- and agent's build is immediately previous.
+			AND EXISTS (
+				SELECT 1
+				FROM workspace_builds latest
+				JOIN provisioner_jobs pj ON pj.id = latest.job_id
+				WHERE latest.workspace_id = workspace_build_with_user.workspace_id
+				AND latest.build_number = workspace_build_with_user.build_number + 1
+				AND latest.build_number = (
+					SELECT MAX(build_number)
+					FROM workspace_builds l2
+					WHERE l2.workspace_id = latest.workspace_id
+				)
+				AND latest.transition = 'stop'
+				AND pj.job_status IN ('pending', 'running')
+			) THEN TRUE
+		ELSE FALSE
+	END
 `
 
-type GetWorkspaceAgentAndLatestBuildByAuthTokenRow struct {
+type GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow struct {
 	WorkspaceTable WorkspaceTable `db:"workspace_table" json:"workspace_table"`
 	WorkspaceAgent WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"`
 	WorkspaceBuild WorkspaceBuild `db:"workspace_build" json:"workspace_build"`
 	TaskID         uuid.NullUUID  `db:"task_id" json:"task_id"`
 }
 
-func (q *sqlQuerier) GetWorkspaceAgentAndLatestBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetWorkspaceAgentAndLatestBuildByAuthTokenRow, error) {
-	row := q.db.QueryRowContext(ctx, getWorkspaceAgentAndLatestBuildByAuthToken, authToken)
-	var i GetWorkspaceAgentAndLatestBuildByAuthTokenRow
+// GetAuthenticatedWorkspaceAgentAndBuildByAuthToken returns an authenticated
+// workspace agent and its associated build. During normal operation, this is
+// the latest build. During shutdown, this may be the previous START build while
+// the STOP build is executing, allowing shutdown scripts to authenticate (see
+// issue #19467).
+func (q *sqlQuerier) GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx context.Context, authToken uuid.UUID) (GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow, error) {
+	row := q.db.QueryRowContext(ctx, getAuthenticatedWorkspaceAgentAndBuildByAuthToken, authToken)
+	var i GetAuthenticatedWorkspaceAgentAndBuildByAuthTokenRow
 	err := row.Scan(
 		&i.WorkspaceTable.ID,
 		&i.WorkspaceTable.CreatedAt,

coderd/database/queries/workspaceagents.sql

@@ -281,7 +281,12 @@ WHERE
 	-- Filter out deleted sub agents.
 	AND workspace_agents.deleted = FALSE;
 
--- name: GetWorkspaceAgentAndLatestBuildByAuthToken :one
+-- GetAuthenticatedWorkspaceAgentAndBuildByAuthToken returns an authenticated
+-- workspace agent and its associated build. During normal operation, this is
+-- the latest build. During shutdown, this may be the previous START build while
+-- the STOP build is executing, allowing shutdown scripts to authenticate (see
+-- issue #19467).
+-- name: GetAuthenticatedWorkspaceAgentAndBuildByAuthToken :one
 SELECT
 	sqlc.embed(workspaces),
 	sqlc.embed(workspace_agents),
@@ -311,17 +316,40 @@ WHERE
 	AND workspaces.deleted = FALSE
 	-- Filter out deleted sub agents.
 	AND workspace_agents.deleted = FALSE
-	-- Filter out builds that are not the latest.
-	AND workspace_build_with_user.build_number = (
-		-- Select from workspace_builds as it's one less join compared
-		-- to workspace_build_with_user.
-		SELECT
-			MAX(build_number)
-		FROM
-			workspace_builds
-		WHERE
-			workspace_id = workspace_build_with_user.workspace_id
-	)
+	-- Filter out builds that are not the latest, with exception for shutdown case.
+	-- Use CASE for short-circuiting: check normal case first (most common), then shutdown case.
+	AND CASE
+		-- Normal case: Agent's build is the latest build.
+		WHEN workspace_build_with_user.build_number = (
+			SELECT
+				MAX(build_number)
+			FROM
+				workspace_builds
+			WHERE
+				workspace_id = workspace_build_with_user.workspace_id
+		) THEN TRUE
+		-- Shutdown case: Agent from previous START build during STOP build execution.
+		WHEN workspace_build_with_user.transition = 'start'
+			-- Agent's START build job succeeded.
+			AND (SELECT job_status FROM provisioner_jobs WHERE id = workspace_build_with_user.job_id) = 'succeeded'
+			-- Latest build is a STOP build whose job is still active,
+			-- and agent's build is immediately previous.
+			AND EXISTS (
+				SELECT 1
+				FROM workspace_builds latest
+				JOIN provisioner_jobs pj ON pj.id = latest.job_id
+				WHERE latest.workspace_id = workspace_build_with_user.workspace_id
+				AND latest.build_number = workspace_build_with_user.build_number + 1
+				AND latest.build_number = (
+					SELECT MAX(build_number)
+					FROM workspace_builds l2
+					WHERE l2.workspace_id = latest.workspace_id
+				)
+				AND latest.transition = 'stop'
+				AND pj.job_status IN ('pending', 'running')
+			) THEN TRUE
+		ELSE FALSE
+	END
 ;
 
 -- name: InsertWorkspaceAgentScriptTimings :one

coderd/httpmw/workspaceagent.go

@@ -92,7 +92,7 @@ func ExtractWorkspaceAgentAndLatestBuild(opts ExtractWorkspaceAgentAndLatestBuil
 			}
 
 			//nolint:gocritic // System needs to be able to get workspace agents.
-			row, err := opts.DB.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), token)
+			row, err := opts.DB.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), token)
 			if err != nil {
 				if errors.Is(err, sql.ErrNoRows) {
 					optionalWrite(http.StatusUnauthorized, codersdk.Response{

coderd/httpmw/workspaceagent_test.go

@@ -1,9 +1,11 @@
 package httpmw_test
 
 import (
+	"database/sql"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
@@ -12,6 +14,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbgen"
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/database/dbtime"
 	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/codersdk"
 )
@@ -95,6 +98,180 @@ func TestWorkspaceAgent(t *testing.T) {
 		t.Cleanup(func() { _ = res.Body.Close() })
 		require.Equal(t, http.StatusUnauthorized, res.StatusCode)
 	})
+
+	t.Run("DuringShutdown", func(t *testing.T) {
+		t.Parallel()
+		db, ps := dbtestutil.NewDB(t)
+		authToken := uuid.New()
+		req, rtr, ws, tpv := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild(
+			httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
+				DB:       db,
+				Optional: false,
+			}),
+		)
+
+		// Create a STOP build with running job (becomes latest).
+		stopJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+			OrganizationID: ws.OrganizationID,
+			JobStatus:      database.ProvisionerJobStatusRunning,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       ws.ID,
+			JobID:             stopJob.ID,
+			TemplateVersionID: tpv.ID,
+			BuildNumber:       2,
+			Transition:        database.WorkspaceTransitionStop,
+		})
+
+		// Agent should still authenticate during shutdown.
+		rw := httptest.NewRecorder()
+		req.Header.Set(codersdk.SessionTokenHeader, authToken.String())
+		rtr.ServeHTTP(rw, req)
+
+		//nolint:bodyclose // Closed in `t.Cleanup`
+		res := rw.Result()
+		t.Cleanup(func() { _ = res.Body.Close() })
+		require.Equal(t, http.StatusOK, res.StatusCode, "agent should authenticate during stop build execution")
+	})
+
+	t.Run("AfterShutdownCompletes", func(t *testing.T) {
+		t.Parallel()
+		db, ps := dbtestutil.NewDB(t)
+		authToken := uuid.New()
+		req, rtr, ws, tpv := setup(t, db, authToken, httpmw.ExtractWorkspaceAgentAndLatestBuild(
+			httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
+				DB:       db,
+				Optional: false,
+			}),
+		)
+
+		// Create a STOP build with completed job (becomes latest).
+		stopJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+			OrganizationID: ws.OrganizationID,
+			JobStatus:      database.ProvisionerJobStatusSucceeded,
+			CompletedAt: sql.NullTime{
+				Time:  dbtime.Now(),
+				Valid: true,
+			},
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       ws.ID,
+			JobID:             stopJob.ID,
+			TemplateVersionID: tpv.ID,
+			BuildNumber:       2,
+			Transition:        database.WorkspaceTransitionStop,
+		})
+
+		// Agent should NOT authenticate after stop job completes.
+		rw := httptest.NewRecorder()
+		req.Header.Set(codersdk.SessionTokenHeader, authToken.String())
+		rtr.ServeHTTP(rw, req)
+
+		//nolint:bodyclose // Closed in `t.Cleanup`
+		res := rw.Result()
+		t.Cleanup(func() { _ = res.Body.Close() })
+		require.Equal(t, http.StatusUnauthorized, res.StatusCode, "agent should not authenticate after stop job completes")
+	})
+
+	t.Run("FailedStartBuild", func(t *testing.T) {
+		t.Parallel()
+		db, ps := dbtestutil.NewDB(t)
+		authToken := uuid.New()
+
+		org := dbgen.Organization(t, db, database.Organization{})
+		user := dbgen.User(t, db, database.User{
+			Status: database.UserStatusActive,
+		})
+		_ = dbgen.OrganizationMember(t, db, database.OrganizationMember{
+			UserID:         user.ID,
+			OrganizationID: org.ID,
+		})
+		templateVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
+			OrganizationID: org.ID,
+			CreatedBy:      user.ID,
+		})
+		template := dbgen.Template(t, db, database.Template{
+			OrganizationID:  org.ID,
+			ActiveVersionID: templateVersion.ID,
+			CreatedBy:       user.ID,
+		})
+		workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
+			OwnerID:        user.ID,
+			OrganizationID: org.ID,
+			TemplateID:     template.ID,
+		})
+
+		// Create START build with FAILED job status.
+		startJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+			OrganizationID: org.ID,
+			JobStatus:      database.ProvisionerJobStatusFailed,
+			StartedAt: sql.NullTime{
+				Time:  dbtime.Now().Add(-time.Minute),
+				Valid: true,
+			},
+			CompletedAt: sql.NullTime{
+				Time:  dbtime.Now(),
+				Valid: true,
+			},
+			Error: sql.NullString{
+				String: "build failed",
+				Valid:  true,
+			},
+			ErrorCode: sql.NullString{
+				String: "FAILED",
+				Valid:  true,
+			},
+		})
+		resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
+			JobID: startJob.ID,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			JobID:             startJob.ID,
+			TemplateVersionID: templateVersion.ID,
+			BuildNumber:       1,
+			Transition:        database.WorkspaceTransitionStart,
+		})
+		_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
+			ResourceID: resource.ID,
+			AuthToken:  authToken,
+		})
+
+		// Create a STOP build with running job.
+		stopJob := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+			OrganizationID: org.ID,
+			JobStatus:      database.ProvisionerJobStatusRunning,
+		})
+		_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+			WorkspaceID:       workspace.ID,
+			JobID:             stopJob.ID,
+			TemplateVersionID: templateVersion.ID,
+			BuildNumber:       2,
+			Transition:        database.WorkspaceTransitionStop,
+		})
+
+		req := httptest.NewRequest("GET", "/", nil)
+		rtr := chi.NewRouter()
+		rtr.Use(httpmw.ExtractWorkspaceAgentAndLatestBuild(
+			httpmw.ExtractWorkspaceAgentAndLatestBuildConfig{
+				DB:       db,
+				Optional: false,
+			}))
+		rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
+			_ = httpmw.WorkspaceAgent(r)
+			rw.WriteHeader(http.StatusOK)
+		})
+
+		// Agent should NOT authenticate (start build failed).
+		rw := httptest.NewRecorder()
+		req.Header.Set(codersdk.SessionTokenHeader, authToken.String())
+		rtr.ServeHTTP(rw, req)
+
+		//nolint:bodyclose // Closed in `t.Cleanup`
+		res := rw.Result()
+		t.Cleanup(func() { _ = res.Body.Close() })
+		require.Equal(t, http.StatusUnauthorized, res.StatusCode, "agent should not authenticate when start build failed")
+	})
 }
 
 func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Handler) http.Handler) (*http.Request, http.Handler, database.WorkspaceTable, database.TemplateVersion) {
@@ -123,6 +300,15 @@ func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Ha
 	})
 	job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
 		OrganizationID: org.ID,
+		JobStatus:      database.ProvisionerJobStatusSucceeded,
+		StartedAt: sql.NullTime{
+			Time:  dbtime.Now().Add(-30 * time.Second),
+			Valid: true,
+		},
+		CompletedAt: sql.NullTime{
+			Time:  dbtime.Now(),
+			Valid: true,
+		},
 	})
 	resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
 		JobID: job.ID,
@@ -131,6 +317,8 @@ func setup(t testing.TB, db database.Store, authToken uuid.UUID, mw func(http.Ha
 		WorkspaceID:       workspace.ID,
 		JobID:             job.ID,
 		TemplateVersionID: templateVersion.ID,
+		BuildNumber:       1,
+		Transition:        database.WorkspaceTransitionStart,
 	})
 	_ = dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
 		ResourceID: resource.ID,

coderd/workspaceagents_test.go

@@ -623,7 +623,7 @@ func TestWorkspaceAgentClientCoordinate_BadVersion(t *testing.T) {
 	ctx := testutil.Context(t, testutil.WaitShort)
 	agentToken, err := uuid.Parse(r.AgentToken)
 	require.NoError(t, err)
-	ao, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentToken)
+	ao, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentToken)
 	require.NoError(t, err)
 
 	//nolint: bodyclose // closed by ReadBodyAsError
@@ -713,7 +713,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) {
 		agentTokenUUID, err := uuid.Parse(r.AgentToken)
 		require.NoError(t, err)
 		ctx := testutil.Context(t, testutil.WaitLong)
-		agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID)
+		agentAndBuild, err := api.Database.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID)
 		require.NoError(t, err)
 
 		// Connect with no resume token, and ensure that the peer ID is set to a
@@ -785,7 +785,7 @@ func TestWorkspaceAgentClientCoordinate_ResumeToken(t *testing.T) {
 		agentTokenUUID, err := uuid.Parse(r.AgentToken)
 		require.NoError(t, err)
 		ctx := testutil.Context(t, testutil.WaitLong)
-		agentAndBuild, err := api.Database.GetWorkspaceAgentAndLatestBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID)
+		agentAndBuild, err := api.Database.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(dbauthz.AsSystemRestricted(ctx), agentTokenUUID)
 		require.NoError(t, err)
 
 		// Connect with no resume token, and ensure that the peer ID is set to a

enterprise/cli/prebuilds_test.go

@@ -435,7 +435,7 @@ func TestSchedulePrebuilds(t *testing.T) {
 
 			// Mark the prebuilt workspace's agent as ready so the prebuild can be claimed
 			ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong))
-			agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken))
+			agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken))
 			require.NoError(t, err)
 			err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{
 				ID:             agent.WorkspaceAgent.ID,

enterprise/coderd/schedule/template_test.go

@@ -1363,7 +1363,7 @@ func TestTemplateUpdatePrebuilds(t *testing.T) {
 			// Mark the prebuilt workspace's agent as ready so the prebuild can be claimed
 			// nolint:gocritic
 			agentCtx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong))
-			agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(agentCtx, uuid.MustParse(workspaceBuild.AgentToken))
+			agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(agentCtx, uuid.MustParse(workspaceBuild.AgentToken))
 			require.NoError(t, err)
 			err = db.UpdateWorkspaceAgentLifecycleStateByID(agentCtx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{
 				ID:             agent.WorkspaceAgent.ID,

enterprise/coderd/workspaceagents_test.go

@@ -187,7 +187,7 @@ func TestReinitializeAgent(t *testing.T) {
 			ctx := testutil.Context(t, testutil.WaitLong)
 			var prebuildID uuid.UUID
 			require.Eventually(t, func() bool {
-				agentAndBuild, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, agentToken)
+				agentAndBuild, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, agentToken)
 				if err != nil {
 					return false
 				}

enterprise/coderd/workspaces_test.go

@@ -567,7 +567,7 @@ func TestCreateUserWorkspace(t *testing.T) {
 		}).Do()
 
 		ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong))
-		agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(r.AgentToken))
+		agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, uuid.MustParse(r.AgentToken))
 		require.NoError(t, err)
 
 		err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{
@@ -2788,7 +2788,7 @@ func TestPrebuildUpdateLifecycleParams(t *testing.T) {
 
 			// Mark the prebuilt workspace's agent as ready so the prebuild can be claimed
 			ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong))
-			agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken))
+			agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, uuid.MustParse(workspaceBuild.AgentToken))
 			require.NoError(t, err)
 			err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{
 				ID:             agent.WorkspaceAgent.ID,
@@ -2887,7 +2887,7 @@ func TestPrebuildActivityBump(t *testing.T) {
 	// Mark the prebuilt workspace's agent as ready so the prebuild can be claimed
 	// nolint:gocritic
 	ctx := dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong))
-	agent, err := db.GetWorkspaceAgentAndLatestBuildByAuthToken(ctx, uuid.MustParse(wb.AgentToken))
+	agent, err := db.GetAuthenticatedWorkspaceAgentAndBuildByAuthToken(ctx, uuid.MustParse(wb.AgentToken))
 	require.NoError(t, err)
 	err = db.UpdateWorkspaceAgentLifecycleStateByID(ctx, database.UpdateWorkspaceAgentLifecycleStateByIDParams{
 		ID:             agent.WorkspaceAgent.ID,