diff --git a/site/src/api/queries/groups.test.ts b/site/src/api/queries/groups.test.ts new file mode 100644 index 0000000000000..74ebf8606b651 --- /dev/null +++ b/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, + ); + }); +}); diff --git a/site/src/api/queries/groups.ts b/site/src/api/queries/groups.ts index 8093330fc244a..6e714f92c7925 100644 --- a/site/src/api/queries/groups.ts +++ b/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[], diff --git a/site/src/api/queries/users.ts b/site/src/api/queries/users.ts index b5a9b2bed6039..47e7f539319a4 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 { 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); }, }; }; diff --git a/site/src/pages/GroupsPage/GroupPage.stories.tsx b/site/src/pages/GroupsPage/GroupPage.stories.tsx index 8b8c15ba3be7a..4b3ad04fe8437 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,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"],