Agents/Providers — Token Management

Key Rotation Health Monitoring Cooldown Tracking File Locking

Contents

Overview

flowchart TD
  subgraph Input["Request Phase"]
    REQ["Agent request\n(provider + model)"]
    STORE["auth-profiles.ts\nLoad auth store\n(file-locked)"]
  end
  subgraph Resolution["Profile Resolution"]
    ORDER["resolveAuthProfileOrder()"]
    FILTER["Filter expired &\ncooldown profiles"]
    SORT["Sort by recency\n+ random for fairness"]
  end
  subgraph Health["Health Check"]
    HEALTH["AuthProfileHealth\nok | expiring | expired\nmissing | static"]
    COOLDOWN["Cooldown tracker\nnow + TTL expiry"]
    SOONEST["getSoonestCooldownExpiry()"]
  end
  subgraph Output["Selected Credential"]
    CRED["API key / token\nfor provider request"]
    FAIL["markAuthProfileFailure()\nSet cooldown on error"]
  end
  REQ --> STORE
  STORE --> ORDER
  ORDER --> FILTER
  FILTER --> SORT
  SORT --> HEALTH
  HEALTH --> CRED
  CRED -->|"401/403/429"| FAIL
  FAIL --> COOLDOWN
  COOLDOWN --> SOONEST
  SOONEST -->|"retry after expiry"| FILTER

The flow begins when the agent runtime needs credentials for a provider request. The auth store is loaded with file locking to prevent corruption from concurrent access. resolveAuthProfileOrder() applies configured ordering, filters out expired and cooling-down profiles, then sorts by recency with random jitter for fair rotation across equivalent profiles.

Decision Points

How It Works

Token management operates in three phases: resolution (selecting the best profile), monitoring (tracking health status), and recovery (cooldown management after failures).

Phase 1: Auth Profile Resolution

The resolveAuthProfileOrder() function implements the core selection algorithm. It loads the auth store, applies any configured ordering preferences, filters out unusable profiles, and returns the best available credential.

// auth-profiles.ts — simplified resolution flow
export async function resolveAuthProfileOrder(
  providerId: string,
  configOrder?: string[],
): Promise<ResolvedAuthProfile | null> {
  // 1. Load auth store with file locking
  const store = await loadAuthStore({
    retries: AUTH_STORE_LOCK_OPTIONS.retries,       // 10
    minTimeout: 100,                                 // ms
    maxTimeout: 10_000,                              // ms
    factor: 2,                                       // exponential backoff
    stale: 30_000,                                   // 30s staleness
  });

  // 2. Get all profiles for this provider
  const profiles = store.profiles.filter(p => p.providerId === providerId);

  // 3. Apply config ordering (if provided)
  const ordered = configOrder
    ? reorderByConfig(profiles, configOrder)
    : profiles;

  // 4. Filter expired and in-cooldown profiles
  const now = Date.now();
  const available = ordered.filter(p => {
    if (p.expiresAt && p.expiresAt < now) return false;       // expired
    if (p.cooldownExpiry && p.cooldownExpiry > now) return false; // in cooldown
    return true;
  });

  // 5. Sort by recency with random jitter for fair rotation
  available.sort((a, b) => {
    const recencyDiff = (b.lastUsedAt ?? 0) - (a.lastUsedAt ?? 0);
    return recencyDiff + (Math.random() - 0.5) * 1000;
  });

  return available[0] ?? null;
}

Phase 2: Health Monitoring

The AuthProfileHealth type tracks the health status of each stored credential. The health check runs before every profile selection and after every API response.

Phase 3: Cooldown & Recovery

When an API request fails with an auth or rate-limit error, markAuthProfileFailure() places the profile into cooldown. The cooldown expiry is calculated as now + ttl where the TTL varies by error type.

Auth Profile State Machine

