diff --git a/vscode/src/chat/agentic/DeepCodyRateLimiter.test.ts b/vscode/src/chat/agentic/DeepCodyRateLimiter.test.ts index 2d386d378b72..bef272f0f6d4 100644 --- a/vscode/src/chat/agentic/DeepCodyRateLimiter.test.ts +++ b/vscode/src/chat/agentic/DeepCodyRateLimiter.test.ts @@ -87,6 +87,18 @@ describe('DeepCodyRateLimiter', () => { expect(rateLimiter.isAtLimit()).toBeUndefined() }) + + it('resets quota after 24 hours of non-use', () => { + rateLimiter = new DeepCodyRateLimiter(50, 1) + const mockUsage = { + quota: 0, // Empty quota + lastUsed: new Date(NOW.getTime() - 25 * 60 * 60 * 1000), // 25 hours ago + } + vi.spyOn(localStorage, 'getDeepCodyUsage').mockImplementation(() => mockUsage) + + expect(rateLimiter.isAtLimit()).toBeUndefined() + expect(localStorage.setDeepCodyUsage).toHaveBeenCalledWith(50, NOW.toISOString()) + }) }) describe('getRateLimitError', () => { diff --git a/vscode/src/chat/agentic/DeepCodyRateLimiter.ts b/vscode/src/chat/agentic/DeepCodyRateLimiter.ts index 3d9c504e66e0..ddf75402331f 100644 --- a/vscode/src/chat/agentic/DeepCodyRateLimiter.ts +++ b/vscode/src/chat/agentic/DeepCodyRateLimiter.ts @@ -15,38 +15,38 @@ export class DeepCodyRateLimiter { private readonly multiplier: number = 1 ) {} - /** - * Returns the number of seconds the user needs to wait before they can use DeepCody again. - */ public isAtLimit(): string | undefined { const DAILY_QUOTA = this.baseQuota * this.multiplier - // If there is no quota set, there is no limit. + // If there is no quota set, there is no limit if (!DAILY_QUOTA) { return undefined } - // Get current quota and last used time, with defaults + const now = new Date().getTime() const { quota, lastUsed } = localStorage.getDeepCodyUsage() - const currentQuota = quota ?? DAILY_QUOTA const lastUsedTime = lastUsed.getTime() - - const now = new Date().getTime() const timeDiff = now - lastUsedTime - // Calculate quota replenishment based on time passed + // Reset quota if more than 24 hours have passed + if (timeDiff >= this.ONE_DAY_MS) { + // Reset to full quota and update last used time + localStorage.setDeepCodyUsage(DAILY_QUOTA, new Date().toISOString()) + return undefined + } + + // Calculate remaining quota with time-based replenishment const quotaToAdd = DAILY_QUOTA * (timeDiff / this.ONE_DAY_MS) + const currentQuota = quota ?? DAILY_QUOTA const newQuota = Math.min(DAILY_QUOTA, currentQuota + quotaToAdd) // If we have at least 1 quota available if (newQuota >= 1) { - // Update quota and timestamp - localStorage?.setDeepCodyUsage(newQuota - 1, new Date().toISOString()) + localStorage.setDeepCodyUsage(newQuota - 1, new Date().toISOString()) return undefined } - // No quota available. - // Calculate how much time after the lastUsedTime we need to wait. + // Calculate wait time if no quota available const timeToWait = this.ONE_DAY_MS - timeDiff return Math.floor(timeToWait / 1000).toString() }