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
+8
View File
@@ -83,6 +83,7 @@ import (
"github.com/coder/coder/v2/coderd/schedule"
"github.com/coder/coder/v2/coderd/telemetry"
"github.com/coder/coder/v2/coderd/updatecheck"
"github.com/coder/coder/v2/coderd/usage"
"github.com/coder/coder/v2/coderd/util/namesgenerator"
"github.com/coder/coder/v2/coderd/util/ptr"
"github.com/coder/coder/v2/coderd/webpush"
@@ -190,6 +191,7 @@ type Options struct {
TelemetryReporter telemetry.Reporter
ProvisionerdServerMetrics *provisionerdserver.Metrics
UsageInserter usage.Inserter
}
// New constructs a codersdk client connected to an in-memory API instance.
@@ -270,6 +272,11 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
}
}
var usageInserter *atomic.Pointer[usage.Inserter]
if options.UsageInserter != nil {
usageInserter = &atomic.Pointer[usage.Inserter]{}
usageInserter.Store(&options.UsageInserter)
}
if options.Database == nil {
options.Database, options.Pubsub = dbtestutil.NewDB(t)
}
@@ -563,6 +570,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
Database: options.Database,
Pubsub: options.Pubsub,
ExternalAuthConfigs: options.ExternalAuthConfigs,
UsageInserter: usageInserter,
Auditor: options.Auditor,
ConnectionLogger: options.ConnectionLogger,
+44
View File
@@ -0,0 +1,44 @@
package coderdtest
import (
"context"
"sync"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/usage"
"github.com/coder/coder/v2/coderd/usage/usagetypes"
)
var _ usage.Inserter = (*UsageInserter)(nil)
type UsageInserter struct {
sync.Mutex
events []usagetypes.DiscreteEvent
}
func NewUsageInserter() *UsageInserter {
return &UsageInserter{
events: []usagetypes.DiscreteEvent{},
}
}
func (u *UsageInserter) InsertDiscreteUsageEvent(_ context.Context, _ database.Store, event usagetypes.DiscreteEvent) error {
u.Lock()
defer u.Unlock()
u.events = append(u.events, event)
return nil
}
func (u *UsageInserter) GetEvents() []usagetypes.DiscreteEvent {
u.Lock()
defer u.Unlock()
eventsCopy := make([]usagetypes.DiscreteEvent, len(u.events))
copy(eventsCopy, u.events)
return eventsCopy
}
func (u *UsageInserter) Reset() {
u.Lock()
defer u.Unlock()
u.events = []usagetypes.DiscreteEvent{}
}
+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) {