stateDiagram-v2
  [*] --> Valid : Profile loaded
  Valid --> InUse : Selected for request
  InUse --> Valid : API success
  InUse --> Cooldown : API error (401/403/429)
  Cooldown --> Probe : shouldAllowCooldownProbe()
  Probe --> Valid : Probe succeeds
  Probe --> Cooldown : Probe fails
  Cooldown --> Valid : Cooldown expiry passes
  Valid --> Expiring : Near expiry threshold
  Expiring --> Expired : expiresAt reached
  Expired --> Valid : Manual refresh / CLI sync
  Expired --> [*] : Removed from store

  note right of Cooldown
    cooldownExpiry = now + ttl
    ttl varies by error type:
    401/403 → longer TTL
    429 → shorter TTL
  end note

  note right of Probe
    Single retry with short timeout
    Only allowed for specific error reasons
  end note

The state machine above shows the full lifecycle of an auth profile. Most profiles spend their time in the Valid ↔ InUse cycle. The Cooldown state is entered on API failure and exited either by timer expiry or by a successful cooldown probe. The probe mechanism is controlled by shouldAllowCooldownProbeForReason() which only allows probes for transient errors (rate limits), not permanent failures (invalid key).

Annotated Code Walkthrough

Walk through the cooldown marking and recovery flow.

// auth-profiles.ts — marking a failure
export async function markAuthProfileFailure(
  profileId: string,
  error: ProviderError,
): Promise<void> {
  const store = await loadAuthStore(AUTH_STORE_LOCK_OPTIONS);
  const profile = store.profiles.find(p => p.id === profileId);
  if (!profile) return;

  const ttl = getCooldownTTL(error.status);
  profile.cooldownExpiry = Date.now() + ttl;
  profile.lastError = {
    status: error.status,
    message: error.message,
    timestamp: Date.now(),
  };

  await saveAuthStore(store);
}

function getCooldownTTL(status: number): number {
  switch (status) {
    case 429: return 30_000;   // Rate limit: 30s cooldown
    case 401: return 300_000;  // Unauthorized: 5min cooldown
    case 403: return 600_000;  // Forbidden: 10min cooldown
    default:  return 60_000;   // Other: 1min cooldown
  }
}

// auth-profiles.ts — finding nearest available retry
export function getSoonestCooldownExpiry(
  profiles: AuthProfile[],
): number | null {
  const now = Date.now();
  let soonest: number | null = null;

  for (const p of profiles) {
    if (!p.cooldownExpiry || p.cooldownExpiry <= now) {
      return null; // At least one profile is already available
    }
    if (soonest === null || p.cooldownExpiry < soonest) {
      soonest = p.cooldownExpiry;
    }
  }

  return soonest;
}

// auth-profiles.ts — external CLI sync
const EXTERNAL_CLI_SYNC_TTL_MS = 900_000; // 15 minutes
let lastExternalSync = 0;

async function maybeSyncExternalCredentials(
  store: AuthStore,
): Promise<boolean> {
  const now = Date.now();
  if (now - lastExternalSync < EXTERNAL_CLI_SYNC_TTL_MS) {
    return false; // TTL not yet expired
  }

  lastExternalSync = now;
  const externalProfiles = await loadFromExternalCLI();

  let changed = false;
  for (const ext of externalProfiles) {
    const existing = store.profiles.find(
      p => p.providerId === ext.providerId && p.keyHash === ext.keyHash,
    );
    if (!existing) {
      store.profiles.push(ext);
      changed = true;
    }
  }

  return changed;
}

File Locking Details

The auth store is a JSON file on disk, and multiple agent processes may attempt to read or write it concurrently. The file-locking mechanism uses exponential backoff with configurable retry parameters:

Edge Cases & Error Handling

Configuration

Related Topics

See also:

• Providers overview — the parent sub-module page
• Adapter Pattern — how provider wrappers consume the selected credentials
• Model Selection — model fallback chains interact with auth cooldowns
• Streaming pipeline — where auth errors trigger cooldown marking

Directory Map

pi-embedded-runner/
├── auth-profiles.ts             ← Auth store management, profile resolution, file locking
├── auth-health.ts               ← AuthProfileHealth status types and health checks
├── model.ts                     ← Model registry (consumes auth profiles for provider auth)
└── run/
    ├── setup.ts                 ← Passes resolved auth to stream wrappers
    └── attempt.ts               ← Triggers markAuthProfileFailure() on API errors