ApiKeyPoolService Algorithm Deep Dive

ApiKeyPoolService implements a load-balancing solution for multiple API keys: weighted round-robin key selection, session-level affinity to maintain Prompt Cache, exponential backoff cooldowns for rate limiting, and permanent disabling on authentication failures.

---

File Locations

File	Path
ApiKeyPoolService	packages/desktop/app/main/services/capabilities/llm/completion/ApiKeyPoolService.ts
API Key DB operations	packages/desktop/app/main/workers/db/apiKeys.ts
DB Worker registration	packages/desktop/app/main/workers/db/index.ts
DB Worker types	packages/desktop/app/main/workers/types.ts
IPC Router	packages/desktop/app/main/services/routers/llm/ApiKeyRouter.ts
Frontend UI	packages/renderer/src/features/settings/components/provider-settings/llm/ApiKeyPoolSection.tsx

---

Architectural Context

graph TB
    subgraph CompletionService
        Resolve[resolveApiKeyForRequest]
        Retry[retry on 429/529]
        Success[reportSuccess]
    end

    subgraph ApiKeyPoolService
        direction TB
        GetKey[getKeyForSession]
        GetKeyNoSession[getKey]
        Report[reportError]
        ReportOk[reportSuccess]
        Select[selectWeightedRoundRobin]
        Available[getAvailableKeys]
        Cooldown[applyCooldown]
        AuthFail[handleAuthFailure]
        Cleanup[cleanupExpiredCooldowns<br/>every 30 seconds]
    end

    subgraph InMemoryStructures ["In-Memory Data Structures"]
        SB["sessionBindings<br/>Map&lt;sessionId, SessionBinding&gt;"]
        RRI["rrIndex<br/>Map&lt;providerId, number&gt;"]
        KC["keyCache<br/>Map&lt;providerId, ApiKeyEntry[]&gt;"]
        CD["cooldowns<br/>Map&lt;keyId, KeyCooldown&gt;"]
    end

    subgraph ExternalDeps ["External Dependencies"]
        DB[(SQLite<br/>llm_provider_api_keys)]
        Loader["loadKeys(providerId)"]
        Disabler["disableKey(keyId)"]
    end

    Resolve --> GetKey
    Retry --> Report
    Success --> ReportOk

    GetKey --> SB
    GetKey --> Available
    Available --> KC
    Available --> CD
    KC --> Loader
    Loader --> DB

    Select --> RRI
    Report --> Cooldown
    Report --> AuthFail
    AuthFail --> Disabler
    Disabler --> DB
    Cleanup --> CD

---

Data Structures

Core Types

// Session binding: locks a session to a specific key
interface SessionBinding {
  keyId: string;        // Bound key ID
  providerId: string;   // Owning provider ID
}

// Key cooldown state
interface KeyCooldown {
  until: number;        // Cooldown expiry timestamp (Date.now() + cooldownMs)
  errors: number;       // Consecutive error count (used for exponential backoff)
}

// API key entry (from database)
interface ApiKeyEntry {
  id: string;           // UUID
  providerId: string;   // Owning provider
  label?: string;       // Display label (e.g. "Production Key #1")
  apiKey: string;       // Actual key value (may start with $ for env var)
  enabled: boolean;     // Whether enabled
  weight: number;       // Weight (1-100)
}

// Key loader function signature
type ApiKeysLoader = (providerId: string) => Promise<ApiKeyEntry[]>;

// Key disabler function signature
type ApiKeyDisabler = (keyId: string) => Promise<boolean>;

// Key resolver (handles $ env variable prefix)
type ApiKeyResolver = (rawKey: string) => string;

In-Memory Data Structures at a Glance

Structure	Type	Purpose	Lifecycle
sessionBindings	Map<sessionId, SessionBinding>	Maps sessions to bound keys	Cleared via releaseSession() when session ends
rrIndex	Map<providerId, number>	Per-provider round-robin index	Persists for the application lifetime
keyCache	Map<providerId, ApiKeyEntry[]>	Key list cache (avoids DB queries on every request)	Cleared via invalidateCache() after CRUD operations
cooldowns	Map<keyId, KeyCooldown>	Key cooldown state	Expired entries cleaned up every 30 seconds

---

Algorithms and Logic

Weighted Round-Robin Algorithm

