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**.
2026-09-24 15:04:27 +08:00 · 2026-01-27 19:01:17 -06:00
parent 3eeeabfd68
commit 799b190dee
7 changed files with 225 additions and 19 deletions
@@ -975,7 +975,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, transition database.WorkspaceTransition) (wsbuilder.UsageCheckResponse, error) {
+func (api *API) CheckBuildUsage(ctx context.Context, store database.Store, templateVersion *database.TemplateVersion, task *database.Task, 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 {
@@ -988,7 +988,7 @@ func (api *API) CheckBuildUsage(ctx context.Context, store database.Store, templ
 		}
 	}

-	resp, err := api.checkAIBuildUsage(ctx, store, templateVersion, transition)
+	resp, err := api.checkAIBuildUsage(ctx, store, task, transition)
 	if err != nil {
 		return wsbuilder.UsageCheckResponse{}, err
 	}
@@ -1001,14 +1001,14 @@ func (api *API) CheckBuildUsage(ctx context.Context, store database.Store, templ

 // 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) {
+func (api *API) checkAIBuildUsage(ctx context.Context, store database.Store, task *database.Task, 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 {
+	if task == nil {
 		return wsbuilder.UsageCheckResponse{Permitted: true}, nil
 	}

@@ -772,6 +772,10 @@ func TestCheckBuildUsage_SkipsAIForNonStartTransitions(t *testing.T) {
 		HasExternalAgent: sql.NullBool{Valid: true, Bool: false},
 	}

+	task := &database.Task{
+		TemplateVersionID: tv.ID,
+	}
+
 	// Mock DB: expect exactly one count call for the "start" transition.
 	mDB := dbmock.NewMockStore(ctrl)
 	mDB.EXPECT().
@@ -782,18 +786,18 @@ func TestCheckBuildUsage_SkipsAIForNonStartTransitions(t *testing.T) {
 	ctx := context.Background()

 	// Start transition: should be not permitted due to limit breach.
-	startResp, err := eapi.CheckBuildUsage(ctx, mDB, tv, database.WorkspaceTransitionStart)
+	startResp, err := eapi.CheckBuildUsage(ctx, mDB, tv, task, database.WorkspaceTransitionStart)
 	require.NoError(t, err)
 	require.False(t, startResp.Permitted)
 	require.Contains(t, startResp.Message, "breached the managed agent limit")

 	// Stop transition: should be permitted and must not trigger additional DB calls.
-	stopResp, err := eapi.CheckBuildUsage(ctx, mDB, tv, database.WorkspaceTransitionStop)
+	stopResp, err := eapi.CheckBuildUsage(ctx, mDB, tv, task, database.WorkspaceTransitionStop)
 	require.NoError(t, err)
 	require.True(t, stopResp.Permitted)

 	// Delete transition: should be permitted and must not trigger additional DB calls.
-	deleteResp, err := eapi.CheckBuildUsage(ctx, mDB, tv, database.WorkspaceTransitionDelete)
+	deleteResp, err := eapi.CheckBuildUsage(ctx, mDB, tv, task, database.WorkspaceTransitionDelete)
 	require.NoError(t, err)
 	require.True(t, deleteResp.Permitted)
 }
@@ -4705,3 +4705,121 @@ func TestWorkspacesSharedWith(t *testing.T) {
 		assert.Equal(t, "/emojis/1f60d.png", groupActor.AvatarURL)
 	})
 }
+
+//nolint:tparallel,paralleltest // Sub tests need to run sequentially.
+func TestWorkspaceAITask(t *testing.T) {
+	t.Parallel()
+
+	usage := coderdtest.NewUsageInserter()
+	owner, _, first := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{
+		Options: &coderdtest.Options{
+			UsageInserter:            usage,
+			IncludeProvisionerDaemon: true,
+		},
+		LicenseOptions: (&coderdenttest.LicenseOptions{
+			Features: license.Features{
+				codersdk.FeatureTemplateRBAC: 1,
+			},
+		}).ManagedAgentLimit(10, 20),
+	})
+
+	client, _ := coderdtest.CreateAnotherUser(t, owner, first.OrganizationID,
+		rbac.RoleTemplateAdmin(), rbac.RoleUserAdmin())
+
+	graphWithTask := []*proto.Response{{
+		Type: &proto.Response_Graph{
+			Graph: &proto.GraphComplete{
+				Error:                 "",
+				Timings:               nil,
+				Resources:             nil,
+				Parameters:            nil,
+				ExternalAuthProviders: nil,
+				Presets:               nil,
+				HasAiTasks:            true,
+				AiTasks: []*proto.AITask{
+					{
+						Id:         "test",
+						SidebarApp: nil,
+						AppId:      "test",
+					},
+				},
+				HasExternalAgents: false,
+			},
+		},
+	}}
+	planWithTask := []*proto.Response{{
+		Type: &proto.Response_Plan{
+			Plan: &proto.PlanComplete{
+				Plan:        []byte("{}"),
+				AiTaskCount: 1,
+			},
+		},
+	}}
+
+	t.Run("CreateWorkspaceWithTaskNormally", func(t *testing.T) {
+		// Creating a workspace that has agentic tasks, but is not launced via task
+		// should not count towards the usage.
+		t.Cleanup(usage.Reset)
+		version := coderdtest.CreateTemplateVersion(t, client, first.OrganizationID, &echo.Responses{
+			Parse:          echo.ParseComplete,
+			ProvisionInit:  echo.InitComplete,
+			ProvisionPlan:  planWithTask,
+			ProvisionApply: echo.ApplyComplete,
+			ProvisionGraph: graphWithTask,
+		})
+		_ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, first.OrganizationID, version.ID)
+		wrk := coderdtest.CreateWorkspace(t, client, template.ID)
+		build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, wrk.LatestBuild.ID)
+		require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
+		require.Len(t, usage.GetEvents(), 0)
+	})
+
+	t.Run("CreateTaskWorkspace", func(t *testing.T) {
+		ctx := testutil.Context(t, testutil.WaitMedium)
+		t.Cleanup(usage.Reset)
+		version := coderdtest.CreateTemplateVersion(t, client, first.OrganizationID, &echo.Responses{
+			Parse:          echo.ParseComplete,
+			ProvisionInit:  echo.InitComplete,
+			ProvisionPlan:  planWithTask,
+			ProvisionApply: echo.ApplyComplete,
+			ProvisionGraph: graphWithTask,
+		})
+		_ = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, first.OrganizationID, version.ID)
+
+		task, err := client.CreateTask(ctx, codersdk.Me, codersdk.CreateTaskRequest{
+			TemplateVersionID: template.ActiveVersionID,
+			Name:              "istask",
+		})
+		require.NoError(t, err)
+
+		wrk, err := client.Workspace(ctx, task.WorkspaceID.UUID)
+		require.NoError(t, err)
+
+		build := coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, wrk.LatestBuild.ID)
+		require.Equal(t, codersdk.WorkspaceStatusRunning, build.Status)
+		require.Len(t, usage.GetEvents(), 1)
+
+		usage.Reset() // Clean slate for easy checks
+		// Stopping the workspace should not create additional usage.
+		build, err = client.CreateWorkspaceBuild(ctx, wrk.ID, codersdk.CreateWorkspaceBuildRequest{
+			TemplateVersionID: wrk.LatestBuild.TemplateVersionID,
+			Transition:        codersdk.WorkspaceTransitionStop,
+		})
+		require.NoError(t, err)
+		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
+		require.Len(t, usage.GetEvents(), 0)
+
+		usage.Reset() // Clean slate for easy checks
+		// Starting the workspace manually **WILL** create usage, as it's
+		// still a task workspace.
+		build, err = client.CreateWorkspaceBuild(ctx, wrk.ID, codersdk.CreateWorkspaceBuildRequest{
+			TemplateVersionID: wrk.LatestBuild.TemplateVersionID,
+			Transition:        codersdk.WorkspaceTransitionStart,
+		})
+		require.NoError(t, err)
+		coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, build.ID)
+		require.Len(t, usage.GetEvents(), 1)
+	})
+}