fix(site): refresh group member budgets after overrides (#27553)

Group member budget rows remained stale after saving or deleting a user override because the cached member-spend query was not refreshed.

Invalidate the affected user's override query and only cached group-member spend queries whose user ID list contains that user. This refreshes relevant rows without invalidating unrelated groups.

site/src/api/queries/groups.test.ts

@@ -0,0 +1,35 @@
+import { QueryClient } from "react-query";
+import { describe, expect, it } from "vitest";
+import {
+	getGroupMembersAISpendQueryKey,
+	invalidateGroupMembersAISpend,
+} from "./groups";
+
+describe("invalidateGroupMembersAISpend", () => {
+	it("invalidates only group member spend queries containing the user", async () => {
+		const queryClient = new QueryClient();
+		const userId = "user-1";
+		const spendWithUser = getGroupMembersAISpendQueryKey("group-1", [
+			"user-2",
+			userId,
+		]);
+		const spendWithoutUser = getGroupMembersAISpendQueryKey("group-2", [
+			"user-2",
+		]);
+		const otherGroupQuery = ["group", "group-1"];
+
+		queryClient.setQueryData(spendWithUser, {});
+		queryClient.setQueryData(spendWithoutUser, {});
+		queryClient.setQueryData(otherGroupQuery, {});
+
+		await invalidateGroupMembersAISpend(queryClient, userId);
+
+		expect(queryClient.getQueryState(spendWithUser)?.isInvalidated).toBe(true);
+		expect(queryClient.getQueryState(spendWithoutUser)?.isInvalidated).toBe(
+			false,
+		);
+		expect(queryClient.getQueryState(otherGroupQuery)?.isInvalidated).toBe(
+			false,
+		);
+	});
+});

site/src/api/queries/groups.ts

@@ -64,6 +64,25 @@ export const getGroupMembersAISpendQueryKey = (
 	userIds: readonly string[],
 ) => ["group", groupId, "members", "aiSpend", [...userIds].sort()];
 
+const isGroupMembersAISpendQueryKey = (
+	queryKey: readonly unknown[],
+	userId: string,
+): boolean =>
+	queryKey[0] === "group" &&
+	queryKey[2] === "members" &&
+	queryKey[3] === "aiSpend" &&
+	Array.isArray(queryKey[4]) &&
+	queryKey[4].includes(userId);
+
+export const invalidateGroupMembersAISpend = (
+	queryClient: QueryClient,
+	userId: string,
+) =>
+	queryClient.invalidateQueries({
+		queryKey: ["group"],
+		predicate: (query) => isGroupMembersAISpendQueryKey(query.queryKey, userId),
+	});
+
 export const groupMembersAISpend = (
 	groupId: string,
 	userIds: readonly string[],

site/src/api/queries/users.ts

@@ -31,6 +31,7 @@ import {
 import type { UsePaginatedQueryOptions } from "#/hooks/usePaginatedQuery";
 import { prepareQuery } from "#/utils/filters";
 import { getAuthorizationKey } from "./authCheck";
+import { invalidateGroupMembersAISpend } from "./groups";
 import { cachedQuery } from "./util";
 
 export function usersKey(req: UsersRequest) {
@@ -203,6 +204,17 @@ export const userAIBudgetOverride = (
 	};
 };
 
+const invalidateUserAIBudgetQueries = (
+	queryClient: QueryClient,
+	userId: string,
+) =>
+	Promise.all([
+		queryClient.invalidateQueries({
+			queryKey: getUserAIBudgetOverrideQueryKey(userId),
+		}),
+		invalidateGroupMembersAISpend(queryClient, userId),
+	]);
+
 export const saveUserAIBudgetOverride = (
 	queryClient: QueryClient,
 	userId: string,
@@ -211,9 +223,7 @@ export const saveUserAIBudgetOverride = (
 		mutationFn: (request: UpsertUserAIBudgetOverrideRequest) =>
 			API.upsertUserAIBudgetOverride(userId, request),
 		onSuccess: async () => {
-			await queryClient.invalidateQueries({
-				queryKey: getUserAIBudgetOverrideQueryKey(userId),
-			});
+			await invalidateUserAIBudgetQueries(queryClient, userId);
 		},
 	};
 };
@@ -225,9 +235,7 @@ export const deleteUserAIBudgetOverride = (
 	return {
 		mutationFn: () => API.deleteUserAIBudgetOverride(userId),
 		onSuccess: async () => {
-			await queryClient.invalidateQueries({
-				queryKey: getUserAIBudgetOverrideQueryKey(userId),
-			});
+			await invalidateUserAIBudgetQueries(queryClient, userId);
 		},
 	};
 };

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

@@ -1,5 +1,5 @@
 import type { Meta, StoryObj } from "@storybook/react-vite";
-import { expect, spyOn, userEvent, within } from "storybook/test";
+import { expect, spyOn, userEvent, waitFor, within } from "storybook/test";
 import {
 	reactRouterOutlet,
 	reactRouterParameters,
@@ -24,6 +24,7 @@ import type {
 	GroupMemberAISpend,
 	GroupMembersAISpend,
 	ReducedUser,
+	UserAIBudgetOverride,
 	UserAISpendStatus,
 } from "#/api/typesGenerated";
 import {
@@ -501,6 +502,140 @@ export const WithMemberAIBudgetInAnotherOrg: Story = {
 	},
 };
 
+const mockUserBudgetOverride: UserAIBudgetOverride = {
+	user_id: MockUserOwner.id,
+	group_id: MockGroupWithoutMembers.id,
+	spend_limit_micros: 12_000_000_000,
+	created_at: "2026-06-01T00:00:00Z",
+	updated_at: "2026-06-01T00:00:00Z",
+};
+
+const mockOwnerSpend: GroupMemberAISpend = {
+	...mockSpend,
+	user_id: MockUserOwner.id,
+};
+
+const mockOwnerOverrideSpend: GroupMemberAISpend = {
+	...mockOwnerSpend,
+	group_budget: {
+		spend_limit_micros: mockUserBudgetOverride.spend_limit_micros,
+		limit_source: "user_override",
+	},
+};
+
+export const SaveMemberAIBudgetRefreshesRow: Story = {
+	beforeEach: () => {
+		spyOn(API, "upsertUserAIBudgetOverride").mockResolvedValue(
+			mockUserBudgetOverride,
+		);
+		spyOn(API, "getUserAIBudgetOverride").mockResolvedValue(
+			mockUserBudgetOverride,
+		);
+		spyOn(API, "getGroupMembersAISpend").mockResolvedValue(
+			membersSpendQuery([mockOwnerOverrideSpend]).data,
+		);
+	},
+	parameters: {
+		features: ["aibridge"],
+		experiments: ["ai-gateway-cost-control"],
+		queries: [
+			groupQuery(MockGroupWithoutMembers),
+			groupMembersQuery({ users: [MockUserOwner], count: 1 }),
+			membersSpendQuery([mockOwnerSpend]),
+			permissionsQuery({ canUpdateGroup: true }),
+			{ key: meAISpendKey, data: mockUserAISpend },
+			{ key: getUserAIBudgetOverrideQueryKey(MockUserOwner.id), data: null },
+			{
+				key: getGroupsForUserQueryKey(
+					MockUserOwner.id,
+					MockGroupWithoutMembers.organization_id,
+				),
+				data: [MockGroup2],
+			},
+			{
+				key: groupAIBudget(MockGroupWithoutMembers.id).queryKey,
+				data: mockGroupBudget,
+			},
+		],
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		const body = within(document.body);
+		const budgetCell = await canvas.findByTestId(
+			`member-ai-budget-${MockUserOwner.id}`,
+		);
+		await expect(budgetCell).toHaveTextContent("Group limit $9,000");
+
+		await userEvent.click(canvas.getByRole("button", { name: "Open menu" }));
+		await userEvent.click(
+			await body.findByRole("menuitem", { name: "Manage AI budget" }),
+		);
+		await userEvent.click(await body.findByText("Override group budget"));
+		const input = body.getByLabelText("Custom monthly budget");
+		await userEvent.clear(input);
+		await userEvent.type(input, "12000");
+		await userEvent.click(body.getByRole("button", { name: "Update" }));
+
+		await waitFor(() =>
+			expect(budgetCell).toHaveTextContent("Custom limit $12,000"),
+		);
+	},
+};
+
+export const DeleteMemberAIBudgetRefreshesRow: Story = {
+	beforeEach: () => {
+		spyOn(API, "deleteUserAIBudgetOverride").mockResolvedValue();
+		spyOn(API, "getGroupMembersAISpend").mockResolvedValue(
+			membersSpendQuery([mockOwnerSpend]).data,
+		);
+	},
+	parameters: {
+		features: ["aibridge"],
+		experiments: ["ai-gateway-cost-control"],
+		queries: [
+			groupQuery(MockGroupWithoutMembers),
+			groupMembersQuery({ users: [MockUserOwner], count: 1 }),
+			membersSpendQuery([mockOwnerOverrideSpend]),
+			permissionsQuery({ canUpdateGroup: true }),
+			{ key: meAISpendKey, data: mockUserAISpend },
+			{
+				key: getUserAIBudgetOverrideQueryKey(MockUserOwner.id),
+				data: mockUserBudgetOverride,
+			},
+			{
+				key: getGroupsForUserQueryKey(
+					MockUserOwner.id,
+					MockGroupWithoutMembers.organization_id,
+				),
+				data: [MockGroup2],
+			},
+			{
+				key: groupAIBudget(MockGroupWithoutMembers.id).queryKey,
+				data: mockGroupBudget,
+			},
+		],
+	},
+	play: async ({ canvasElement }) => {
+		const canvas = within(canvasElement);
+		const body = within(document.body);
+		const budgetCell = await canvas.findByTestId(
+			`member-ai-budget-${MockUserOwner.id}`,
+		);
+		await expect(budgetCell).toHaveTextContent("Custom limit $12,000");
+
+		await userEvent.click(canvas.getByRole("button", { name: "Open menu" }));
+		await userEvent.click(
+			await body.findByRole("menuitem", { name: "Manage AI budget" }),
+		);
+		await userEvent.click(body.getByRole("checkbox"));
+		await userEvent.click(body.getByRole("button", { name: "Update" }));
+
+		await waitFor(() =>
+			expect(budgetCell).toHaveTextContent("Group limit $9,000"),
+		);
+	},
+};
+
 export const OpenAIBudgetForCurrentGroupMember: Story = {
 	parameters: {
 		features: ["aibridge"],