Skip to content

Commit

Permalink
[Backport vscode-v1.52.x] fix(deep-cody): Rate limiter synchronizatio…
Browse files Browse the repository at this point in the history 
…n issue (#6378)

Follow up on #6366

This PR fixes several issues with the Deep Cody rate limiter
implementation:

Root Cause Fix: Rate Limit State Synchronization

- Fixed race condition in LocalStorageProvider.ts where concurrent
Promise.all() usage led to inconsistent state between quota and
timestamp storage
- Changed to sequential storage operations to ensure timestamp is always
saved after quota update
- This resolves the core issue where users were locked out due to
incorrect time calculations (CODY-4529)

## Test plan



- Added precise tests for time-based quota calculations
- Enhanced test coverage for various timing scenarios

Before this change, the Deep Cody rate limit error would always return
the "Retry time" with current time + 24 hours to indicate the
current time is being used as the last used time.

Submitted at 6:44


![image](https://github.com/user-attachments/assets/d9b50ff8-bc09-473b-a387-497d1ccbe0c2)

Submitted again at 6:45


![image](https://github.com/user-attachments/assets/e1b0266b-cfbc-40f8-807b-a9d4c93be682)

After this change, the "Retry Time" should stick with the last
used time until it's been reset after 24 hours.

## Changelog


 <br> Backport a7339ef from #6377

Co-authored-by: Beatrix <[email protected]>
  • Loading branch information
@sourcegraph-release-bot @abeatrix
sourcegraph-release-bot and abeatrix authored Dec 17, 2024
1 parent da995cb commit fb81339
Show file tree
Hide file tree
Showing 4 changed files with 39 additions and 110 deletions.
83 changes: 26 additions & 57 deletions vscode/src/chat/agentic/DeepCodyRateLimiter.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,6 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { localStorage } from './../../services/LocalStorageProvider'
import { DeepCodyRateLimiter } from './DeepCodyRateLimiter'

// Create a mock type for localStorage
vi.mock('./../../services/LocalStorageProvider', () => ({
localStorage: {
getDeepCodyUsage: vi.fn(),
setDeepCodyUsage: vi.fn(),
setStorage: vi.fn(),
},
}))

describe('DeepCodyRateLimiter', () => {
let rateLimiter: DeepCodyRateLimiter
const NOW = new Date('2024-01-01T12:00:00Z')
Expand All @@ -36,68 +27,46 @@ describe('DeepCodyRateLimiter', () => {
expect(rateLimiter.isAtLimit()).toBeUndefined()
})

it('allows usage when quota available', () => {
rateLimiter = new DeepCodyRateLimiter(10, 1)

// Set up mock return value
const mockUsage = {
quota: 5,
lastUsed: new Date(NOW.getTime() - 3600000), // 1 hour ago
}
vi.spyOn(localStorage, 'getDeepCodyUsage').mockImplementation(() => mockUsage)

it('allows usage when quota available', async () => {
rateLimiter = new DeepCodyRateLimiter(6, 1)
expect(rateLimiter.isAtLimit()).toBeUndefined()
expect(localStorage.setDeepCodyUsage).toHaveBeenCalled()
})

it('blocks usage when no quota available', () => {
rateLimiter = new DeepCodyRateLimiter(10, 1)

const mockUsage = {
quota: 0,
lastUsed: new Date(NOW.getTime() - 3600000),
}
vi.spyOn(localStorage, 'getDeepCodyUsage').mockImplementation(() => mockUsage)

const result = rateLimiter.isAtLimit()
expect(result).toBeDefined()
expect(Number(result)).toBeGreaterThan(0)
})

it('correctly calculates quota replenishment', () => {
rateLimiter = new DeepCodyRateLimiter(24, 1) // 24 tokens per day = 1 per hour

const mockUsage = {
quota: 0,
lastUsed: new Date(NOW.getTime() - 3600000),
}
vi.spyOn(localStorage, 'getDeepCodyUsage').mockImplementation(() => mockUsage)

const { quota, lastUsed } = localStorage.getDeepCodyUsage()
expect(Math.ceil(quota!)).toBe(5)
expect(lastUsed).toBe(NOW.toISOString())
expect(rateLimiter.isAtLimit()).toBeUndefined()
})

it('respects multiplier in quota calculation', () => {
it('respects multiplier in quota calculation', async () => {
const { quota, lastUsed } = localStorage.getDeepCodyUsage()
expect(Math.ceil(quota!)).toBe(4)
expect(lastUsed).toBe(NOW.toISOString())
rateLimiter = new DeepCodyRateLimiter(10, 2) // 20 tokens per day

const mockUsage = {
quota: 0,
lastUsed: new Date(NOW.getTime() - 43200000), // 12 hours ago
}
vi.spyOn(localStorage, 'getDeepCodyUsage').mockImplementation(() => mockUsage)

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)
it('resets quota after 24 hours of non-use', async () => {
const { quota, lastUsed } = localStorage.getDeepCodyUsage()
expect(Math.ceil(quota!)).toBe(3)
expect(lastUsed).toBe(NOW.toISOString())
rateLimiter = new DeepCodyRateLimiter(3, 1)
expect(rateLimiter.isAtLimit()).toBeUndefined()
})

it('blocks usage when no quota available', async () => {
rateLimiter = new DeepCodyRateLimiter(1, 1)
const ONE_DAY_MS = 24 * 60 * 60 * 1000
const ONE_HOUR_MS = ONE_DAY_MS / 24
expect(rateLimiter.isAtLimit()).toBeUndefined()
expect(localStorage.setDeepCodyUsage).toHaveBeenCalledWith(50, NOW.toISOString())
// It should be 24 hours after last used time (which is current)
expect(Number(rateLimiter.isAtLimit())).toBe((ONE_HOUR_MS * 24) / 1000)
// Fake an hour has passed.
vi.advanceTimersByTime(ONE_HOUR_MS)
// Check if the time to wait has decreased by an hour.
expect(Number(rateLimiter.isAtLimit())).toBe((ONE_HOUR_MS * 23) / 1000)
})
})

Expand Down
16 changes: 8 additions & 8 deletions vscode/src/chat/agentic/DeepCodyRateLimiter.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,26 +23,26 @@ export class DeepCodyRateLimiter {
return undefined
}

const now = new Date().getTime()
const now = new Date()
const { quota, lastUsed } = localStorage.getDeepCodyUsage()
const lastUsedTime = lastUsed.getTime()
const timeDiff = now - lastUsedTime

// 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())
// Reset for cases where lastUsed was not stored properly but quota was.
if (quota !== undefined && lastUsed === undefined) {
localStorage.setDeepCodyUsage(DAILY_QUOTA - 1, now.toISOString())
return undefined
}

const lastUsedTime = new Date(lastUsed ?? now.toISOString()).getTime()
const timeDiff = now.getTime() - lastUsedTime

// 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) {
localStorage.setDeepCodyUsage(newQuota - 1, new Date().toISOString())
localStorage.setDeepCodyUsage(newQuota - 1, now.toISOString())
return undefined
}

Expand Down
2 changes: 1 addition & 1 deletion vscode/src/chat/chat-view/ChatController.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -811,7 +811,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
this.featureDeepCodyRateLimitMultiplier.value.last ? 2 : 1
)
const deepCodyLimit = deepCodyRateLimiter.isAtLimit()
if (isDeepCodyModel && isDeepCodyEnabled && deepCodyLimit) {
if (isDeepCodyEnabled && deepCodyLimit) {
this.postError(deepCodyRateLimiter.getRateLimitError(deepCodyLimit), 'transcript')
this.handleAbort()
return
Expand Down
48 changes: 4 additions & 44 deletions vscode/src/services/LocalStorageProvider.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -347,56 +347,16 @@ class LocalStorage implements LocalStorageForModelPreferences {
await this.set(this.CODY_CHAT_MEMORY, memories)
}

public getDeepCodyUsage(): { quota: number | undefined; lastUsed: Date } {
public getDeepCodyUsage(): { quota: number | undefined; lastUsed: string | undefined } {
const quota = this.get<number>(this.keys.deepCodyDailyUsageCount) ?? undefined
const lastUsed = new Date(
this.get<string>(this.keys.deepCodyLastUsedDate) ?? new Date().toISOString()
)
const lastUsed = this.get<string>(this.keys.deepCodyLastUsedDate) ?? undefined

return { quota, lastUsed }
}

public async setDeepCodyUsage(newQuota: number, lastUsed: string): Promise<void> {
await Promise.all([
localStorage.set(localStorage.keys.deepCodyDailyUsageCount, newQuota - 1),
localStorage.set(localStorage.keys.deepCodyLastUsedDate, lastUsed),
])
}

public isAtDeepCodyDailyLimit(DAILY_QUOTA?: number): string | undefined {
if (!DAILY_QUOTA) {
return undefined
}

const ONE_DAY_MS = 24 * 60 * 60 * 1000

// Get current quota and last used time, with defaults
const currentQuota = this.get<number>(this.keys.deepCodyDailyUsageCount) ?? DAILY_QUOTA
const lastUsedTime = new Date(
this.get<string>(this.keys.deepCodyLastUsedDate) ?? new Date().toISOString()
).getTime()

const now = new Date().getTime()
const timeDiff = now - lastUsedTime

// Calculate quota replenishment based on time passed
const quotaToAdd = DAILY_QUOTA * (timeDiff / ONE_DAY_MS)
const newQuota = Math.min(DAILY_QUOTA, currentQuota + quotaToAdd)

// If we have at least 1 quota available
if (newQuota >= 1) {
// Update quota and timestamp
Promise.all([
this.set(this.keys.deepCodyDailyUsageCount, newQuota - 1),
this.set(this.keys.deepCodyLastUsedDate, new Date().toISOString()),
])
return undefined
}

// No quota available.
// Calculate how much time after the lastUsedTime we need to wait.
const timeToWait = ONE_DAY_MS - timeDiff
return Math.floor(timeToWait / 1000).toString()
await this.set(this.keys.deepCodyDailyUsageCount, newQuota)
await this.set(this.keys.deepCodyLastUsedDate, lastUsed)
}

public get<T>(key: string): T | null {
Expand Down

0 comments on commit fb81339

Please sign in to comment.