From 7a59826d2d787fe0a400f8b1d2d5f002e3a3a201 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 27 Jul 2026 20:47:52 +0000 Subject: [PATCH 1/3] fix(site): refresh member budgets after overrides --- site/src/api/queries/groups.ts | 7 + site/src/api/queries/users.ts | 23 ++- .../pages/GroupsPage/GroupPage.stories.tsx | 150 +++++++++++++++++- 3 files changed, 173 insertions(+), 7 deletions(-) diff --git a/site/src/api/queries/groups.ts b/site/src/api/queries/groups.ts index 8093330fc24..7bd95be7dd0 100644 --- a/site/src/api/queries/groups.ts +++ b/site/src/api/queries/groups.ts @@ -64,6 +64,13 @@ export const getGroupMembersAISpendQueryKey = ( userIds: readonly string[], ) => ["group", groupId, "members", "aiSpend", [...userIds].sort()]; +export const isGroupMembersAISpendQueryKey = ( + queryKey: readonly unknown[], +): boolean => + queryKey[0] === "group" && + queryKey[2] === "members" && + queryKey[3] === "aiSpend"; + export const groupMembersAISpend = ( groupId: string, userIds: readonly string[], diff --git a/site/src/api/queries/users.ts b/site/src/api/queries/users.ts index b5a9b2bed60..2856f03dc3f 100644 --- a/site/src/api/queries/users.ts +++ b/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 { isGroupMembersAISpendQueryKey } from "./groups"; import { cachedQuery } from "./util"; export function usersKey(req: UsersRequest) { @@ -211,9 +212,14 @@ export const saveUserAIBudgetOverride = ( mutationFn: (request: UpsertUserAIBudgetOverrideRequest) => API.upsertUserAIBudgetOverride(userId, request), onSuccess: async () => { - await queryClient.invalidateQueries({ - queryKey: getUserAIBudgetOverrideQueryKey(userId), - }); + await Promise.all([ + queryClient.invalidateQueries({ + queryKey: getUserAIBudgetOverrideQueryKey(userId), + }), + queryClient.invalidateQueries({ + predicate: (query) => isGroupMembersAISpendQueryKey(query.queryKey), + }), + ]); }, }; }; @@ -225,9 +231,14 @@ export const deleteUserAIBudgetOverride = ( return { mutationFn: () => API.deleteUserAIBudgetOverride(userId), onSuccess: async () => { - await queryClient.invalidateQueries({ - queryKey: getUserAIBudgetOverrideQueryKey(userId), - }); + await Promise.all([ + queryClient.invalidateQueries({ + queryKey: getUserAIBudgetOverrideQueryKey(userId), + }), + queryClient.invalidateQueries({ + predicate: (query) => isGroupMembersAISpendQueryKey(query.queryKey), + }), + ]); }, }; }; diff --git a/site/src/pages/GroupsPage/GroupPage.stories.tsx b/site/src/pages/GroupsPage/GroupPage.stories.tsx index 8b8c15ba3be..f86d9396663 100644 --- a/site/src/pages/GroupsPage/GroupPage.stories.tsx +++ b/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 { @@ -503,6 +504,153 @@ 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 updatedMembersSpend = ( + groupBudget: GroupMemberAISpend["group_budget"], +): GroupMembersAISpend => ({ + period_start: "2026-06-01T00:00:00Z", + period_end: "2026-07-01T00:00:00Z", + members: [ + { + ...mockSpend, + user_id: MockUserOwner.id, + group_budget: groupBudget, + }, + ], +}); + +export const SaveMemberAIBudgetRefreshesRow: Story = { + beforeEach: () => { + spyOn(API, "upsertUserAIBudgetOverride").mockResolvedValue( + mockUserBudgetOverride, + ); + spyOn(API, "getUserAIBudgetOverride").mockResolvedValue( + mockUserBudgetOverride, + ); + spyOn(API, "getGroupMembersAISpend").mockResolvedValue( + updatedMembersSpend({ + spend_limit_micros: mockUserBudgetOverride.spend_limit_micros, + limit_source: "user_override", + }), + ); + }, + parameters: { + features: ["aibridge"], + experiments: ["ai-gateway-cost-control"], + queries: [ + groupQuery(MockGroupWithoutMembers), + groupMembersQuery({ users: [MockUserOwner], count: 1 }), + membersSpendQuery([{ ...mockSpend, user_id: MockUserOwner.id }]), + 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( + updatedMembersSpend(mockSpend.group_budget), + ); + }, + parameters: { + features: ["aibridge"], + experiments: ["ai-gateway-cost-control"], + queries: [ + groupQuery(MockGroupWithoutMembers), + groupMembersQuery({ users: [MockUserOwner], count: 1 }), + membersSpendQuery([ + { + ...mockSpend, + user_id: MockUserOwner.id, + group_budget: { + spend_limit_micros: mockUserBudgetOverride.spend_limit_micros, + limit_source: "user_override", + }, + }, + ]), + 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"], From 52339380904751eb54a7c1cc49506710643eaf61 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 28 Jul 2026 13:27:37 +0000 Subject: [PATCH 2/3] fix(site): target member budget cache invalidation --- site/src/api/queries/groups.test.ts | 27 ++++++++++++ site/src/api/queries/groups.ts | 5 ++- site/src/api/queries/users.ts | 6 ++- .../pages/GroupsPage/GroupPage.stories.tsx | 44 ++++++++++--------- 4 files changed, 58 insertions(+), 24 deletions(-) create mode 100644 site/src/api/queries/groups.test.ts diff --git a/site/src/api/queries/groups.test.ts b/site/src/api/queries/groups.test.ts new file mode 100644 index 00000000000..3bddb54af7c --- /dev/null +++ b/site/src/api/queries/groups.test.ts @@ -0,0 +1,27 @@ +import { describe, expect, it } from "vitest"; +import { + getGroupMembersAISpendQueryKey, + isGroupMembersAISpendQueryKey, +} from "./groups"; + +describe("isGroupMembersAISpendQueryKey", () => { + it("matches only group member spend queries containing the user", () => { + const userId = "user-1"; + + expect( + isGroupMembersAISpendQueryKey( + getGroupMembersAISpendQueryKey("group-1", ["user-2", userId]), + userId, + ), + ).toBe(true); + expect( + isGroupMembersAISpendQueryKey( + getGroupMembersAISpendQueryKey("group-1", ["user-2"]), + userId, + ), + ).toBe(false); + expect(isGroupMembersAISpendQueryKey(["group", "group-1"], userId)).toBe( + false, + ); + }); +}); diff --git a/site/src/api/queries/groups.ts b/site/src/api/queries/groups.ts index 7bd95be7dd0..4f1b9dbb74e 100644 --- a/site/src/api/queries/groups.ts +++ b/site/src/api/queries/groups.ts @@ -66,10 +66,13 @@ export const getGroupMembersAISpendQueryKey = ( export const isGroupMembersAISpendQueryKey = ( queryKey: readonly unknown[], + userId: string, ): boolean => queryKey[0] === "group" && queryKey[2] === "members" && - queryKey[3] === "aiSpend"; + queryKey[3] === "aiSpend" && + Array.isArray(queryKey[4]) && + queryKey[4].includes(userId); export const groupMembersAISpend = ( groupId: string, diff --git a/site/src/api/queries/users.ts b/site/src/api/queries/users.ts index 2856f03dc3f..9369a4b8a53 100644 --- a/site/src/api/queries/users.ts +++ b/site/src/api/queries/users.ts @@ -217,7 +217,8 @@ export const saveUserAIBudgetOverride = ( queryKey: getUserAIBudgetOverrideQueryKey(userId), }), queryClient.invalidateQueries({ - predicate: (query) => isGroupMembersAISpendQueryKey(query.queryKey), + predicate: (query) => + isGroupMembersAISpendQueryKey(query.queryKey, userId), }), ]); }, @@ -236,7 +237,8 @@ export const deleteUserAIBudgetOverride = ( queryKey: getUserAIBudgetOverrideQueryKey(userId), }), queryClient.invalidateQueries({ - predicate: (query) => isGroupMembersAISpendQueryKey(query.queryKey), + predicate: (query) => + isGroupMembersAISpendQueryKey(query.queryKey, userId), }), ]); }, diff --git a/site/src/pages/GroupsPage/GroupPage.stories.tsx b/site/src/pages/GroupsPage/GroupPage.stories.tsx index f86d9396663..92aa55e0811 100644 --- a/site/src/pages/GroupsPage/GroupPage.stories.tsx +++ b/site/src/pages/GroupsPage/GroupPage.stories.tsx @@ -512,19 +512,16 @@ const mockUserBudgetOverride: UserAIBudgetOverride = { updated_at: "2026-06-01T00:00:00Z", }; -const updatedMembersSpend = ( - groupBudget: GroupMemberAISpend["group_budget"], -): GroupMembersAISpend => ({ +const mockOwnerSpend: GroupMemberAISpend = { + ...mockSpend, + user_id: MockUserOwner.id, +}; + +const mockOwnerMembersSpend: GroupMembersAISpend = { period_start: "2026-06-01T00:00:00Z", period_end: "2026-07-01T00:00:00Z", - members: [ - { - ...mockSpend, - user_id: MockUserOwner.id, - group_budget: groupBudget, - }, - ], -}); + members: [mockOwnerSpend], +}; export const SaveMemberAIBudgetRefreshesRow: Story = { beforeEach: () => { @@ -534,12 +531,18 @@ export const SaveMemberAIBudgetRefreshesRow: Story = { spyOn(API, "getUserAIBudgetOverride").mockResolvedValue( mockUserBudgetOverride, ); - spyOn(API, "getGroupMembersAISpend").mockResolvedValue( - updatedMembersSpend({ - spend_limit_micros: mockUserBudgetOverride.spend_limit_micros, - limit_source: "user_override", - }), - ); + spyOn(API, "getGroupMembersAISpend").mockResolvedValue({ + ...mockOwnerMembersSpend, + members: [ + { + ...mockOwnerSpend, + group_budget: { + spend_limit_micros: mockUserBudgetOverride.spend_limit_micros, + limit_source: "user_override", + }, + }, + ], + }); }, parameters: { features: ["aibridge"], @@ -547,7 +550,7 @@ export const SaveMemberAIBudgetRefreshesRow: Story = { queries: [ groupQuery(MockGroupWithoutMembers), groupMembersQuery({ users: [MockUserOwner], count: 1 }), - membersSpendQuery([{ ...mockSpend, user_id: MockUserOwner.id }]), + membersSpendQuery([mockOwnerSpend]), permissionsQuery({ canUpdateGroup: true }), { key: meAISpendKey, data: mockUserAISpend }, { key: getUserAIBudgetOverrideQueryKey(MockUserOwner.id), data: null }, @@ -592,7 +595,7 @@ export const DeleteMemberAIBudgetRefreshesRow: Story = { beforeEach: () => { spyOn(API, "deleteUserAIBudgetOverride").mockResolvedValue(); spyOn(API, "getGroupMembersAISpend").mockResolvedValue( - updatedMembersSpend(mockSpend.group_budget), + mockOwnerMembersSpend, ); }, parameters: { @@ -603,8 +606,7 @@ export const DeleteMemberAIBudgetRefreshesRow: Story = { groupMembersQuery({ users: [MockUserOwner], count: 1 }), membersSpendQuery([ { - ...mockSpend, - user_id: MockUserOwner.id, + ...mockOwnerSpend, group_budget: { spend_limit_micros: mockUserBudgetOverride.spend_limit_micros, limit_source: "user_override", From 018f91bb113ea47e693dfc86c2618628257a7047 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Wed, 29 Jul 2026 15:57:30 +0300 Subject: [PATCH 3/3] refactor(site): extract AI budget invalidation helpers --- site/src/api/queries/groups.test.ts | 40 +++++++++++-------- site/src/api/queries/groups.ts | 11 ++++- site/src/api/queries/users.ts | 33 +++++++-------- .../pages/GroupsPage/GroupPage.stories.tsx | 37 +++++------------ 4 files changed, 59 insertions(+), 62 deletions(-) diff --git a/site/src/api/queries/groups.test.ts b/site/src/api/queries/groups.test.ts index 3bddb54af7c..74ebf8606b6 100644 --- a/site/src/api/queries/groups.test.ts +++ b/site/src/api/queries/groups.test.ts @@ -1,26 +1,34 @@ +import { QueryClient } from "react-query"; import { describe, expect, it } from "vitest"; import { getGroupMembersAISpendQueryKey, - isGroupMembersAISpendQueryKey, + invalidateGroupMembersAISpend, } from "./groups"; -describe("isGroupMembersAISpendQueryKey", () => { - it("matches only group member spend queries containing the user", () => { +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"]; - expect( - isGroupMembersAISpendQueryKey( - getGroupMembersAISpendQueryKey("group-1", ["user-2", userId]), - userId, - ), - ).toBe(true); - expect( - isGroupMembersAISpendQueryKey( - getGroupMembersAISpendQueryKey("group-1", ["user-2"]), - userId, - ), - ).toBe(false); - expect(isGroupMembersAISpendQueryKey(["group", "group-1"], userId)).toBe( + 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, ); }); diff --git a/site/src/api/queries/groups.ts b/site/src/api/queries/groups.ts index 4f1b9dbb74e..6e714f92c79 100644 --- a/site/src/api/queries/groups.ts +++ b/site/src/api/queries/groups.ts @@ -64,7 +64,7 @@ export const getGroupMembersAISpendQueryKey = ( userIds: readonly string[], ) => ["group", groupId, "members", "aiSpend", [...userIds].sort()]; -export const isGroupMembersAISpendQueryKey = ( +const isGroupMembersAISpendQueryKey = ( queryKey: readonly unknown[], userId: string, ): boolean => @@ -74,6 +74,15 @@ export const isGroupMembersAISpendQueryKey = ( 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[], diff --git a/site/src/api/queries/users.ts b/site/src/api/queries/users.ts index 9369a4b8a53..47e7f539319 100644 --- a/site/src/api/queries/users.ts +++ b/site/src/api/queries/users.ts @@ -31,7 +31,7 @@ import { import type { UsePaginatedQueryOptions } from "#/hooks/usePaginatedQuery"; import { prepareQuery } from "#/utils/filters"; import { getAuthorizationKey } from "./authCheck"; -import { isGroupMembersAISpendQueryKey } from "./groups"; +import { invalidateGroupMembersAISpend } from "./groups"; import { cachedQuery } from "./util"; export function usersKey(req: UsersRequest) { @@ -204,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, @@ -212,15 +223,7 @@ export const saveUserAIBudgetOverride = ( mutationFn: (request: UpsertUserAIBudgetOverrideRequest) => API.upsertUserAIBudgetOverride(userId, request), onSuccess: async () => { - await Promise.all([ - queryClient.invalidateQueries({ - queryKey: getUserAIBudgetOverrideQueryKey(userId), - }), - queryClient.invalidateQueries({ - predicate: (query) => - isGroupMembersAISpendQueryKey(query.queryKey, userId), - }), - ]); + await invalidateUserAIBudgetQueries(queryClient, userId); }, }; }; @@ -232,15 +235,7 @@ export const deleteUserAIBudgetOverride = ( return { mutationFn: () => API.deleteUserAIBudgetOverride(userId), onSuccess: async () => { - await Promise.all([ - queryClient.invalidateQueries({ - queryKey: getUserAIBudgetOverrideQueryKey(userId), - }), - queryClient.invalidateQueries({ - predicate: (query) => - isGroupMembersAISpendQueryKey(query.queryKey, userId), - }), - ]); + await invalidateUserAIBudgetQueries(queryClient, userId); }, }; }; diff --git a/site/src/pages/GroupsPage/GroupPage.stories.tsx b/site/src/pages/GroupsPage/GroupPage.stories.tsx index 92aa55e0811..4b3ad04fe84 100644 --- a/site/src/pages/GroupsPage/GroupPage.stories.tsx +++ b/site/src/pages/GroupsPage/GroupPage.stories.tsx @@ -517,10 +517,12 @@ const mockOwnerSpend: GroupMemberAISpend = { user_id: MockUserOwner.id, }; -const mockOwnerMembersSpend: GroupMembersAISpend = { - period_start: "2026-06-01T00:00:00Z", - period_end: "2026-07-01T00:00:00Z", - members: [mockOwnerSpend], +const mockOwnerOverrideSpend: GroupMemberAISpend = { + ...mockOwnerSpend, + group_budget: { + spend_limit_micros: mockUserBudgetOverride.spend_limit_micros, + limit_source: "user_override", + }, }; export const SaveMemberAIBudgetRefreshesRow: Story = { @@ -531,18 +533,9 @@ export const SaveMemberAIBudgetRefreshesRow: Story = { spyOn(API, "getUserAIBudgetOverride").mockResolvedValue( mockUserBudgetOverride, ); - spyOn(API, "getGroupMembersAISpend").mockResolvedValue({ - ...mockOwnerMembersSpend, - members: [ - { - ...mockOwnerSpend, - group_budget: { - spend_limit_micros: mockUserBudgetOverride.spend_limit_micros, - limit_source: "user_override", - }, - }, - ], - }); + spyOn(API, "getGroupMembersAISpend").mockResolvedValue( + membersSpendQuery([mockOwnerOverrideSpend]).data, + ); }, parameters: { features: ["aibridge"], @@ -595,7 +588,7 @@ export const DeleteMemberAIBudgetRefreshesRow: Story = { beforeEach: () => { spyOn(API, "deleteUserAIBudgetOverride").mockResolvedValue(); spyOn(API, "getGroupMembersAISpend").mockResolvedValue( - mockOwnerMembersSpend, + membersSpendQuery([mockOwnerSpend]).data, ); }, parameters: { @@ -604,15 +597,7 @@ export const DeleteMemberAIBudgetRefreshesRow: Story = { queries: [ groupQuery(MockGroupWithoutMembers), groupMembersQuery({ users: [MockUserOwner], count: 1 }), - membersSpendQuery([ - { - ...mockOwnerSpend, - group_budget: { - spend_limit_micros: mockUserBudgetOverride.spend_limit_micros, - limit_source: "user_override", - }, - }, - ]), + membersSpendQuery([mockOwnerOverrideSpend]), permissionsQuery({ canUpdateGroup: true }), { key: meAISpendKey, data: mockUserAISpend }, {