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
@@ -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)
+	})
+}