feat: add per-group AI budget table and endpoints (#25203)

Closes https://linear.app/codercom/issue/AIGOV-284/add-group-budgets-table-and-crud-api ## Summary Adds the `group_ai_budgets` table and the following endpoints: - `GET /api/v2/groups/{group}/ai/budget` - `PUT /api/v2/groups/{group}/ai/budget` - `DELETE /api/v2/groups/{group}/ai/budget` Each group may have at most one budget row. If no row exists, no budget is enforced. ### Feature gate Added `RequireFeatureMW(FeatureAIBridge)` on the `/ai/budget` sub-route. ## RBAC Authorization reuses `rbac.ResourceGroup` with the existing `.InOrganization(...).WithID(...)` scoping model. The `dbauthz` wrappers load the parent `groups` row and authorize against it. No new resource type is introduced. As a result, anyone with `group:update` permissions (Owner, OrgAdmin, or UserAdmin within the organization) can manage AI budgets for that group. ## Read access for group members `database.Group.RBACObject()` grants `policy.ActionRead` to all members of the group through the group ACL: ```go func (g Group) RBACObject() rbac.Object { return rbac.ResourceGroup.WithID(g.ID). InOrg(g.OrganizationID). // Group members can read the group. WithGroupACL(map[string][]policy.Action{ g.ID.String(): { policy.ActionRead, }, }) } ``` Because the `GET` endpoint authorizes against the same loaded `Group` object, any group member can call: ```text GET /api/v2/groups/{group}/ai/budget ``` `PUT` and `DELETE` remain admin-only. The group ACL grants only `ActionRead`, so write operations continue to require role-based `group:update` permissions. ## Alternative considered A dedicated `rbac.ResourceGroupAiBudget` resource would allow budget management to be separated from general group administration. We decided not to add that complexity for now.
2026-06-02 20:48:20 +00:00 · 2026-05-14 15:54:37 -04:00
parent d97f5ae2a6
commit 238968cfa0
25 changed files with 1040 additions and 0 deletions
@@ -697,3 +697,90 @@ func populatedAndConvertAIBridgeInterceptions(ctx context.Context, db database.S

 	return items, nil
 }
+
+// @Summary Get group AI budget
+// @ID get-group-ai-budget
+// @Security CoderSessionToken
+// @Produce json
+// @Tags Enterprise
+// @Param group path string true "Group ID" format(uuid)
+// @Success 200 {object} codersdk.GroupAIBudget
+// @Router /api/v2/groups/{group}/ai/budget [get]
+func (api *API) groupAIBudget(rw http.ResponseWriter, r *http.Request) {
+	ctx := r.Context()
+	group := httpmw.GroupParam(r)
+
+	budget, err := api.Database.GetGroupAIBudget(ctx, group.ID)
+	if httpapi.Is404Error(err) {
+		httpapi.ResourceNotFound(rw)
+		return
+	}
+	if err != nil {
+		api.Logger.Error(ctx, "get group AI budget", slog.Error(err))
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
+	httpapi.Write(ctx, rw, http.StatusOK, db2sdk.GroupAIBudget(budget))
+}
+
+// @Summary Upsert group AI budget
+// @ID upsert-group-ai-budget
+// @Security CoderSessionToken
+// @Accept json
+// @Produce json
+// @Tags Enterprise
+// @Param group path string true "Group ID" format(uuid)
+// @Param request body codersdk.UpsertGroupAIBudgetRequest true "Upsert group AI budget request"
+// @Success 200 {object} codersdk.GroupAIBudget
+// @Router /api/v2/groups/{group}/ai/budget [put]
+func (api *API) upsertGroupAIBudget(rw http.ResponseWriter, r *http.Request) {
+	ctx := r.Context()
+	group := httpmw.GroupParam(r)
+
+	var req codersdk.UpsertGroupAIBudgetRequest
+	if !httpapi.Read(ctx, rw, r, &req) {
+		return
+	}
+
+	budget, err := api.Database.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
+		GroupID:          group.ID,
+		SpendLimitMicros: req.SpendLimitMicros,
+	})
+	if httpapi.Is404Error(err) {
+		httpapi.ResourceNotFound(rw)
+		return
+	}
+	if err != nil {
+		api.Logger.Error(ctx, "upsert group AI budget", slog.Error(err))
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
+	httpapi.Write(ctx, rw, http.StatusOK, db2sdk.GroupAIBudget(budget))
+}
+
+// @Summary Delete group AI budget
+// @ID delete-group-ai-budget
+// @Security CoderSessionToken
+// @Tags Enterprise
+// @Param group path string true "Group ID" format(uuid)
+// @Success 204
+// @Router /api/v2/groups/{group}/ai/budget [delete]
+func (api *API) deleteGroupAIBudget(rw http.ResponseWriter, r *http.Request) {
+	ctx := r.Context()
+	group := httpmw.GroupParam(r)
+
+	_, err := api.Database.DeleteGroupAIBudget(ctx, group.ID)
+	if httpapi.Is404Error(err) {
+		httpapi.ResourceNotFound(rw)
+		return
+	}
+	if err != nil {
+		api.Logger.Error(ctx, "delete group AI budget", slog.Error(err))
+		httpapi.InternalServerError(rw, err)
+		return
+	}
+
+	rw.WriteHeader(http.StatusNoContent)
+}
@@ -2603,3 +2603,201 @@ func TestAIBridgeAllowBYOK(t *testing.T) {
 		})
 	}
 }
