Thanks to visit codestin.com
Credit goes to github.com

Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
EhabY merged 3 commits into from
Jul 29, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7a59826
fix(site): refresh member budgets after overrides
EhabY Jul 27, 2026
5233938
fix(site): target member budget cache invalidation
EhabY Jul 28, 2026
018f91b
refactor(site): extract AI budget invalidation helpers
EhabY Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions site/src/api/queries/groups.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
);
});
});
19 changes: 19 additions & 0 deletions site/src/api/queries/groups.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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[],
Expand Down
20 changes: 14 additions & 6 deletions site/src/api/queries/users.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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,
Expand All @@ -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);
},
};
};
Expand All @@ -225,9 +235,7 @@ export const deleteUserAIBudgetOverride = (
return {
mutationFn: () => API.deleteUserAIBudgetOverride(userId),
onSuccess: async () => {
await queryClient.invalidateQueries({
queryKey: getUserAIBudgetOverrideQueryKey(userId),
});
await invalidateUserAIBudgetQueries(queryClient, userId);
},
};
};
Expand Down
137 changes: 136 additions & 1 deletion site/src/pages/GroupsPage/GroupPage.stories.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -24,6 +24,7 @@ import type {
GroupMemberAISpend,
GroupMembersAISpend,
ReducedUser,
UserAIBudgetOverride,
UserAISpendStatus,
} from "#/api/typesGenerated";
import {
Expand Down Expand Up @@ -503,6 +504,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"],
Expand Down
Loading