diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 508f4d8277..00a80c9540 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -987,8 +987,8 @@ func (api *API) CheckBuildUsage( _ context.Context, _ database.Store, templateVersion *database.TemplateVersion, - _ *database.Task, - _ database.WorkspaceTransition, + task *database.Task, + transition database.WorkspaceTransition, ) (wsbuilder.UsageCheckResponse, error) { // If the template version has an external agent, we need to check that the // license is entitled to this feature. @@ -1002,6 +1002,25 @@ func (api *API) CheckBuildUsage( } } + // Verify managed agent entitlement for AI task builds. + // The count/limit check is intentionally omitted — breaching the + // limit is advisory only and surfaced as a warning via entitlements. + if transition != database.WorkspaceTransitionStart || task == nil { + return wsbuilder.UsageCheckResponse{Permitted: true}, nil + } + + if !api.Entitlements.HasLicense() { + return wsbuilder.UsageCheckResponse{Permitted: true}, nil + } + + managedAgentLimit, ok := api.Entitlements.Feature(codersdk.FeatureManagedAgentLimit) + if !ok || !managedAgentLimit.Enabled { + return wsbuilder.UsageCheckResponse{ + Permitted: false, + Message: "Your license is not entitled to managed agents. Please contact sales to continue using managed agents.", + }, nil + } + return wsbuilder.UsageCheckResponse{Permitted: true}, nil } diff --git a/enterprise/coderd/coderd_test.go b/enterprise/coderd/coderd_test.go index 7662078c7e..e9c4d22779 100644 --- a/enterprise/coderd/coderd_test.go +++ b/enterprise/coderd/coderd_test.go @@ -849,6 +849,88 @@ func TestCheckBuildUsage_NeverBlocksOnManagedAgentLimit(t *testing.T) { require.True(t, deleteResp.Permitted) } +func TestCheckBuildUsage_BlocksWithoutManagedAgentEntitlement(t *testing.T) { + t.Parallel() + + tv := &database.TemplateVersion{ + HasAITask: sql.NullBool{Valid: true, Bool: true}, + HasExternalAgent: sql.NullBool{Valid: true, Bool: false}, + } + task := &database.Task{ + TemplateVersionID: tv.ID, + } + + // Both "feature absent" and "feature explicitly disabled" should + // block AI task builds on licensed deployments. + tests := []struct { + name string + setupEnts func(e *codersdk.Entitlements) + }{ + { + name: "FeatureAbsent", + setupEnts: func(e *codersdk.Entitlements) { + e.HasLicense = true + }, + }, + { + name: "FeatureDisabled", + setupEnts: func(e *codersdk.Entitlements) { + e.HasLicense = true + e.Features[codersdk.FeatureManagedAgentLimit] = codersdk.Feature{ + Enabled: false, + } + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + entSet := entitlements.New() + entSet.Modify(tc.setupEnts) + + agpl := &agplcoderd.API{ + Options: &agplcoderd.Options{ + Entitlements: entSet, + }, + } + eapi := &coderd.API{ + AGPL: agpl, + Options: &coderd.Options{Options: agpl.Options}, + } + + mDB := dbmock.NewMockStore(ctrl) + ctx := context.Background() + + // Start transition with a task: should be blocked because the + // license doesn't include the managed agent entitlement. + resp, err := eapi.CheckBuildUsage(ctx, mDB, tv, task, database.WorkspaceTransitionStart) + require.NoError(t, err) + require.False(t, resp.Permitted) + require.Contains(t, resp.Message, "not entitled to managed agents") + + // Stop and delete transitions should still be permitted so + // that existing workspaces can be stopped/cleaned up. + stopResp, err := eapi.CheckBuildUsage(ctx, mDB, tv, task, database.WorkspaceTransitionStop) + require.NoError(t, err) + require.True(t, stopResp.Permitted) + + deleteResp, err := eapi.CheckBuildUsage(ctx, mDB, tv, task, database.WorkspaceTransitionDelete) + require.NoError(t, err) + require.True(t, deleteResp.Permitted) + + // Start transition without a task: should be permitted (not + // an AI task build, so the entitlement check doesn't apply). + noTaskResp, err := eapi.CheckBuildUsage(ctx, mDB, tv, nil, database.WorkspaceTransitionStart) + require.NoError(t, err) + require.True(t, noTaskResp.Permitted) + }) + } +} + // testDBAuthzRole returns a context with a subject that has a role // with permissions required for test setup. func testDBAuthzRole(ctx context.Context) context.Context {