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

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
+3 -3
View File
@@ -93,7 +93,7 @@ type Builder struct {
}
type UsageChecker interface {
CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (UsageCheckResponse, error)
CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, transition database.WorkspaceTransition) (UsageCheckResponse, error)
}
type UsageCheckResponse struct {
@@ -105,7 +105,7 @@ type NoopUsageChecker struct{}
var _ UsageChecker = NoopUsageChecker{}
func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion) (UsageCheckResponse, error) {
func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ database.WorkspaceTransition) (UsageCheckResponse, error) {
return UsageCheckResponse{
Permitted: true,
}, nil
@@ -1307,7 +1307,7 @@ func (b *Builder) checkUsage() error {
return BuildError{http.StatusInternalServerError, "Failed to fetch template version", err}
}
resp, err := b.usageChecker.CheckBuildUsage(b.ctx, b.store, templateVersion)
resp, err := b.usageChecker.CheckBuildUsage(b.ctx, b.store, templateVersion, b.trans)
if err != nil {
return BuildError{http.StatusInternalServerError, "Failed to check build usage", err}
}
+5 -5
View File
@@ -1049,7 +1049,7 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) {
var calls int64
fakeUsageChecker := &fakeUsageChecker{
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error) {
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
atomic.AddInt64(&calls, 1)
return wsbuilder.UsageCheckResponse{Permitted: true}, nil
},
@@ -1126,7 +1126,7 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) {
var calls int64
fakeUsageChecker := &fakeUsageChecker{
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error) {
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
atomic.AddInt64(&calls, 1)
return c.response, c.responseErr
},
@@ -1577,11 +1577,11 @@ func expectFindMatchingPresetID(id uuid.UUID, err error) func(mTx *dbmock.MockSt
}
type fakeUsageChecker struct {
checkBuildUsageFunc func(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error)
checkBuildUsageFunc func(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, transition database.WorkspaceTransition) (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)
func (f *fakeUsageChecker) CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
return f.checkBuildUsageFunc(ctx, store, templateVersion, transition)
}
func withNoTask(mTx *dbmock.MockStore) {