diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts index 2f6d9bc47b..e1ed0500a2 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts @@ -252,7 +252,7 @@ export function useWand({ }) setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() }) + queryClient.invalidateQueries({ queryKey: subscriptionKeys.all }) }, 1000) } catch (error: any) { if (error.name === 'AbortError') { diff --git a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts index f2d9cc2f76..61af9f0624 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts +++ b/apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-workflow-execution.ts @@ -573,7 +573,7 @@ export function useWorkflowExecution() { // Invalidate subscription queries to update usage setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() }) + queryClient.invalidateQueries({ queryKey: subscriptionKeys.all }) }, 1000) safeEnqueue(encodeSSE({ event: 'final', data: result })) @@ -646,7 +646,7 @@ export function useWorkflowExecution() { // Invalidate subscription queries to update usage setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() }) + queryClient.invalidateQueries({ queryKey: subscriptionKeys.all }) }, 1000) } return result diff --git a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/components/cancel-subscription/cancel-subscription.tsx b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/components/cancel-subscription/cancel-subscription.tsx index c778f06603..2461f3e739 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/components/cancel-subscription/cancel-subscription.tsx +++ b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/components/cancel-subscription/cancel-subscription.tsx @@ -165,7 +165,7 @@ export function CancelSubscription({ subscription, subscriptionData }: CancelSub logger.info('Subscription restored successfully', result) } - await queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() }) + await queryClient.invalidateQueries({ queryKey: subscriptionKeys.all }) if (activeOrgId) { await queryClient.invalidateQueries({ queryKey: organizationKeys.detail(activeOrgId) }) await queryClient.invalidateQueries({ queryKey: organizationKeys.billing(activeOrgId) }) diff --git a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/usage-indicator/usage-indicator.tsx b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/usage-indicator/usage-indicator.tsx index dd1df673ae..97f4dc8c36 100644 --- a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/usage-indicator/usage-indicator.tsx +++ b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/usage-indicator/usage-indicator.tsx @@ -199,7 +199,7 @@ export function UsageIndicator({ onClick }: UsageIndicatorProps) { useEffect(() => { const handleOperationConfirmed = () => { setTimeout(() => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() }) + queryClient.invalidateQueries({ queryKey: subscriptionKeys.all }) }, 1000) } onOperationConfirmed(handleOperationConfirmed) diff --git a/apps/sim/hooks/queries/subscription.ts b/apps/sim/hooks/queries/subscription.ts index 750d5f5251..89ded91231 100644 --- a/apps/sim/hooks/queries/subscription.ts +++ b/apps/sim/hooks/queries/subscription.ts @@ -98,13 +98,13 @@ export function useUpdateUsageLimit() { return response.json() }, onMutate: async ({ limit }) => { - await queryClient.cancelQueries({ queryKey: subscriptionKeys.user() }) - await queryClient.cancelQueries({ queryKey: subscriptionKeys.usage() }) + await queryClient.cancelQueries({ queryKey: subscriptionKeys.all }) - const previousSubscriptionData = queryClient.getQueryData(subscriptionKeys.user()) + const previousSubscriptionData = queryClient.getQueryData(subscriptionKeys.user(false)) + const previousSubscriptionDataWithOrg = queryClient.getQueryData(subscriptionKeys.user(true)) const previousUsageData = queryClient.getQueryData(subscriptionKeys.usage()) - queryClient.setQueryData(subscriptionKeys.user(), (old: any) => { + const updateSubscriptionData = (old: any) => { if (!old) return old const currentUsage = old.data?.usage?.current || 0 const newPercentUsed = limit > 0 ? (currentUsage / limit) * 100 : 0 @@ -120,7 +120,10 @@ export function useUpdateUsageLimit() { }, }, } - }) + } + + queryClient.setQueryData(subscriptionKeys.user(false), updateSubscriptionData) + queryClient.setQueryData(subscriptionKeys.user(true), updateSubscriptionData) queryClient.setQueryData(subscriptionKeys.usage(), (old: any) => { if (!old) return old @@ -133,19 +136,24 @@ export function useUpdateUsageLimit() { } }) - return { previousSubscriptionData, previousUsageData } + return { previousSubscriptionData, previousSubscriptionDataWithOrg, previousUsageData } }, onError: (_err, _variables, context) => { if (context?.previousSubscriptionData) { - queryClient.setQueryData(subscriptionKeys.user(), context.previousSubscriptionData) + queryClient.setQueryData(subscriptionKeys.user(false), context.previousSubscriptionData) + } + if (context?.previousSubscriptionDataWithOrg) { + queryClient.setQueryData( + subscriptionKeys.user(true), + context.previousSubscriptionDataWithOrg + ) } if (context?.previousUsageData) { queryClient.setQueryData(subscriptionKeys.usage(), context.previousUsageData) } }, onSettled: () => { - queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() }) - queryClient.invalidateQueries({ queryKey: subscriptionKeys.usage() }) + queryClient.invalidateQueries({ queryKey: subscriptionKeys.all }) }, }) } diff --git a/apps/sim/stores/panel/copilot/store.ts b/apps/sim/stores/panel/copilot/store.ts index f963983ce0..8a5a634af2 100644 --- a/apps/sim/stores/panel/copilot/store.ts +++ b/apps/sim/stores/panel/copilot/store.ts @@ -2661,7 +2661,7 @@ export const useCopilotStore = create()( // Invalidate subscription queries to update usage setTimeout(() => { const queryClient = getQueryClient() - queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() }) + queryClient.invalidateQueries({ queryKey: subscriptionKeys.all }) }, 1000) } finally { clearTimeout(timeoutId)