Each key's weight determines the number of "slots" it occupies in a round-robin cycle.

Steps:

selectWeightedRoundRobin(providerId, keys):
  1. If only 1 key → return directly
  2. Calculate totalWeight = sum(keys[i].weight)
  3. Advance round-robin index: idx = (rrIndex[providerId] + 1) % totalWeight
  4. Save new index: rrIndex[providerId] = idx
  5. Accumulate and iterate:
     accum = 0
     for each key in keys:
       accum += key.weight
       if idx < accum:
         return key
  6. Fallback: return keys[0]

Example:

Three keys: A(weight=3), B(weight=1), C(weight=2), totalWeight=6

Round-robin index (idx)	Accumulated	Selected key
0	A: 3	A (0 < 3)
1	A: 3	A (1 < 3)
2	A: 3	A (2 < 3)
3	A: 3, B: 4	B (3 < 4)
4	A: 3, B: 4, C: 6	C (4 < 6)
5	A: 3, B: 4, C: 6	C (5 < 6)
0	(cycle repeats)	A

Weight meaning: a key with weight=3 is selected 3 times per cycle, a key with weight=1 is selected once.

---

Session Affinity (Session Binding)

flowchart TD
    Start[getKeyForSession] --> CheckBinding{Session has binding?}
    CheckBinding -->|Yes| CheckProvider{providerId matches?}
    CheckProvider -->|Yes| CheckAvailable{Bound key available?}
    CheckAvailable -->|Yes| Return[Return bound key]
    CheckAvailable -->|No| ReBind[Re-bind]
    CheckProvider -->|No| ReBind
    CheckBinding -->|No| ReBind

    ReBind --> GetKeys[getAvailableKeys]
    GetKeys --> Empty{Key list empty?}
    Empty -->|Yes| ReturnEmpty[Return empty string]
    Empty -->|No| WRR[selectWeightedRoundRobin]
    WRR --> Bind[sessionBindings.set]
    Bind --> ReturnNew[Return new key]

Why session affinity is needed:

• Providers such as Anthropic implement Prompt Cache
• Sending requests with the same API Key can hit the cache, saving cost and time
• Switching keys causes cache misses
• Therefore, within a single session, the same key should be used whenever possible

Handling an unavailable key:

getKeyForSession(providerId, sessionId):
  binding = sessionBindings.get(sessionId)
  if binding && binding.providerId === providerId:
    keys = getAvailableKeys(providerId)  // filter: enabled + not in cooldown
    boundKey = keys.find(k.id === binding.keyId)
    if boundKey:
      return resolveKey(boundKey.apiKey)  // hit: return
    // Key disabled/deleted/in cooldown → re-bind
    log.info("Session key no longer available, re-binding")

  // Select a new key and bind
  keys = getAvailableKeys(providerId)
  if keys.length === 0: return ''
  selected = selectWeightedRoundRobin(providerId, keys)
  sessionBindings.set(sessionId, { keyId: selected.id, providerId })
  return resolveKey(selected.apiKey)

---

Cooldown / Backoff Mechanism

Exponential Backoff Formula

cooldownMs = min(DEFAULT_COOLDOWN_MS * 2^(errors - 1), MAX_COOLDOWN_MS)

Consecutive errors	Calculation	Cooldown duration
1	60,000 * 2^0	60 s (1 minute)
2	60,000 * 2^1	120 s (2 minutes)
3	60,000 * 2^2	240 s (4 minutes)
4	60,000 * 2^3	480 s (8 minutes)
5+	60,000 * 2^4	900 s (15 minutes cap)

Constant configuration:

Constant	Value	Description
DEFAULT_COOLDOWN_MS	60,000 (60 s)	Base cooldown duration
MAX_COOLDOWN_MS	900,000 (15 min)	Maximum cooldown duration
COOLDOWN_MULTIPLIER	2	Exponential base
Cleanup interval	30,000 (30 s)	cleanupExpiredCooldowns() interval

Cooldown Application Flow

applyCooldown(keyId, providerId, statusCode):
  current = cooldowns.get(keyId)
  errors = (current?.errors ?? 0) + 1
  cooldownMs = min(60_000 * 2^(errors-1), 900_000)
  cooldowns.set(keyId, {
    until: Date.now() + cooldownMs,
    errors: errors
  })

Cooldown Reset

