Skip to content

Commit 9a6dd73

Browse files
deansheatherdeansheather
authored
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
1 parent aa1a985 commit 9a6dd73

24 files changed

+586
-74
lines changed

‎cli/server.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1101,7 +1101,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
11011101
autobuildTicker := time.NewTicker(vals.AutobuildPollInterval.Value())
11021102
defer autobuildTicker.Stop()
11031103
autobuildExecutor := autobuild.NewExecutor(
1104-
ctx, options.Database, options.Pubsub, coderAPI.FileCache, options.PrometheusRegistry, coderAPI.TemplateScheduleStore, &coderAPI.Auditor, coderAPI.AccessControlStore, logger, autobuildTicker.C, options.NotificationsEnqueuer, coderAPI.Experiments)
1104+
ctx, options.Database, options.Pubsub, coderAPI.FileCache, options.PrometheusRegistry, coderAPI.TemplateScheduleStore, &coderAPI.Auditor, coderAPI.AccessControlStore, coderAPI.BuildUsageChecker, logger, autobuildTicker.C, options.NotificationsEnqueuer, coderAPI.Experiments)
11051105
autobuildExecutor.Run()
11061106

11071107
jobReaperTicker := time.NewTicker(vals.JobReaperDetectorInterval.Value())

‎coderd/autobuild/lifecycle_executor.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ type Executor struct {
4242
templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore]
4343
accessControlStore *atomic.Pointer[dbauthz.AccessControlStore]
4444
auditor *atomic.Pointer[audit.Auditor]
45+
buildUsageChecker *atomic.Pointer[wsbuilder.UsageChecker]
4546
log slog.Logger
4647
tick <-chan time.Time
4748
statsCh chan<- Stats
@@ -65,7 +66,7 @@ type Stats struct {
6566
}
6667

