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

feat: add managed agent license limit checks (#18937)

Browse Source 
- Adds a query for counting managed agent workspace builds between two
timestamps
- The "Actual" field in the feature entitlement for managed agents is
now populated with the value read from the database
- The wsbuilder package now validates AI agent usage against the limit
when a license is installed

Closes coder/internal#777
This commit is contained in:
Dean Sheather
2025-07-22 13:39:26 +10:00
committed by GitHub
parent aa1a985381
commit 9a6dd73f68
24 changed files with 586 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/wsbuilder/wsbuilder.go
+42 -3
View File
@@ -56,6 +56,7 @@ type Builder struct {
logLevel string
deploymentValues *codersdk.DeploymentValues
experiments codersdk.Experiments
usageChecker UsageChecker
richParameterValues []codersdk.WorkspaceBuildParameter
initiator uuid.UUID
@@ -89,7 +90,24 @@ type Builder struct {
verifyNoLegacyParametersOnce bool
}
type Option func(Builder) Builder
type UsageChecker interface {
CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (UsageCheckResponse, error)
}
type UsageCheckResponse struct {
Permitted bool
Message string
}
type NoopUsageChecker struct{}
var _ UsageChecker = NoopUsageChecker{}
func (NoopUsageChecker) CheckBuildUsage(_ context.Context, _ database.Store, _ *database.TemplateVersion) (UsageCheckResponse, error) {
return UsageCheckResponse{
Permitted: true,
}, nil
}
// versionTarget expresses how to determine the template version for the build.
//
@@ -121,8 +139,8 @@ type stateTarget struct {
explicit *[]byte
}
func New(w database.Workspace, t database.WorkspaceTransition) Builder {
return Builder{workspace: w, trans: t}
func New(w database.Workspace, t database.WorkspaceTransition, uc UsageChecker) Builder {
return Builder{workspace: w, trans: t, usageChecker: uc}
}
// Methods that customize the build are public, have a struct receiver and return a new Builder.
@@ -321,6 +339,10 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
if err != nil {
return nil, nil, nil, err
}
err = b.checkUsage()
if err != nil {
return nil, nil, nil, err
}
err = b.checkRunningBuild()
if err != nil {
return nil, nil, nil, err
@@ -1253,6 +1275,23 @@ func (b *Builder) checkTemplateJobStatus() error {
return nil
}
func (b *Builder) checkUsage() error {
templateVersion, err := b.getTemplateVersion()
if err != nil {
return BuildError{http.StatusInternalServerError, "Failed to fetch template version", err}
}
resp, err := b.usageChecker.CheckBuildUsage(b.ctx, b.store, templateVersion)
if err != nil {
return BuildError{http.StatusInternalServerError, "Failed to check build usage", err}
}
if !resp.Permitted {
return BuildError{http.StatusForbidden, "Build is not permitted: " + resp.Message, nil}
}
return nil
}
func (b *Builder) checkRunningBuild() error {
job, err := b.getLastBuildJob()
if xerrors.Is(err, sql.ErrNoRows) {
coderd/wsbuilder/wsbuilder_test.go
+149 -23
View File
@@ -5,30 +5,30 @@ import (
"database/sql"
"encoding/json"
"net/http"
"sync/atomic"
"testing"
"time"
"github.com/prometheus/client_golang/prometheus"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/files"
"github.com/coder/coder/v2/coderd/httpapi/httperror"
"github.com/coder/coder/v2/provisionersdk"
"github.com/google/uuid"
"github.com/prometheus/client_golang/prometheus"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/propagation"
"go.uber.org/mock/gomock"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/audit"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbmock"
"github.com/coder/coder/v2/coderd/database/dbtime"
"github.com/coder/coder/v2/coderd/files"
"github.com/coder/coder/v2/coderd/httpapi/httperror"
"github.com/coder/coder/v2/coderd/provisionerdserver"
"github.com/coder/coder/v2/coderd/wsbuilder"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/provisionersdk"
)
var (
@@ -102,7 +102,7 @@ func TestBuilder_NoOptions(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart)
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{})
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
req.NoError(err)
@@ -142,7 +142,8 @@ func TestBuilder_Initiator(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).Initiator(otherUserID)
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
Initiator(otherUserID)
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
req.NoError(err)
@@ -188,7 +189,8 @@ func TestBuilder_Baggage(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).Initiator(otherUserID)
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
Initiator(otherUserID)
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{IP: "127.0.0.1"})
req.NoError(err)
@@ -227,7 +229,8 @@ func TestBuilder_Reason(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).Reason(database.BuildReasonAutostart)
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
Reason(database.BuildReasonAutostart)
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
req.NoError(err)
@@ -271,7 +274,8 @@ func TestBuilder_ActiveVersion(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).ActiveVersion()
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
ActiveVersion()
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
req.NoError(err)
@@ -386,7 +390,8 @@ func TestWorkspaceBuildWithTags(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).RichParameterValues(buildParameters)
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
RichParameterValues(buildParameters)
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
req.NoError(err)
@@ -469,7 +474,8 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).RichParameterValues(nextBuildParameters)
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
RichParameterValues(nextBuildParameters)
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
req.NoError(err)
@@ -517,7 +523,8 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).RichParameterValues(nextBuildParameters)
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
RichParameterValues(nextBuildParameters)
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
req.NoError(err)
@@ -555,7 +562,8 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart)
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{})
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
bldErr := wsbuilder.BuildError{}
req.ErrorAs(err, &bldErr)
@@ -591,7 +599,8 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).RichParameterValues(nextBuildParameters)
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
RichParameterValues(nextBuildParameters)
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
bldErr := wsbuilder.BuildError{}
@@ -656,7 +665,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
RichParameterValues(nextBuildParameters).
VersionID(activeVersionID)
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
@@ -720,7 +729,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
RichParameterValues(nextBuildParameters).
VersionID(activeVersionID)
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
@@ -782,7 +791,7 @@ func TestWorkspaceBuildWithRichParameters(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
RichParameterValues(nextBuildParameters).
VersionID(activeVersionID)
// nolint: dogsled
@@ -849,7 +858,7 @@ func TestWorkspaceBuildWithPreset(t *testing.T) {
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart).
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, wsbuilder.NoopUsageChecker{}).
ActiveVersion().
TemplateVersionPresetID(presetID)
// nolint: dogsled
@@ -916,7 +925,7 @@ func TestWorkspaceBuildDeleteOrphan(t *testing.T) {
)
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionDelete).Orphan()
uut := wsbuilder.New(ws, database.WorkspaceTransitionDelete, wsbuilder.NoopUsageChecker{}).Orphan()
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
// nolint: dogsled
@@ -993,7 +1002,7 @@ func TestWorkspaceBuildDeleteOrphan(t *testing.T) {
)
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionDelete).Orphan()
uut := wsbuilder.New(ws, database.WorkspaceTransitionDelete, wsbuilder.NoopUsageChecker{}).Orphan()
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
@@ -1001,6 +1010,115 @@ func TestWorkspaceBuildDeleteOrphan(t *testing.T) {
})
}
func TestWorkspaceBuildUsageChecker(t *testing.T) {
t.Parallel()
t.Run("Permitted", func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
var calls int64
fakeUsageChecker := &fakeUsageChecker{
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error) {
atomic.AddInt64(&calls, 1)
return wsbuilder.UsageCheckResponse{Permitted: true}, nil
},
}
mDB := expectDB(t,
// Inputs
withTemplate,
withInactiveVersion(nil),
withLastBuildFound,
withTemplateVersionVariables(inactiveVersionID, nil),
withRichParameters(nil),
withParameterSchemas(inactiveJobID, nil),
withWorkspaceTags(inactiveVersionID, nil),
withProvisionerDaemons([]database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow{}),
// Outputs
expectProvisionerJob(func(job database.InsertProvisionerJobParams) {}),
withInTx,
expectBuild(func(bld database.InsertWorkspaceBuildParams) {}),
withBuild,
expectBuildParameters(func(params database.InsertWorkspaceBuildParametersParams) {}),
)
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, fakeUsageChecker)
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
require.NoError(t, err)
require.EqualValues(t, 1, calls)
})
// The failure cases are mostly identical from a test perspective.
const message = "fake test message"
cases := []struct {
name string
response wsbuilder.UsageCheckResponse
responseErr error
assertions func(t *testing.T, err error)
}{
{
name: "NotPermitted",
response: wsbuilder.UsageCheckResponse{
Permitted: false,
Message: message,
},
assertions: func(t *testing.T, err error) {
require.ErrorContains(t, err, message)
var buildErr wsbuilder.BuildError
require.ErrorAs(t, err, &buildErr)
require.Equal(t, http.StatusForbidden, buildErr.Status)
},
},
{
name: "Error",
responseErr: xerrors.New("fake error"),
assertions: func(t *testing.T, err error) {
require.ErrorContains(t, err, "fake error")
require.ErrorAs(t, err, &wsbuilder.BuildError{})
},
},
}
for _, c := range cases {
c := c
t.Run(c.name, func(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
var calls int64
fakeUsageChecker := &fakeUsageChecker{
checkBuildUsageFunc: func(_ context.Context, _ database.Store, templateVersion *database.TemplateVersion) (wsbuilder.UsageCheckResponse, error) {
atomic.AddInt64(&calls, 1)
return c.response, c.responseErr
},
}
mDB := expectDB(t,
withTemplate,
withInactiveVersionNoParams(),
)
fc := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
ws := database.Workspace{ID: workspaceID, TemplateID: templateID, OwnerID: userID}
uut := wsbuilder.New(ws, database.WorkspaceTransitionStart, fakeUsageChecker).
VersionID(inactiveVersionID)
// nolint: dogsled
_, _, _, err := uut.Build(ctx, mDB, fc, nil, audit.WorkspaceBuildBaggage{})
c.assertions(t, err)
require.EqualValues(t, 1, calls)
})
}
}
func TestWsbuildError(t *testing.T) {
t.Parallel()
@@ -1366,3 +1484,11 @@ func withProvisionerDaemons(provisionerDaemons []database.GetEligibleProvisioner
mTx.EXPECT().GetEligibleProvisionerDaemonsByProvisionerJobIDs(gomock.Any(), gomock.Any()).Return(provisionerDaemons, nil)
}
}
type fakeUsageChecker struct {
checkBuildUsageFunc func(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion) (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)
}