improvement(usage-indicator): update query invalidation for usage to update in realtime (#2607)

* improvement(usage-indicator): update query invalidation for usage to update in realtime

* ack PR comments
This commit is contained in:
Waleed
2025-12-27 15:23:02 -08:00
committed by GitHub
parent fdba1cfac2
commit aa9cc5604a
6 changed files with 23 additions and 15 deletions

View File

@@ -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') {

View File

@@ -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

View File

@@ -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) })

View File

@@ -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)

View File

@@ -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 })
},
})
}

View File

@@ -2661,7 +2661,7 @@ export const useCopilotStore = create<CopilotStore>()(
// Invalidate subscription queries to update usage
setTimeout(() => {
const queryClient = getQueryClient()
queryClient.invalidateQueries({ queryKey: subscriptionKeys.user() })
queryClient.invalidateQueries({ queryKey: subscriptionKeys.all })
}, 1000)
} finally {
clearTimeout(timeoutId)