feat: add managed agent license limit checks (#18937)

- Adds a query for counting managed agent workspace builds between two
timestamps
- The "Actual" field in the feature entitlement for managed agents is
now populated with the value read from the database
- The wsbuilder package now validates AI agent usage against the limit
when a license is installed

Closes coder/internal#777

coderd/autobuild/lifecycle_executor.go

@@ -42,6 +42,7 @@ type Executor struct {
 	templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
 	accessControlStore    *atomic.Pointer[dbauthz.AccessControlStore]
 	auditor               *atomic.Pointer[audit.Auditor]
+	buildUsageChecker     *atomic.Pointer[wsbuilder.UsageChecker]
 	log                   slog.Logger
 	tick                  <-chan time.Time
 	statsCh               chan<- Stats
@@ -65,7 +66,7 @@ type Stats struct {
 }
 
 // New returns a new wsactions executor.
-func NewExecutor(ctx context.Context, db database.Store, ps pubsub.Pubsub, fc *files.Cache, reg prometheus.Registerer, tss *atomic.Pointer[schedule.TemplateScheduleStore], auditor *atomic.Pointer[audit.Auditor], acs *atomic.Pointer[dbauthz.AccessControlStore], log slog.Logger, tick <-chan time.Time, enqueuer notifications.Enqueuer, exp codersdk.Experiments) *Executor {
+func NewExecutor(ctx context.Context, db database.Store, ps pubsub.Pubsub, fc *files.Cache, reg prometheus.Registerer, tss *atomic.Pointer[schedule.TemplateScheduleStore], auditor *atomic.Pointer[audit.Auditor], acs *atomic.Pointer[dbauthz.AccessControlStore], buildUsageChecker *atomic.Pointer[wsbuilder.UsageChecker], log slog.Logger, tick <-chan time.Time, enqueuer notifications.Enqueuer, exp codersdk.Experiments) *Executor {
 	factory := promauto.With(reg)
 	le := &Executor{
 		//nolint:gocritic // Autostart has a limited set of permissions.
@@ -78,6 +79,7 @@ func NewExecutor(ctx context.Context, db database.Store, ps pubsub.Pubsub, fc *f
 		log:                   log.Named("autobuild"),
 		auditor:               auditor,
 		accessControlStore:    acs,
+		buildUsageChecker:     buildUsageChecker,
 		notificationsEnqueuer: enqueuer,
 		reg:                   reg,
 		experiments:           exp,
@@ -279,7 +281,7 @@ func (e *Executor) runOnce(t time.Time) Stats {
 					}
 
 					if nextTransition != "" {
-						builder := wsbuilder.New(ws, nextTransition).
+						builder := wsbuilder.New(ws, nextTransition, *e.buildUsageChecker.Load()).
 							SetLastWorkspaceBuildInTx(&latestBuild).
 							SetLastWorkspaceBuildJobInTx(&latestJob).
 							Experiments(e.experiments).

coderd/coderd.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/coder/coder/v2/coderd/oauth2provider"
 	"github.com/coder/coder/v2/coderd/prebuilds"
+	"github.com/coder/coder/v2/coderd/wsbuilder"
 
 	"github.com/andybalholm/brotli"
 	"github.com/go-chi/chi/v5"
@@ -559,6 +560,13 @@ func New(options *Options) *API {
 	// bugs that may only occur when a key isn't precached in tests and the latency cost is minimal.
 	cryptokeys.StartRotator(ctx, options.Logger, options.Database)
 
+	// AGPL uses a no-op build usage checker as there are no license
+	// entitlements to enforce. This is swapped out in
+	// enterprise/coderd/coderd.go.
+	var buildUsageChecker atomic.Pointer[wsbuilder.UsageChecker]
+	var noopUsageChecker wsbuilder.UsageChecker = wsbuilder.NoopUsageChecker{}
+	buildUsageChecker.Store(&noopUsageChecker)
+
 	api := &API{
 		ctx:          ctx,
 		cancel:       cancel,
@@ -579,6 +587,7 @@ func New(options *Options) *API {
 		TemplateScheduleStore:       options.TemplateScheduleStore,
 		UserQuietHoursScheduleStore: options.UserQuietHoursScheduleStore,
 		AccessControlStore:          options.AccessControlStore,
+		BuildUsageChecker:           &buildUsageChecker,
 		FileCache:                   files.New(options.PrometheusRegistry, options.Authorizer),
 		Experiments:                 experiments,
 		WebpushDispatcher:           options.WebPushDispatcher,
@@ -1650,6 +1659,9 @@ type API struct {
 	FileCache           *files.Cache
 	PrebuildsClaimer    atomic.Pointer[prebuilds.Claimer]
 	PrebuildsReconciler atomic.Pointer[prebuilds.ReconciliationOrchestrator]
+	// BuildUsageChecker is a pointer as it's passed around to multiple
+	// components.
+	BuildUsageChecker *atomic.Pointer[wsbuilder.UsageChecker]
 
 	UpdatesProvider tailnet.WorkspaceUpdatesProvider
 

coderd/coderdtest/coderdtest.go

@@ -55,6 +55,7 @@ import (
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/v2/archive"
 	"github.com/coder/coder/v2/coderd/files"
+	"github.com/coder/coder/v2/coderd/wsbuilder"
 	"github.com/coder/quartz"
 
 	"github.com/coder/coder/v2/coderd"
@@ -364,6 +365,10 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 	}
 	connectionLogger.Store(&options.ConnectionLogger)
 
+	var buildUsageChecker atomic.Pointer[wsbuilder.UsageChecker]
+	var noopUsageChecker wsbuilder.UsageChecker = wsbuilder.NoopUsageChecker{}
+	buildUsageChecker.Store(&noopUsageChecker)
+
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	experiments := coderd.ReadExperiments(*options.Logger, options.DeploymentValues.Experiments)
 	lifecycleExecutor := autobuild.NewExecutor(
@@ -375,6 +380,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 		&templateScheduleStore,
 		&auditor,
 		accessControlStore,
+		&buildUsageChecker,
 		*options.Logger,
 		options.AutobuildTicker,
 		options.NotificationsEnqueuer,

coderd/database/dbauthz/dbauthz.go

@@ -2193,6 +2193,14 @@ func (q *querier) GetLogoURL(ctx context.Context) (string, error) {
 	return q.db.GetLogoURL(ctx)
 }
 
+func (q *querier) GetManagedAgentCount(ctx context.Context, arg database.GetManagedAgentCountParams) (int64, error) {
+	// Must be able to read all workspaces to check usage.
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceWorkspace); err != nil {
+		return 0, xerrors.Errorf("authorize read all workspaces: %w", err)
+	}
+	return q.db.GetManagedAgentCount(ctx, arg)
+}
+
 func (q *querier) GetNotificationMessagesByStatus(ctx context.Context, arg database.GetNotificationMessagesByStatusParams) ([]database.NotificationMessage, error) {
 	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceNotificationMessage); err != nil {
 		return nil, err

coderd/database/dbauthz/dbauthz_test.go

@@ -17,20 +17,18 @@ import (
 	"golang.org/x/xerrors"
 
 	"cdr.dev/slog"
-
-	"github.com/coder/coder/v2/coderd/database/db2sdk"
-	"github.com/coder/coder/v2/coderd/notifications"
-	"github.com/coder/coder/v2/coderd/rbac/policy"
-	"github.com/coder/coder/v2/codersdk"
-
 	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
+	"github.com/coder/coder/v2/coderd/database/db2sdk"
 	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"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/notifications"
 	"github.com/coder/coder/v2/coderd/rbac"
+	"github.com/coder/coder/v2/coderd/rbac/policy"
 	"github.com/coder/coder/v2/coderd/util/slice"
+	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/provisionersdk"
 	"github.com/coder/coder/v2/testutil"
 )
@@ -903,6 +901,14 @@ func (s *MethodTestSuite) TestLicense() {
 		require.NoError(s.T(), err)
 		check.Args().Asserts().Returns("value")
 	}))
+	s.Run("GetManagedAgentCount", s.Subtest(func(db database.Store, check *expects) {
+		start := dbtime.Now()
+		end := start.Add(time.Hour)
+		check.Args(database.GetManagedAgentCountParams{
+			StartTime: start,
+			EndTime:   end,
+		}).Asserts(rbac.ResourceWorkspace, policy.ActionRead).Returns(int64(0))
+	}))
 }
 
 func (s *MethodTestSuite) TestOrganization() {

coderd/database/dbmetrics/querymetrics.go

@@ -964,6 +964,13 @@ func (m queryMetricsStore) GetLogoURL(ctx context.Context) (string, error) {
 	return url, err
 }
 
+func (m queryMetricsStore) GetManagedAgentCount(ctx context.Context, arg database.GetManagedAgentCountParams) (int64, error) {
+	start := time.Now()
+	r0, r1 := m.s.GetManagedAgentCount(ctx, arg)
+	m.queryLatencies.WithLabelValues("GetManagedAgentCount").Observe(time.Since(start).Seconds())
+	return r0, r1
+}
+
 func (m queryMetricsStore) GetNotificationMessagesByStatus(ctx context.Context, arg database.GetNotificationMessagesByStatusParams) ([]database.NotificationMessage, error) {
 	start := time.Now()
 	r0, r1 := m.s.GetNotificationMessagesByStatus(ctx, arg)

coderd/database/dbmock/dbmock.go

@@ -2012,6 +2012,21 @@ func (mr *MockStoreMockRecorder) GetLogoURL(ctx any) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetLogoURL", reflect.TypeOf((*MockStore)(nil).GetLogoURL), ctx)
 }
 
+// GetManagedAgentCount mocks base method.
+func (m *MockStore) GetManagedAgentCount(ctx context.Context, arg database.GetManagedAgentCountParams) (int64, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetManagedAgentCount", ctx, arg)
+	ret0, _ := ret[0].(int64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetManagedAgentCount indicates an expected call of GetManagedAgentCount.
+func (mr *MockStoreMockRecorder) GetManagedAgentCount(ctx, arg any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetManagedAgentCount", reflect.TypeOf((*MockStore)(nil).GetManagedAgentCount), ctx, arg)
+}
+
 // GetNotificationMessagesByStatus mocks base method.
 func (m *MockStore) GetNotificationMessagesByStatus(ctx context.Context, arg database.GetNotificationMessagesByStatusParams) ([]database.NotificationMessage, error) {
 	m.ctrl.T.Helper()

coderd/database/querier.go

@@ -216,6 +216,8 @@ type sqlcQuerier interface {
 	GetLicenseByID(ctx context.Context, id int32) (License, error)
 	GetLicenses(ctx context.Context) ([]License, error)
 	GetLogoURL(ctx context.Context) (string, error)
+	// This isn't strictly a license query, but it's related to license enforcement.
+	GetManagedAgentCount(ctx context.Context, arg GetManagedAgentCountParams) (int64, error)
 	GetNotificationMessagesByStatus(ctx context.Context, arg GetNotificationMessagesByStatusParams) ([]NotificationMessage, error)
 	// Fetch the notification report generator log indicating recent activity.
 	GetNotificationReportGeneratorLogByTemplate(ctx context.Context, templateID uuid.UUID) (NotificationReportGeneratorLog, error)

coderd/database/queries.sql.go

@@ -4286,6 +4286,44 @@ func (q *sqlQuerier) GetLicenses(ctx context.Context) ([]License, error) {
 	return items, nil
 }
 
+const getManagedAgentCount = `-- name: GetManagedAgentCount :one
+SELECT
+	COUNT(DISTINCT wb.id) AS count
+FROM
+	workspace_builds AS wb
+JOIN
+	provisioner_jobs AS pj
+ON
+	wb.job_id = pj.id
+WHERE
+	wb.transition = 'start'::workspace_transition
+	AND wb.has_ai_task = true
+	-- Only count jobs that are pending, running or succeeded. Other statuses
+	-- like cancel(ed|ing), failed or unknown are not considered as managed
+	-- agent usage. These workspace builds are typically unusable anyway.
+	AND pj.job_status IN (
+		'pending'::provisioner_job_status,
+		'running'::provisioner_job_status,
+		'succeeded'::provisioner_job_status
+	)
+	-- Jobs are counted at the time they are created, not when they are
+	-- completed, as pending jobs haven't completed yet.
+	AND wb.created_at BETWEEN $1::timestamptz AND $2::timestamptz
+`
+
+type GetManagedAgentCountParams struct {
+	StartTime time.Time `db:"start_time" json:"start_time"`
+	EndTime   time.Time `db:"end_time" json:"end_time"`
+}
+
+// This isn't strictly a license query, but it's related to license enforcement.
+func (q *sqlQuerier) GetManagedAgentCount(ctx context.Context, arg GetManagedAgentCountParams) (int64, error) {
+	row := q.db.QueryRowContext(ctx, getManagedAgentCount, arg.StartTime, arg.EndTime)
+	var count int64
+	err := row.Scan(&count)
+	return count, err
+}
+
 const getUnexpiredLicenses = `-- name: GetUnexpiredLicenses :many
 SELECT id, uploaded_at, jwt, exp, uuid
 FROM licenses

coderd/database/queries/licenses.sql

@@ -35,3 +35,28 @@ DELETE
 FROM licenses
 WHERE id = $1
 RETURNING id;
+
+-- name: GetManagedAgentCount :one
+-- This isn't strictly a license query, but it's related to license enforcement.
+SELECT
+	COUNT(DISTINCT wb.id) AS count
+FROM
+	workspace_builds AS wb
+JOIN
+	provisioner_jobs AS pj
+ON
+	wb.job_id = pj.id
+WHERE
+	wb.transition = 'start'::workspace_transition
+	AND wb.has_ai_task = true
+	-- Only count jobs that are pending, running or succeeded. Other statuses
+	-- like cancel(ed|ing), failed or unknown are not considered as managed
+	-- agent usage. These workspace builds are typically unusable anyway.
+	AND pj.job_status IN (
+		'pending'::provisioner_job_status,
+		'running'::provisioner_job_status,
+		'succeeded'::provisioner_job_status
+	)
+	-- Jobs are counted at the time they are created, not when they are
+	-- completed, as pending jobs haven't completed yet.
+	AND wb.created_at BETWEEN @start_time::timestamptz AND @end_time::timestamptz;

coderd/workspacebuilds.go

@@ -335,7 +335,7 @@ func (api *API) postWorkspaceBuilds(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	builder := wsbuilder.New(workspace, database.WorkspaceTransition(createBuild.Transition)).
+	builder := wsbuilder.New(workspace, database.WorkspaceTransition(createBuild.Transition), *api.BuildUsageChecker.Load()).
 		Initiator(apiKey.UserID).
 		RichParameterValues(createBuild.RichParameterValues).
 		LogLevel(string(createBuild.LogLevel)).

coderd/workspaces.go

@@ -701,7 +701,7 @@ func createWorkspace(
 			return xerrors.Errorf("get workspace by ID: %w", err)
 		}
 
-		builder := wsbuilder.New(workspace, database.WorkspaceTransitionStart).
+		builder := wsbuilder.New(workspace, database.WorkspaceTransitionStart, *api.BuildUsageChecker.Load()).
 			Reason(database.BuildReasonInitiator).
 			Initiator(initiatorID).
 			ActiveVersion().

coderd/wsbuilder/wsbuilder.go

@@ -56,6 +56,7 @@ type Builder struct {
 	logLevel         string
 	deploymentValues *codersdk.DeploymentValues
 	experiments      codersdk.Experiments
+	usageChecker     UsageChecker
 
 	richParameterValues     []codersdk.WorkspaceBuildParameter
 	initiator               uuid.UUID
@@ -89,7 +90,24 @@ type Builder struct {
 	verifyNoLegacyParametersOnce bool
 }
 
-type Option func(Builder) Builder
+type UsageChecker interface {
+	CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (UsageCheckResponse, error)
+}
+
+type UsageCheckResponse struct {
+	Permitted bool
+	Message   string
+}
+
+type NoopUsageChecker struct{}
+
+var _ UsageChecker = NoopUsageChecker{}
+
+func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion) (UsageCheckResponse, error) {
+	return UsageCheckResponse{
+		Permitted: true,
+	}, nil
+}
 
 // versionTarget expresses how to determine the template version for the build.
 //
@@ -121,8 +139,8 @@ type stateTarget struct {
 	explicit *[]byte
 }
 
-func New(w database.Workspace, t database.WorkspaceTransition) Builder {
-	return Builder{workspace: w, trans: t}
+func New(w database.Workspace, t database.WorkspaceTransition, uc UsageChecker) Builder {
+	return Builder{workspace: w, trans: t, usageChecker: uc}
 }
 
 // Methods that customize the build are public, have a struct receiver and return a new Builder.
@@ -321,6 +339,10 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
 	if err != nil {
 		return nil, nil, nil, err
 	}
+	err = b.checkUsage()
+	if err != nil {
+		return nil, nil, nil, err
+	}
 	err = b.checkRunningBuild()
 	if err != nil {
 		return nil, nil, nil, err
@@ -1253,6 +1275,23 @@ func (b *Builder) checkTemplateJobStatus() error {
 	return nil
 }
 
+func (b *Builder) checkUsage() error {
+	templateVersion, err := b.getTemplateVersion()
+	if err != nil {
+		return BuildError{http.StatusInternalServerError, "Failed to fetch template version", err}
+	}
+
+	resp, err := b.usageChecker.CheckBuildUsage(b.ctx, b.store, templateVersion)
+	if err != nil {
+		return BuildError{http.StatusInternalServerError, "Failed to check build usage", err}
+	}
+	if !resp.Permitted {
+		return BuildError{http.StatusForbidden, "Build is not permitted: " + resp.Message, nil}
+	}
+
+	return nil
+}
+
 func (b *Builder) checkRunningBuild() error {
 	job, err := b.getLastBuildJob()
 	if xerrors.Is(err, sql.ErrNoRows) {

coderd/wsbuilder/wsbuilder_test.go

@@ -5,30 +5,30 @@ import (
 	"database/sql"
 	"encoding/json"
 	"net/http"
+	"sync/atomic"
 	"testing"
 	"time"
 
-	"github.com/prometheus/client_golang/prometheus"
-
-	"github.com/coder/coder/v2/coderd/coderdtest"
-	"github.com/coder/coder/v2/coderd/files"
-	"github.com/coder/coder/v2/coderd/httpapi/httperror"
-	"github.com/coder/coder/v2/provisionersdk"
-
 	"github.com/google/uuid"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/propagation"
 	"go.uber.org/mock/gomock"
+	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/coderd/audit"
+	"github.com/coder/coder/v2/coderd/coderdtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbmock"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/files"
+	"github.com/coder/coder/v2/coderd/httpapi/httperror"
 	"github.com/coder/coder/v2/coderd/provisionerdserver"
 	"github.com/coder/coder/v2/coderd/wsbuilder"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/provisionersdk"
 )
 
 var (
@@ -102,7 +102,7 @@ func TestBuilder_NoOptions(t *testing.T) {
 	fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 	ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart)
+	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{})
 	// nolint: dogsled
 	_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
 	req.NoError(err)
@@ -142,7 +142,8 @@ func TestBuilder_Initiator(t *testing.T) {
 	fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 	ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).Initiator(otherUserID)
+	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
+		Initiator(otherUserID)
 	// nolint: dogsled
 	_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
 	req.NoError(err)
@@ -188,7 +189,8 @@ func TestBuilder_Baggage(t *testing.T) {
 	fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 	ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).Initiator(otherUserID)
+	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
+		Initiator(otherUserID)
 	// nolint: dogsled
 	_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{IP: "127.0.0.1"})
 	req.NoError(err)
@@ -227,7 +229,8 @@ func TestBuilder_Reason(t *testing.T) {
 	fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 	ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).Reason(database.BuildReasonAutostart)
+	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
+		Reason(database.BuildReasonAutostart)
 	// nolint: dogsled
 	_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
 	req.NoError(err)
@@ -271,7 +274,8 @@ func TestBuilder_ActiveVersion(t *testing.T) {
 	fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 	ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).ActiveVersion()
+	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
+		ActiveVersion()
 	// nolint: dogsled
 	_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
 	req.NoError(err)
@@ -386,7 +390,8 @@ func TestWorkspaceBuildWithTags(t *testing.T) {
 	fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 	ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).RichParameterValues(buildParameters)
+	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
+		RichParameterValues(buildParameters)
 	// nolint: dogsled
 	_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
 	req.NoError(err)
@@ -469,7 +474,8 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
 		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).RichParameterValues(nextBuildParameters)
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
+			RichParameterValues(nextBuildParameters)
 		// nolint: dogsled
 		_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
 		req.NoError(err)
@@ -517,7 +523,8 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
 		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).RichParameterValues(nextBuildParameters)
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
+			RichParameterValues(nextBuildParameters)
 		// nolint: dogsled
 		_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
 		req.NoError(err)
@@ -555,7 +562,8 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
 		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart)
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{})
+		// nolint: dogsled
 		_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
 		bldErr := wsbuilder.BuildError{}
 		req.ErrorAs(err, &bldErr)
@@ -591,7 +599,8 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
 		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).RichParameterValues(nextBuildParameters)
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
+			RichParameterValues(nextBuildParameters)
 		// nolint: dogsled
 		_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
 		bldErr := wsbuilder.BuildError{}
@@ -656,7 +665,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
 		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
 			RichParameterValues(nextBuildParameters).
 			VersionID(activeVersionID)
 		_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
@@ -720,7 +729,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
 		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
 			RichParameterValues(nextBuildParameters).
 			VersionID(activeVersionID)
 		_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
@@ -782,7 +791,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
 		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
 			RichParameterValues(nextBuildParameters).
 			VersionID(activeVersionID)
 		// nolint: dogsled
@@ -849,7 +858,7 @@ func TestWorkspaceBuildWithPreset(t *testing.T) {
 	fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 	ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).
+	uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
 		ActiveVersion().
 		TemplateVersionPresetID(presetID)
 	// nolint: dogsled
@@ -916,7 +925,7 @@ func TestWorkspaceBuildDeleteOrphan(t *testing.T) {
 		)
 
 		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-		uut := wsbuilder.New(ws, database.WorkspaceTransitionDelete).Orphan()
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionDelete, wsbuilder.NoopUsageChecker{}).Orphan()
 		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 
 		// nolint: dogsled
@@ -993,7 +1002,7 @@ func TestWorkspaceBuildDeleteOrphan(t *testing.T) {
 		)
 
 		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
-		uut := wsbuilder.New(ws, database.WorkspaceTransitionDelete).Orphan()
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionDelete, wsbuilder.NoopUsageChecker{}).Orphan()
 		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 		// nolint: dogsled
 		_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
@@ -1001,6 +1010,115 @@ func TestWorkspaceBuildDeleteOrphan(t *testing.T) {
 	})
 }
 
+func TestWorkspaceBuildUsageChecker(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Permitted", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		var calls int64
+		fakeUsageChecker := &fakeUsageChecker{
+			checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error) {
+				atomic.AddInt64(&calls, 1)
+				return wsbuilder.UsageCheckResponse{Permitted: true}, nil
+			},
+		}
+
+		mDB := expectDB(t,
+			// Inputs
+			withTemplate,
+			withInactiveVersion(nil),
+			withLastBuildFound,
+			withTemplateVersionVariables(inactiveVersionID, nil),
+			withRichParameters(nil),
+			withParameterSchemas(inactiveJobID, nil),
+			withWorkspaceTags(inactiveVersionID, nil),
+			withProvisionerDaemons([]database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow{}),
+
+			// Outputs
+			expectProvisionerJob(func(job database.InsertProvisionerJobParams) {}),
+			withInTx,
+			expectBuild(func(bld database.InsertWorkspaceBuildParams) {}),
+			withBuild,
+			expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) {}),
+		)
+		fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
+
+		ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
+		uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, fakeUsageChecker)
+		// nolint: dogsled
+		_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
+		require.NoError(t, err)
+		require.EqualValues(t, 1, calls)
+	})
+
+	// The failure cases are mostly identical from a test perspective.
+	const message = "fake test message"
+	cases := []struct {
+		name        string
+		response    wsbuilder.UsageCheckResponse
+		responseErr error
+		assertions  func(t *testing.T, err error)
+	}{
+		{
+			name: "NotPermitted",
+			response: wsbuilder.UsageCheckResponse{
+				Permitted: false,
+				Message:   message,
+			},
+			assertions: func(t *testing.T, err error) {
+				require.ErrorContains(t, err, message)
+				var buildErr wsbuilder.BuildError
+				require.ErrorAs(t, err, &buildErr)
+				require.Equal(t, http.StatusForbidden, buildErr.Status)
+			},
+		},
+		{
+			name:        "Error",
+			responseErr: xerrors.New("fake error"),
+			assertions: func(t *testing.T, err error) {
+				require.ErrorContains(t, err, "fake error")
+				require.ErrorAs(t, err, &wsbuilder.BuildError{})
+			},
+		},
+	}
+
+	for _, c := range cases {
+		c := c
+		t.Run(c.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			var calls int64
+			fakeUsageChecker := &fakeUsageChecker{
+				checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error) {
+					atomic.AddInt64(&calls, 1)
+					return c.response, c.responseErr
+				},
+			}
+
+			mDB := expectDB(t,
+				withTemplate,
+				withInactiveVersionNoParams(),
+			)
+			fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
+
+			ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
+			uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, fakeUsageChecker).
+				VersionID(inactiveVersionID)
+			// nolint: dogsled
+			_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
+			c.assertions(t, err)
+			require.EqualValues(t, 1, calls)
+		})
+	}
+}
+
 func TestWsbuildError(t *testing.T) {
 	t.Parallel()
 
@@ -1366,3 +1484,11 @@ func withProvisionerDaemons(provisionerDaemons []database.GetEligibleProvisioner
 		mTx.EXPECT().GetEligibleProvisionerDaemonsByProvisionerJobIDs(gomock.Any(), gomock.Any()).Return(provisionerDaemons, nil)
 	}
 }
+
+type fakeUsageChecker struct {
+	checkBuildUsageFunc func(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error)
+}
+
+func (f *fakeUsageChecker) CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error) {
+	return f.checkBuildUsageFunc(ctx, store, templateVersion)
+}