fix: do not enforce managed agent limit for non-task workspaces (#21689)

Only task workspaces have the checks in wsbuilder for violating the
managed agent caps in the license.

Stopped tasks that are resumed with a regular workspace start **still
count as usage**.
This commit is contained in:
Steven Masley
2026-01-27 19:01:17 -06:00
committed by GitHub
parent 3eeeabfd68
commit 799b190dee
7 changed files with 225 additions and 19 deletions
+36 -7
View File
@@ -6,7 +6,6 @@ import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"net/http"
"time"
@@ -85,13 +84,15 @@ type Builder struct {
templateVersionPresetParameterValues *[]database.TemplateVersionPresetParameter
parameterRender dynamicparameters.Renderer
workspaceTags *map[string]string
task *database.Task
hasTask *bool // A workspace without a task will have a nil `task` and false `hasTask`.
prebuiltWorkspaceBuildStage sdkproto.PrebuiltWorkspaceBuildStage
verifyNoLegacyParametersOnce bool
}
type UsageChecker interface {
CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, transition database.WorkspaceTransition) (UsageCheckResponse, error)
CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, task *database.Task, transition database.WorkspaceTransition) (UsageCheckResponse, error)
}
type UsageCheckResponse struct {
@@ -103,7 +104,7 @@ type NoopUsageChecker struct{}
var _ UsageChecker = NoopUsageChecker{}
func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ database.WorkspaceTransition) (UsageCheckResponse, error) {
func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ *database.Task, _ database.WorkspaceTransition) (UsageCheckResponse, error) {
return UsageCheckResponse{
Permitted: true,
}, nil
@@ -487,8 +488,12 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
return BuildError{code, "insert workspace build", err}
}
task, err := b.getWorkspaceTask()
if err != nil {
return BuildError{http.StatusInternalServerError, "get task by workspace id", err}
}
// If this is a task workspace, link it to the latest workspace build.
if task, err := store.GetTaskByWorkspaceID(b.ctx, b.workspace.ID); err == nil {
if task != nil {
_, err = store.UpsertTaskWorkspaceApp(b.ctx, database.UpsertTaskWorkspaceAppParams{
TaskID: task.ID,
WorkspaceBuildNumber: buildNum,
@@ -498,8 +503,6 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
if err != nil {
return BuildError{http.StatusInternalServerError, "upsert task workspace app", err}
}
} else if !errors.Is(err, sql.ErrNoRows) {
return BuildError{http.StatusInternalServerError, "get task by workspace id", err}
}
err = store.InsertWorkspaceBuildParameters(b.ctx, database.InsertWorkspaceBuildParametersParams{
@@ -632,6 +635,27 @@ func (b *Builder) getTemplateVersionID() (uuid.UUID, error) {
return bld.TemplateVersionID, nil
}
// getWorkspaceTask returns the task associated with the workspace, if any.
// If no task exists, it returns (nil, nil).
func (b *Builder) getWorkspaceTask() (*database.Task, error) {
if b.hasTask != nil {
return b.task, nil
}
t, err := b.store.GetTaskByWorkspaceID(b.ctx, b.workspace.ID)
if err != nil {
if xerrors.Is(err, sql.ErrNoRows) {
b.hasTask = ptr.Ref(false)
//nolint:nilnil // No task exists.
return nil, nil
}
return nil, xerrors.Errorf("get task: %w", err)
}
b.task = &t
b.hasTask = ptr.Ref(true)
return b.task, nil
}
func (b *Builder) getTemplateTerraformValues() (*database.TemplateVersionTerraformValue, error) {
if b.terraformValues != nil {
return b.terraformValues, nil
@@ -1313,7 +1337,12 @@ 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, b.trans)
task, err := b.getWorkspaceTask()
if err != nil {
return BuildError{http.StatusInternalServerError, "Failed to fetch workspace task", err}
}
resp, err := b.usageChecker.CheckBuildUsage(b.ctx, b.store, templateVersion, task, b.trans)
if err != nil {
return BuildError{http.StatusInternalServerError, "Failed to check build usage", err}
}
+8 -5
View File
@@ -570,6 +570,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
mDB := expectDB(t,
// Inputs
withTemplate,
withNoTask,
withInactiveVersionNoParams(),
withLastBuildFound,
withTemplateVersionVariables(inactiveVersionID, nil),
@@ -605,6 +606,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
withTemplate,
withInactiveVersion(richParameters),
withLastBuildFound,
withNoTask,
withTemplateVersionVariables(inactiveVersionID, nil),
withRichParameters(initialBuildParameters),
withParameterSchemas(inactiveJobID, nil),
@@ -1049,7 +1051,7 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) {
var calls int64
fakeUsageChecker := &fakeUsageChecker{
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
checkBuildUsageFunc: func(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ *database.Task, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
atomic.AddInt64(&calls, 1)
return wsbuilder.UsageCheckResponse{Permitted: true}, nil
},
@@ -1126,7 +1128,7 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) {
var calls int64
fakeUsageChecker := &fakeUsageChecker{
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
checkBuildUsageFunc: func(_ context.Context, _ database.Store, _ *database.TemplateVersion, _ *database.Task, _ database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
atomic.AddInt64(&calls, 1)
return c.response, c.responseErr
},
@@ -1134,6 +1136,7 @@ func TestWorkspaceBuildUsageChecker(t *testing.T) {
mDB := expectDB(t,
withTemplate,
withNoTask,
withInactiveVersionNoParams(),
)
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
@@ -1577,11 +1580,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, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error)
checkBuildUsageFunc func(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, task *database.Task, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error)
}
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 (f *fakeUsageChecker) CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, task *database.Task, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
return f.checkBuildUsageFunc(ctx, store, templateVersion, task, transition)
}
func withNoTask(mTx *dbmock.MockStore) {