+
+func TestGroupAIBudget(t *testing.T) {
+	t.Parallel()
+
+	t.Run("Upsert", func(t *testing.T) {
+		t.Parallel()
+
+		adminClient, group := setupGroupAIBudgetTest(t)
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		// First upsert creates the budget.
+		newBudget, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{
+			SpendLimitMicros: 500_000_000,
+		})
+		require.NoError(t, err)
+		require.Equal(t, group.ID, newBudget.GroupID)
+		require.EqualValues(t, 500_000_000, newBudget.SpendLimitMicros)
+
+		// Second upsert updates the existing budget.
+		updatedBudget, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{
+			SpendLimitMicros: 1_000_000_000,
+		})
+		require.NoError(t, err)
+		require.EqualValues(t, 1_000_000_000, updatedBudget.SpendLimitMicros)
+
+		// GET returns the latest value.
+		currentBudget, err := adminClient.GroupAIBudget(ctx, group.ID)
+		require.NoError(t, err)
+		require.EqualValues(t, 1_000_000_000, currentBudget.SpendLimitMicros)
+	})
+
+	t.Run("GetWhenAbsent_404", func(t *testing.T) {
+		t.Parallel()
+
+		adminClient, group := setupGroupAIBudgetTest(t)
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		_, err := adminClient.GroupAIBudget(ctx, group.ID)
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
+	})
+
+	t.Run("DeleteWhenAbsent_404", func(t *testing.T) {
+		t.Parallel()
+
+		adminClient, group := setupGroupAIBudgetTest(t)
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		err := adminClient.DeleteGroupAIBudget(ctx, group.ID)
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
+	})
+
+	t.Run("DeleteWhenPresent", func(t *testing.T) {
+		t.Parallel()
+
+		adminClient, group := setupGroupAIBudgetTest(t)
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		_, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{
+			SpendLimitMicros: 500_000_000,
+		})
+		require.NoError(t, err)
+
+		require.NoError(t, adminClient.DeleteGroupAIBudget(ctx, group.ID))
+
+		_, err = adminClient.GroupAIBudget(ctx, group.ID)
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
+	})
+
+	t.Run("RejectsNegativeSpendLimit", func(t *testing.T) {
+		t.Parallel()
+
+		adminClient, group := setupGroupAIBudgetTest(t)
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		_, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{
+			SpendLimitMicros: -1,
+		})
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
+	})
+
+	t.Run("AcceptsZeroSpendLimitToBlock", func(t *testing.T) {
+		t.Parallel()
+
+		adminClient, group := setupGroupAIBudgetTest(t)
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		// 0 is a valid value: it blocks all spend for the group's members.
+		budget, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{
+			SpendLimitMicros: 0,
+		})
+		require.NoError(t, err)
+		require.EqualValues(t, 0, budget.SpendLimitMicros)
+	})
+
+	t.Run("UnknownGroup_404", func(t *testing.T) {
+		t.Parallel()
+
+		adminClient, _ := setupGroupAIBudgetTest(t)
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		_, err := adminClient.GroupAIBudget(ctx, uuid.New())
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
+	})
+
+	t.Run("GroupMemberCanReadButNotWrite", func(t *testing.T) {
+		t.Parallel()
+
+		dv := coderdtest.DeploymentValues(t)
+		dv.AI.BridgeConfig.Enabled = serpent.Bool(true)
+		ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{DeploymentValues: dv},
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureTemplateRBAC: 1,
+					codersdk.FeatureAIBridge:     1,
+				},
+			},
+		})
+		adminClient, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin())
+		memberClient, member := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID)
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+		group, err := adminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{
+			Name: "budget-group",
+		})
+		require.NoError(t, err)
+
+		// Add the member to the group so the Group.RBACObject ACL grants them read.
+		_, err = adminClient.PatchGroup(ctx, group.ID, codersdk.PatchGroupRequest{
+			AddUsers: []string{member.ID.String()},
+		})
+		require.NoError(t, err)
+
+		// Admin sets the budget so there is a row to read.
+		_, err = adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{
+			SpendLimitMicros: 500_000_000,
+		})
+		require.NoError(t, err)
+
+		// Group members can read the budget.
+		got, err := memberClient.GroupAIBudget(ctx, group.ID)
+		require.NoError(t, err)
+		require.EqualValues(t, 500_000_000, got.SpendLimitMicros)
+
+		// Group members cannot write the budget.
+		_, err = memberClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{
+			SpendLimitMicros: 1_000_000_000,
+		})
+		var sdkErr *codersdk.Error
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
+
+		// Group members cannot delete the budget.
+		err = memberClient.DeleteGroupAIBudget(ctx, group.ID)
+		require.ErrorAs(t, err, &sdkErr)
+		require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
+
+		// The failed upsert and delete left the budget untouched.
+		got, err = memberClient.GroupAIBudget(ctx, group.ID)
+		require.NoError(t, err)
+		require.EqualValues(t, 500_000_000, got.SpendLimitMicros)
+	})
+}
+
+// setupGroupAIBudgetTest returns an Admin client along with a newly created group inside it.
+func setupGroupAIBudgetTest(t *testing.T) (adminClient *codersdk.Client, group codersdk.Group) {
+	t.Helper()
+
+	dv := coderdtest.DeploymentValues(t)
+	dv.AI.BridgeConfig.Enabled = serpent.Bool(true)
+	ownerClient, owner := coderdenttest.New(t, &coderdenttest.Options{
+		Options: &coderdtest.Options{DeploymentValues: dv},
+		LicenseOptions: &coderdenttest.LicenseOptions{
+			Features: license.Features{
+				codersdk.FeatureTemplateRBAC: 1,
+				codersdk.FeatureAIBridge:     1,
+			},
+		},
+	})
+	adminClient, _ = coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID, rbac.RoleUserAdmin())
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	g, err := adminClient.CreateGroup(ctx, owner.OrganizationID, codersdk.CreateGroupRequest{
+		Name: "budget-test-group",
+	})
+	require.NoError(t, err)
+	return adminClient, g
+}
@@ -549,6 +549,13 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
 				r.Patch("/", api.patchGroup)
 				r.Delete("/", api.deleteGroup)
 				r.Get("/members", api.groupMembers)
+				r.Route("/ai/budget", func(r chi.Router) {
+					// AI cost controls are a paid feature (AI Governance add-on).
+					r.Use(api.RequireFeatureMW(codersdk.FeatureAIBridge))
+					r.Get("/", api.groupAIBudget)
+					r.Put("/", api.upsertGroupAIBudget)
+					r.Delete("/", api.deleteGroupAIBudget)
+				})
 			})
 		})
 		r.Route("/workspace-quota", func(r chi.Router) {