• Call reportSuccess(sessionId) on a successful request
• If the bound key has a cooldown entry, delete it immediately

---

Authentication Failure Handling

For HTTP 401 and 403 errors, the key is treated as permanently invalid:

flowchart TD
    Error[reportError] --> Check{HTTP status code}
    Check -->|429/529| Cooldown[applyCooldown<br/>exponential backoff]
    Check -->|401/403| AuthFail[handleAuthFailure]
    Check -->|other| Ignore[Ignore]

    AuthFail --> Disable[disableKey<br/>disable in database]
    Disable --> InvalidateCache[keyCache.delete<br/>clear cache]

    Cooldown --> Rebind[Re-select key]
    InvalidateCache --> Rebind

    Rebind --> HasMore{More keys available?}
    HasMore -->|Yes| Return[Return new key]
    HasMore -->|No| ReturnNull[Return null]

HTTP status code classification:

Status code	Class	Handling
401	AUTH_FAILURE	Permanently disable key
403	AUTH_FAILURE	Permanently disable key
429	RATE_LIMIT	Exponential backoff cooldown
529	RATE_LIMIT	Exponential backoff cooldown (Anthropic overload)
Other	—	No action, return null

---

Key Availability Filtering

getAvailableKeys(providerId):
  all = getAllKeys(providerId)  // load from cache or database
  now = Date.now()
  return all.filter(key =>
    key.enabled === true               // must be enabled
    && !(cooldowns[key.id]?.until > now)  // not in cooldown
  )

---

Cleanup Interval

cleanupExpiredCooldowns() runs every 30 seconds:

cleanupExpiredCooldowns():
  now = Date.now()
  for each [keyId, cd] in cooldowns:
    if cd.until <= now:
      cooldowns.delete(keyId)

---

IPC Integration Table

IPC Channel	Direction	Parameters	Zod Schema	Description
llmConfig:getApiKeys	R → M	providerId: string	None	Get all keys for a provider
llmConfig:addApiKey	R → M	{ providerId, label?, apiKey, enabled?, weight? }	AddApiKeySchema	Add a key (default weight=1, enabled=true)
llmConfig:updateApiKey	R → M	{ id, label?, apiKey?, enabled?, weight? }	UpdateApiKeySchema	Update key information
llmConfig:deleteApiKey	R → M	{ id: string }	None	Delete a key
llmConfig:toggleApiKey	R → M	{ id, enabled }	ToggleApiKeySchema	Enable/disable a key

Zod validation rules:

// weight range constraint
weight: z.number().int().min(1).max(100)

// apiKey non-empty
apiKey: z.string().min(1)

// providerId non-empty
providerId: z.string().min(1)

Cache invalidation: All write operations (add/update/delete/toggle) call apiKeyPool.invalidateCache(providerId) upon completion to ensure the next request reloads from the database.

---

Extension Points

Tuning the Cooldown Strategy

Modify the constants in ApiKeyPoolService:

// More aggressive cooldown (for low-QPS scenarios)
private readonly DEFAULT_COOLDOWN_MS = 30_000;  // 30 s
private readonly MAX_COOLDOWN_MS = 5 * 60_000;  // 5 minutes

// More lenient cooldown (for high-QPS scenarios)
private readonly DEFAULT_COOLDOWN_MS = 120_000;  // 2 minutes
private readonly MAX_COOLDOWN_MS = 30 * 60_000;  // 30 minutes

Custom Key Selection Strategy

The current strategy is weighted round-robin. To use an alternative (e.g., least connections, random weighted), replace the selectWeightedRoundRobin() method.

Environment Variable Keys

Key values prefixed with $ are automatically expanded to environment variables:

$OPENAI_API_KEY → process.env.OPENAI_API_KEY

This is handled by the ApiKeyResolver function, implemented in CompletionService.resolveApiKey().

---

Related Files

File	Relationship
capabilities/llm/completion/CompletionService.ts	Consumer: calls the pool via resolveApiKeyForRequest()
workers/db/apiKeys.ts	Data source: provides loadKeys and CRUD operations
workers/types.ts	DB Worker type definitions
routers/llm/ApiKeyRouter.ts	IPC layer: frontend key management operations
renderer/.../ApiKeyPoolSection.tsx	Frontend UI: key list CRUD
shared/llm-config.ts	ApiKeyEntry type definition