opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: allow stops and deletes after breaching AI limit (#21186)

Browse Source 
Fixes a bug a customer encountered once they breached their limit. Adds
a test.
This commit is contained in:
Dean Sheather
2025-12-09 11:05:12 +00:00
committed by GitHub
parent 97bc7eb9e5
commit b199eb1c38
4 changed files with 109 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/coderd/coderd.go
+28 -14
View File
@@ -971,7 +971,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
var _ wsbuilder.UsageChecker = &API{}
func (api *API) CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error) {
func (api *API) CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, 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.
if templateVersion.HasExternalAgent.Valid && templateVersion.HasExternalAgent.Bool {
@@ -984,16 +984,31 @@ func (api *API) CheckBuildUsage(ctx context.Context, store database.Store, templ
}
}
// If the template version doesn't have an AI task, we don't need to check
// usage.
if !templateVersion.HasAITask.Valid || !templateVersion.HasAITask.Bool {
return wsbuilder.UsageCheckResponse{
Permitted: true,
}, nil
resp, err := api.checkAIBuildUsage(ctx, store, templateVersion, transition)
if err != nil {
return wsbuilder.UsageCheckResponse{}, err
}
if !resp.Permitted {
return resp, nil
}
// When unlicensed, we need to check that we haven't breached the managed agent
// limit.
return wsbuilder.UsageCheckResponse{Permitted: true}, nil
}
// checkAIBuildUsage validates AI-related usage constraints. It is a no-op
// unless the transition is "start" and the template version has an AI task.
func (api *API) checkAIBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
// Only check AI usage rules for start transitions.
if transition != database.WorkspaceTransitionStart {
return wsbuilder.UsageCheckResponse{Permitted: true}, nil
}
// If the template version doesn't have an AI task, we don't need to check usage.
if !templateVersion.HasAITask.Valid || !templateVersion.HasAITask.Bool {
return wsbuilder.UsageCheckResponse{Permitted: true}, nil
}
// When licensed, ensure we haven't breached the managed agent limit.
// Unlicensed deployments are allowed to use unlimited managed agents.
if api.Entitlements.HasLicense() {
managedAgentLimit, ok := api.Entitlements.Feature(codersdk.FeatureManagedAgentLimit)
@@ -1004,8 +1019,9 @@ func (api *API) CheckBuildUsage(ctx context.Context, store database.Store, templ
}, nil
}
// This check is intentionally not committed to the database. It's fine if
// it's not 100% accurate or allows for minor breaches due to build races.
// This check is intentionally not committed to the database. It's fine
// if it's not 100% accurate or allows for minor breaches due to build
// races.
// nolint:gocritic // Requires permission to read all usage events.
managedAgentCount, err := store.GetTotalUsageDCManagedAgentsV1(agpldbauthz.AsSystemRestricted(ctx), database.GetTotalUsageDCManagedAgentsV1Params{
StartDate: managedAgentLimit.UsagePeriod.Start,
@@ -1023,9 +1039,7 @@ func (api *API) CheckBuildUsage(ctx context.Context, store database.Store, templ
}
}
return wsbuilder.UsageCheckResponse{
Permitted: true,
}, nil
return wsbuilder.UsageCheckResponse{Permitted: true}, nil
}
// getProxyDERPStartingRegionID returns the starting region ID that should be