fix: ai cost control cap configurable AI spend limit (#27640)

## Problem

A configured AI spend limit was only validated as `gte=0`, with no upper
bound. The group spend query multiplies the per-member limit by the
number of attributed members, so a large enough limit overflows `bigint`
and fails the whole query, returning an error for every group in the
request rather than just the misconfigured one.

## Changes

- Add `MaxAISpendLimitMicros`, $1,000,000 per member per budget period.
- Reject group budgets and per-user overrides above the maximum with a
400 naming the limit.
- Bound both budget forms in the UI so they show the valid range before
submitting.

Follow-up
https://github.com/coder/coder/pull/27589#discussion_r3668956350
Depends on https://github.com/coder/coder/pull/27589

> [!NOTE]
> Initially generated by Claude Opus 5, modified and reviewed by
@ssncferreira

coderd/apidoc/docs.go

@@ -26320,6 +26320,7 @@ const docTemplate = `{
             "type": "object",
             "properties": {
                 "spend_limit_micros": {
+                    "description": "SpendLimitMicros must not exceed MaxAISpendLimitMicros.",
                     "type": "integer",
                     "minimum": 0
                 }
@@ -26337,6 +26338,7 @@ const docTemplate = `{
                     "format": "uuid"
                 },
                 "spend_limit_micros": {
+                    "description": "SpendLimitMicros must not exceed MaxAISpendLimitMicros.",
                     "type": "integer",
                     "minimum": 0
                 }

coderd/apidoc/swagger.json

@@ -24200,6 +24200,7 @@
 			"type": "object",
 			"properties": {
 				"spend_limit_micros": {
+					"description": "SpendLimitMicros must not exceed MaxAISpendLimitMicros.",
 					"type": "integer",
 					"minimum": 0
 				}
@@ -24215,6 +24216,7 @@
 					"format": "uuid"
 				},
 				"spend_limit_micros": {
+					"description": "SpendLimitMicros must not exceed MaxAISpendLimitMicros.",
 					"type": "integer",
 					"minimum": 0
 				}

codersdk/aibridge.go

@@ -15,6 +15,10 @@ import (
 	"github.com/coder/coder/v2/coderd/util/slice"
 )
 
+// MaxAISpendLimitMicros is the highest AI spend limit that can be configured,
+// $1,000,000 per member per budget period.
+const MaxAISpendLimitMicros int64 = 1_000_000_000_000
+
 // AIBudgetLimitSource identifies which tier produced the user's
 // effective budget limit.
 type AIBudgetLimitSource string
@@ -418,6 +422,7 @@ type GroupAIBudget struct {
 }
 
 type UpsertGroupAIBudgetRequest struct {
+	// SpendLimitMicros must not exceed MaxAISpendLimitMicros.
 	SpendLimitMicros int64 `json:"spend_limit_micros" validate:"gte=0"`
 }
 
@@ -485,8 +490,9 @@ type UserAIBudgetOverride struct {
 type UpsertUserAIBudgetOverrideRequest struct {
 	// GroupID is the group the user's spend is attributed to. The user must
 	// be a member of this group.
-	GroupID          uuid.UUID `json:"group_id" format:"uuid" validate:"required"`
-	SpendLimitMicros int64     `json:"spend_limit_micros" validate:"gte=0"`
+	GroupID uuid.UUID `json:"group_id" format:"uuid" validate:"required"`
+	// SpendLimitMicros must not exceed MaxAISpendLimitMicros.
+	SpendLimitMicros int64 `json:"spend_limit_micros" validate:"gte=0"`
 }
 
 // UserAIBudgetOverride returns the AI spend budget override configured for the given user.

docs/reference/api/schemas.md

@@ -14302,9 +14302,9 @@ If the schedule is empty, the user will be updated to use the default schedule.|
 
 ### Properties
 
-| Name                 | Type    | Required | Restrictions | Description |
-|----------------------|---------|----------|--------------|-------------|
-| `spend_limit_micros` | integer | false    |              |             |
+| Name                 | Type    | Required | Restrictions | Description                                               |
+|----------------------|---------|----------|--------------|-----------------------------------------------------------|
+| `spend_limit_micros` | integer | false    |              | Spend limit micros must not exceed MaxAISpendLimitMicros. |
 
 ## codersdk.UpsertUserAIBudgetOverrideRequest
 
@@ -14320,7 +14320,7 @@ If the schedule is empty, the user will be updated to use the default schedule.|
 | Name                 | Type    | Required | Restrictions | Description                                                                                       |
 |----------------------|---------|----------|--------------|---------------------------------------------------------------------------------------------------|
 | `group_id`           | string  | true     |              | Group ID is the group the user's spend is attributed to. The user must be a member of this group. |
-| `spend_limit_micros` | integer | false    |              |                                                                                                   |
+| `spend_limit_micros` | integer | false    |              | Spend limit micros must not exceed MaxAISpendLimitMicros.                                         |
 
 ## codersdk.UpsertWorkspaceAgentPortShareRequest
 

enterprise/coderd/aibridge.go

@@ -611,6 +611,22 @@ func (api *API) groupAIBudget(rw http.ResponseWriter, r *http.Request) {
 	httpapi.Write(ctx, rw, http.StatusOK, db2sdk.GroupAIBudget(groupBudget))
 }
 
+// validAISpendLimit reports whether the limit is within the configurable
+// maximum, writing a 400 when it is not.
+func validAISpendLimit(ctx context.Context, rw http.ResponseWriter, spendLimitMicros int64) bool {
+	if spendLimitMicros <= codersdk.MaxAISpendLimitMicros {
+		return true
+	}
+	httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+		Message: "Invalid AI spend limit.",
+		Validations: []codersdk.ValidationError{{
+			Field:  "spend_limit_micros",
+			Detail: fmt.Sprintf("Must not exceed %d.", codersdk.MaxAISpendLimitMicros),
+		}},
+	})
+	return false
+}
+
 // @Summary Upsert group AI budget
 // @ID upsert-group-ai-budget
 // @Security CoderSessionToken