6768
// New returns a new wsactions executor.
68-
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 {
69+
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 {
6970
factory := promauto.With(reg)
7071
le := &Executor{
7172
//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
7879
log: log.Named("autobuild"),
7980
auditor: auditor,
8081
accessControlStore: acs,
82+
buildUsageChecker: buildUsageChecker,
8183
notificationsEnqueuer: enqueuer,
8284
reg: reg,
8385
experiments: exp,
@@ -279,7 +281,7 @@ func (e *Executor) runOnce(t time.Time) Stats {
279281
}
280282

281283
if nextTransition != "" {
282-
builder := wsbuilder.New(ws, nextTransition).
284+
builder := wsbuilder.New(ws, nextTransition, *e.buildUsageChecker.Load()).
283285
SetLastWorkspaceBuildInTx(&latestBuild).
284286
SetLastWorkspaceBuildJobInTx(&latestJob).
285287
Experiments(e.experiments).

‎coderd/coderd.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121

2222
"github.com/coder/coder/v2/coderd/oauth2provider"
2323
"github.com/coder/coder/v2/coderd/prebuilds"
24+
"github.com/coder/coder/v2/coderd/wsbuilder"
2425

2526
"github.com/andybalholm/brotli"
2627
"github.com/go-chi/chi/v5"
@@ -559,6 +560,13 @@ func New(options *Options) *API {
559560
// bugs that may only occur when a key isn't precached in tests and the latency cost is minimal.
560561
cryptokeys.StartRotator(ctx, options.Logger, options.Database)
561562

563+
// AGPL uses a no-op build usage checker as there are no license
564+
// entitlements to enforce. This is swapped out in
565+
// enterprise/coderd/coderd.go.
566+
var buildUsageChecker atomic.Pointer[wsbuilder.UsageChecker]
567+
var noopUsageChecker wsbuilder.UsageChecker = wsbuilder.NoopUsageChecker{}
568+
buildUsageChecker.Store(&noopUsageChecker)
569+
562570
api := &API{
563571
ctx: ctx,
564572
cancel: cancel,
@@ -579,6 +587,7 @@ func New(options *Options) *API {
579587
TemplateScheduleStore: options.TemplateScheduleStore,
580588
UserQuietHoursScheduleStore: options.UserQuietHoursScheduleStore,
581589
AccessControlStore: options.AccessControlStore,
590+
BuildUsageChecker: &buildUsageChecker,
582591
FileCache: files.New(options.PrometheusRegistry, options.Authorizer),
583592
Experiments: experiments,
584593
WebpushDispatcher: options.WebPushDispatcher,
@@ -1650,6 +1659,9 @@ type API struct {
16501659
FileCache *files.Cache
16511660
PrebuildsClaimer atomic.Pointer[prebuilds.Claimer]
16521661
PrebuildsReconciler atomic.Pointer[prebuilds.ReconciliationOrchestrator]
1662+
// BuildUsageChecker is a pointer as it's passed around to multiple
1663+
// components.
1664+
BuildUsageChecker *atomic.Pointer[wsbuilder.UsageChecker]
16531665

16541666
UpdatesProvider tailnet.WorkspaceUpdatesProvider
16551667

‎coderd/coderdtest/coderdtest.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ import (
5555
"cdr.dev/slog/sloggers/slogtest"
5656
"github.com/coder/coder/v2/archive"
5757
"github.com/coder/coder/v2/coderd/files"
58+
"github.com/coder/coder/v2/coderd/wsbuilder"
5859
"github.com/coder/quartz"
5960

6061
"github.com/coder/coder/v2/coderd"
@@ -364,6 +365,10 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
364365
}
365366
connectionLogger.Store(&options.ConnectionLogger)
366367

368+
var buildUsageChecker atomic.Pointer[wsbuilder.UsageChecker]
369+
var noopUsageChecker wsbuilder.UsageChecker = wsbuilder.NoopUsageChecker{}
370+
buildUsageChecker.Store(&noopUsageChecker)
371+
367372
ctx, cancelFunc := context.WithCancel(context.Background())
368373
experiments := coderd.ReadExperiments(*options.Logger, options.DeploymentValues.Experiments)
369374
lifecycleExecutor := autobuild.NewExecutor(
@@ -375,6 +380,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
375380
&templateScheduleStore,
376381
&auditor,
377382
accessControlStore,
383+
&buildUsageChecker,
378384
*options.Logger,
379385
options.AutobuildTicker,
380386
options.NotificationsEnqueuer,

‎coderd/database/dbauthz/dbauthz.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2193,6 +2193,14 @@ func (q *querier) GetLogoURL(ctx context.Context) (string, error) {
21932193
return q.db.GetLogoURL(ctx)
21942194
}
21952195

2196+
func (q *querier) GetManagedAgentCount(ctx context.Context, arg database.GetManagedAgentCountParams) (int64, error) {
2197+
// Must be able to read all workspaces to check usage.
2198+
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceWorkspace); err != nil {
2199+
return 0, xerrors.Errorf("authorize read all workspaces: %w", err)
2200+
}
2201+
return q.db.GetManagedAgentCount(ctx, arg)
2202+
}
2203+
21962204
func (q *querier) GetNotificationMessagesByStatus(ctx context.Context, arg database.GetNotificationMessagesByStatusParams) ([]database.NotificationMessage, error) {
21972205
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceNotificationMessage); err != nil {
21982206
return nil, err

‎coderd/database/dbauthz/dbauthz_test.go

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,20 +17,18 @@ import (
1717
"golang.org/x/xerrors"
1818

1919
"cdr.dev/slog"
20-
21-
"github.com/coder/coder/v2/coderd/database/db2sdk"
22-
"github.com/coder/coder/v2/coderd/notifications"
23-
"github.com/coder/coder/v2/coderd/rbac/policy"
24-
"github.com/coder/coder/v2/codersdk"
25-
2620
"github.com/coder/coder/v2/coderd/coderdtest"
2721
"github.com/coder/coder/v2/coderd/database"
22+
"github.com/coder/coder/v2/coderd/database/db2sdk"
2823
"github.com/coder/coder/v2/coderd/database/dbauthz"
2924
"github.com/coder/coder/v2/coderd/database/dbgen"
3025
"github.com/coder/coder/v2/coderd/database/dbtestutil"
3126
"github.com/coder/coder/v2/coderd/database/dbtime"
27+
"github.com/coder/coder/v2/coderd/notifications"
3228
"github.com/coder/coder/v2/coderd/rbac"
29+
"github.com/coder/coder/v2/coderd/rbac/policy"
3330
"github.com/coder/coder/v2/coderd/util/slice"
31+
"github.com/coder/coder/v2/codersdk"
3432
"github.com/coder/coder/v2/provisionersdk"
3533
"github.com/coder/coder/v2/testutil"
3634
)
@@ -903,6 +901,14 @@ func (s *MethodTestSuite) TestLicense() {
903901
require.NoError(s.T(), err)
904902
check.Args().Asserts().Returns("value")
905903
}))
904+
s.Run("GetManagedAgentCount", s.Subtest(func(db database.Store, check *expects) {
905+
start := dbtime.Now()
906+
end := start.Add(time.Hour)
907+
check.Args(database.GetManagedAgentCountParams{
908+
StartTime: start,
909+
EndTime: end,
910+
}).Asserts(rbac.ResourceWorkspace, policy.ActionRead).Returns(int64(0))
911+
}))
906912
}
907913

908914
func (s *MethodTestSuite) TestOrganization() {

‎coderd/database/dbmetrics/querymetrics.go

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/querier.go

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go

Lines changed: 38 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)