@@ -640,6 +656,9 @@ func (api *API) upsertGroupAIBudget(rw http.ResponseWriter, r *http.Request) {
 	if !httpapi.Read(ctx, rw, r, &req) {
 		return
 	}
+	if !validAISpendLimit(ctx, rw, req.SpendLimitMicros) {
+		return
+	}
 
 	// Capture the existing budget (if any) so the audit log records the
 	// before-state. An absent row leaves aReq.Old as the zero value.
@@ -750,6 +769,9 @@ func (api *API) upsertUserAIBudgetOverride(rw http.ResponseWriter, r *http.Reque
 	if !httpapi.Read(ctx, rw, r, &req) {
 		return
 	}
+	if !validAISpendLimit(ctx, rw, req.SpendLimitMicros) {
+		return
+	}
 
 	// Look up the new group first so a missing or forbidden group_id
 	// returns 404. We also need the group for the audit log.

enterprise/coderd/aibridge_test.go

@@ -2289,6 +2289,40 @@ func TestGroupAIBudget(t *testing.T) {
 		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
 	})
 
+	t.Run("SpendLimitMaximum", func(t *testing.T) {
+		t.Parallel()
+
+		cases := []struct {
+			name      string
+			limit     int64
+			wantError bool
+		}{
+			{name: "AtMaximum", limit: codersdk.MaxAISpendLimitMicros},
+			{name: "AboveMaximum", limit: codersdk.MaxAISpendLimitMicros + 1, wantError: true},
+		}
+
+		for _, tc := range cases {
+			t.Run(tc.name, func(t *testing.T) {
+				t.Parallel()
+
+				adminClient, group := setupGroupAIBudgetTest(t)
+				ctx := testutil.Context(t, testutil.WaitLong)
+
+				budget, err := adminClient.UpsertGroupAIBudget(ctx, group.ID, codersdk.UpsertGroupAIBudgetRequest{
+					SpendLimitMicros: tc.limit,
+				})
+				if tc.wantError {
+					var sdkErr *codersdk.Error
+					require.ErrorAs(t, err, &sdkErr)
+					require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
+					return
+				}
+				require.NoError(t, err)
+				require.Equal(t, tc.limit, budget.SpendLimitMicros)
+			})
+		}
+	})
+
 	t.Run("AcceptsZeroSpendLimitToBlock", func(t *testing.T) {
 		t.Parallel()
 
@@ -2590,6 +2624,41 @@ func TestUserAIBudgetOverride(t *testing.T) {
 		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
 	})
 
+	t.Run("Upsert/SpendLimitMaximum", func(t *testing.T) {
+		t.Parallel()
+
+		cases := []struct {
+			name      string
+			limit     int64
+			wantError bool
+		}{
+			{name: "AtMaximum", limit: codersdk.MaxAISpendLimitMicros},
+			{name: "AboveMaximum", limit: codersdk.MaxAISpendLimitMicros + 1, wantError: true},
+		}
+
+		for _, tc := range cases {
+			t.Run(tc.name, func(t *testing.T) {
+				t.Parallel()
+
+				adminClient, targetUser, group := setupAICostControlTest(t, aiCostControlTestOptions{GroupName: "override-max-group"})
+				ctx := testutil.Context(t, testutil.WaitLong)
+
+				override, err := adminClient.UpsertUserAIBudgetOverride(ctx, targetUser.ID, codersdk.UpsertUserAIBudgetOverrideRequest{
+					GroupID:          group.ID,
+					SpendLimitMicros: tc.limit,
+				})
+				if tc.wantError {
+					var sdkErr *codersdk.Error
+					require.ErrorAs(t, err, &sdkErr)
+					require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
+					return
+				}
+				require.NoError(t, err)
+				require.Equal(t, tc.limit, override.SpendLimitMicros)
+			})
+		}
+	})
+
 	t.Run("Upsert/RejectsUnknownGroup", func(t *testing.T) {
 		t.Parallel()
 

site/src/api/typesGenerated.ts

@@ -6086,6 +6086,13 @@ export interface MatchedProvisioners {
 	readonly most_recently_seen?: string;
 }
 
+// From codersdk/aibridge.go
+/**
+ * MaxAISpendLimitMicros is the highest AI spend limit that can be configured,
+ * $1,000,000 per member per budget period.
+ */
+export const MaxAISpendLimitMicros = 1000000000000;
+
 // From codersdk/chats.go
 /**
  * MaxChatFileIDs is the maximum number of file IDs that can be
@@ -10218,6 +10225,9 @@ export interface UpsertChatUsageLimitOverrideRequest {
 
 // From codersdk/aibridge.go
 export interface UpsertGroupAIBudgetRequest {
+	/**
+	 * SpendLimitMicros must not exceed MaxAISpendLimitMicros.
+	 */
 	readonly spend_limit_micros: number;
 }
 
@@ -10228,6 +10238,9 @@ export interface UpsertUserAIBudgetOverrideRequest {
 	 * be a member of this group.
 	 */
 	readonly group_id: string;
+	/**
+	 * SpendLimitMicros must not exceed MaxAISpendLimitMicros.
+	 */
 	readonly spend_limit_micros: number;
 }
 

site/src/modules/groups.ts

@@ -1,10 +1,18 @@
-import type {
-	Group,
-	OrganizationMemberWithUserData,
-	ReducedUser,
-	User,
-	WorkspaceUser,
+import {
+	type Group,
+	MaxAISpendLimitMicros,
+	type OrganizationMemberWithUserData,
+	type ReducedUser,
+	type User,
+	type WorkspaceUser,
 } from "#/api/typesGenerated";
+import { MICROS_PER_DOLLAR, usdBudgetFormatter } from "#/utils/currency";
+
+/** Highest AI budget that can be configured for a group or member, in dollars. */
+export const maxAIBudgetDollars = MaxAISpendLimitMicros / MICROS_PER_DOLLAR;
+
+/** Shown when an entered AI budget falls outside the configurable range. */
+export const aiBudgetRangeError = `Enter an amount between 0 and ${usdBudgetFormatter.format(maxAIBudgetDollars)}.`;
 
 /**
  * Union of all user-like types that can be distinguished from Group.

site/src/pages/GroupsPage/GroupSettingsPageView.stories.tsx

@@ -1,5 +1,6 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
 import { expect, fn, userEvent, waitFor, within } from "storybook/test";
+import { maxAIBudgetDollars } from "#/modules/groups";
 import { MockGroup } from "#/testHelpers/entities";
 import GroupSettingsPageView from "./GroupSettingsPageView";
 
@@ -40,10 +41,8 @@ export const WithAIBudget: Story = {
 		await expect(canvas.getByLabelText("Monthly limit per member")).toHaveValue(
 			1000,
 		);
-		const helper = canvas.getByText(/month maximum/i);
-		await expect(helper).toHaveTextContent(
-			"$7,000/month maximum, based on 7 members.",
-		);
+		const helper = canvas.getByText(/month, based on/i);
+		await expect(helper).toHaveTextContent("$7,000/month, based on 7 members.");
 	},
 };
 
@@ -91,10 +90,30 @@ export const AIBudgetDecimal: Story = {
 	play: async ({ canvasElement }) => {
 		const canvas = within(canvasElement);
 		// Cents are kept when the amount is not a whole dollar.
-		const helper = canvas.getByText(/month maximum/i);
-		await expect(helper).toHaveTextContent(
-			"$99.99/month maximum, based on 1 member.",
-		);
+		const helper = canvas.getByText(/month, based on/i);
+		await expect(helper).toHaveTextContent("$99.99/month, based on 1 member.");
+	},
+};
+
+// A budget above the configurable maximum blocks saving.
+export const AIBudgetAboveMaximum: Story = {
+	args: {
+		showAISettings: true,
+		initialBudgetDollars: null,
+	},
+	play: async ({ canvasElement, args }) => {
+		const canvas = within(canvasElement);
+		const input = canvas.getByLabelText("Monthly limit per member");
+
+		await userEvent.type(input, String(maxAIBudgetDollars + 1));
+		// Blur to surface the error, matching the touched-then-validate flow.
+		await userEvent.tab();
+		await expect(
+			await canvas.findByText("Enter an amount between 0 and $1,000,000."),
+		).toBeInTheDocument();
+
+		await userEvent.click(canvas.getByRole("button", { name: "Save" }));
+		await expect(args.onSubmit).not.toHaveBeenCalled();
 	},
 };
 

site/src/pages/GroupsPage/GroupSettingsPageView.tsx

@@ -14,7 +14,11 @@ import {
 } from "#/components/InputGroup/InputGroup";
 import { Label } from "#/components/Label/Label";
 import { Spinner } from "#/components/Spinner/Spinner";
-import { isEveryoneGroup } from "#/modules/groups";
+import {
+	aiBudgetRangeError,
+	isEveryoneGroup,
+	maxAIBudgetDollars,
+} from "#/modules/groups";
 import { usdBudgetFormatter } from "#/utils/currency";
 import {
 	getFormHelpers,
@@ -34,10 +38,11 @@ type FormData = {
 const validationSchema = Yup.object({
 	name: nameValidator("Name"),
 	quota_allowance: Yup.number().required().min(0).integer(),
-	// Optional: empty is unlimited. A value must be zero or more; 0 disables.
+	// Optional: empty is unlimited. A value must be within the range; 0 disables.
 	monthly_budget_per_member: Yup.number()
 		.transform((value, original) => (original === "" ? undefined : value))
-		.min(0, "Enter an amount of zero or more."),
+		.min(0, aiBudgetRangeError)
+		.max(maxAIBudgetDollars, aiBudgetRangeError),
 });
 
 interface AIBudgetFeedbackProps {
@@ -94,7 +99,7 @@ const AIBudgetFeedback: FC<AIBudgetFeedbackProps> = ({
 				<span className="font-medium text-content-primary">
 					{usdBudgetFormatter.format(budgetAmount * memberCount)}
 				</span>
-				/month maximum, based on{" "}
+				/month, based on{" "}
 				<span className="font-medium text-content-primary">{memberCount}</span>{" "}
 				{memberCount === 1 ? "member" : "members"}.
 			</span>
@@ -244,6 +249,7 @@ const UpdateGroupForm: FC<UpdateGroupFormProps> = ({
 									onBlur={budgetField.onBlur}
 									type="number"
 									min="0"
+									max={maxAIBudgetDollars}
 									step="1"
 									placeholder="unlimited"
 									aria-invalid={budgetField.error}

site/src/pages/GroupsPage/UserAIBudgetOverrideDialog.stories.tsx

@@ -4,6 +4,7 @@ import { API } from "#/api/api";
 import { groupAIBudget, groupsForUser } from "#/api/queries/groups";
 import { getUserAIBudgetOverrideQueryKey } from "#/api/queries/users";
 import type { GroupAIBudget, UserAIBudgetOverride } from "#/api/typesGenerated";
+import { maxAIBudgetDollars } from "#/modules/groups";
 import { MockGroup, MockGroup2, MockUserMember } from "#/testHelpers/entities";
 import { UserAIBudgetOverrideDialog } from "./UserAIBudgetOverrideDialog";
 
@@ -122,12 +123,12 @@ export const Uncapped: Story = {
 			async () => {
 				const budgetInput = body.getByLabelText("Custom monthly budget");
 				await expect(
-					body.queryByText("Enter a monthly budget of 0 or more."),
+					body.queryByText("Enter an amount between 0 and $1,000,000."),
 				).not.toBeInTheDocument();
 				await userEvent.click(budgetInput);
 				await userEvent.tab();
 				await expect(
-					await body.findByText("Enter a monthly budget of 0 or more."),
+					await body.findByText("Enter an amount between 0 and $1,000,000."),
 				).toBeInTheDocument();
 			},
 		);
@@ -222,7 +223,7 @@ export const SubmitRequiresValueOrUncheck: Story = {
 			// Blur to surface the error, matching the touched-then-validate flow.
 			await userEvent.tab();
 			await expect(
-				await body.findByText("Enter a monthly budget of 0 or more."),
+				await body.findByText("Enter an amount between 0 and $1,000,000."),
 			).toBeInTheDocument();
 			await expect(updateButton).toBeDisabled();
 		});
@@ -244,6 +245,40 @@ export const SubmitRequiresValueOrUncheck: Story = {
 	},
 };
 
+// A budget above the configurable maximum blocks submit.
+export const SubmitBlockedAboveMaximum: Story = {
+	parameters: {
+		queries: [
+			{
+				key: getUserAIBudgetOverrideQueryKey(MockUserMember.id),
+				data: mockOverride,
+			},
+			...groupQueries,
+		],
+	},
+	play: async ({ step }) => {
+		const body = within(document.body);
+		const budgetInput = await body.findByLabelText("Custom monthly budget");
+		const updateButton = body.getByRole("button", { name: "Update" });
+
+		await step("the maximum itself is submittable", async () => {
+			await userEvent.clear(budgetInput);
+			await userEvent.type(budgetInput, String(maxAIBudgetDollars));
+			await expect(updateButton).toBeEnabled();
+		});
+
+		await step("one dollar above the maximum blocks submit", async () => {
+			await userEvent.clear(budgetInput);
+			await userEvent.type(budgetInput, String(maxAIBudgetDollars + 1));
+			await userEvent.tab();
+			await expect(
+				await body.findByText("Enter an amount between 0 and $1,000,000."),
+			).toBeInTheDocument();
+			await expect(updateButton).toBeDisabled();
+		});
+	},
+};
+
 export const Loading: Story = {
 	beforeEach: () => {
 		spyOn(API, "getUserAIBudgetOverride").mockReturnValue(

site/src/pages/GroupsPage/UserAIBudgetOverrideDialog.tsx

@@ -52,6 +52,7 @@ import {
 import { Label } from "#/components/Label/Label";
 import { Separator } from "#/components/Separator/Separator";
 import { Spinner } from "#/components/Spinner/Spinner";
+import { aiBudgetRangeError, maxAIBudgetDollars } from "#/modules/groups";
 import { cn } from "#/utils/cn";
 import {
 	dollarsToMicros,
@@ -213,9 +214,13 @@ const OverrideForm: FC<OverrideFormProps> = ({
 	const selectedGroup = groupOptions.find((g) => g.id === selectedGroupId);
 	const overrideGroup = groupOptions.find((g) => g.id === override?.group_id);
 
-	// A "0" budget is valid and disables AI; empty or negative is not.
+	// A "0" budget is valid and disables AI. Empty, negative, or above the
+	// configurable maximum is not.
 	const budgetAmount = Number(budgetDollars);
-	const budgetValid = budgetDollars.trim() !== "" && budgetAmount >= 0;
+	const budgetValid =
+		budgetDollars.trim() !== "" &&
+		budgetAmount >= 0 &&
+		budgetAmount <= maxAIBudgetDollars;
 	// Hold the error until the field is touched, so it doesn't flag immediately.
 	const budgetInvalid = overrideEnabled && budgetTouched && !budgetValid;
 	const budgetDisablesAI = budgetValid && budgetAmount === 0;
@@ -319,6 +324,7 @@ const OverrideForm: FC<OverrideFormProps> = ({
 								onBlur={() => setBudgetTouched(true)}
 								type="number"
 								min="0"
+								max={maxAIBudgetDollars}
 								step="1"
 								aria-invalid={budgetInvalid}
 								aria-describedby={
@@ -332,7 +338,7 @@ const OverrideForm: FC<OverrideFormProps> = ({
 								id={`${budgetId}-error`}
 								className="m-0 text-sm text-content-destructive"
 							>
-								Enter a monthly budget of 0 or more.
+								{aiBudgetRangeError}
 							</p>
 						)}
 					</div>