feat(agent-core): add API key pool for concurrent subagent key rotation

.changeset/subagent-api-key-pool.md

@@ -0,0 +1,6 @@
+---
+"@moonshot-ai/agent-core": minor
+"@moonshot-ai/kimi-code": minor
+---
+
+Add API key pool for parallel subagent execution. When multiple `KIMI_API_KEY*` environment variables are configured, subagents rotate through them to avoid rate-limit contention, and failed keys are temporarily cooled down on retryable errors.

packages/agent-core/src/rpc/core-impl.ts

@@ -29,6 +29,7 @@ import {
   ProviderManager, type BearerTokenProvider,
   type OAuthTokenProviderResolver
 } from '../session/provider-manager';
+import { ApiKeyPool } from '../session/api-key-pool';
 import { SessionAPIImpl } from '../session/rpc';
 import { normalizeWorkDir, SessionStore } from '../session/store';
 import { noopTelemetryClient, withTelemetryContext, type TelemetryClient } from '../telemetry';
@@ -118,6 +119,7 @@ export class KimiCore implements PromisableMethods<CoreAPI> {
   private readonly resolveOAuthTokenProvider: OAuthTokenProviderResolver | undefined;
   private readonly skillDirs: readonly string[];
   private readonly sessionStore: SessionStore;
+  private readonly apiKeyPool: ApiKeyPool | undefined;
   readonly plugins: PluginManager;
   private pluginsReady: Promise<void>;
   private pluginsLoadError: Error | undefined;
@@ -143,6 +145,7 @@ export class KimiCore implements PromisableMethods<CoreAPI> {
     this.resolveOAuthTokenProvider = options.resolveOAuthTokenProvider;
     this.skillDirs = options.skillDirs ?? [];
     this.telemetry = options.telemetry ?? noopTelemetryClient;
+    this.apiKeyPool = process.env['KIMI_API_KEY_POOL'] !== undefined ? (ApiKeyPool.fromEnv() ?? undefined) : undefined;
     ensureKimiHome(this.homeDir);
     this.config = readConfigFile(this.configPath);
     this.sessionStore = new SessionStore(this.homeDir);
@@ -669,6 +672,7 @@ export class KimiCore implements PromisableMethods<CoreAPI> {
       kimiRequestHeaders: this.kimiRequestHeaders,
       resolveOAuthTokenProvider: this.resolveOAuthTokenProvider,
       promptCacheKey: sessionId,
+      apiKeyPool: this.apiKeyPool,
     });
   }
 

packages/agent-core/src/session/api-key-pool.ts

@@ -0,0 +1,135 @@
+/**
+ * APIKeyPool — round-robin allocator for multiple API keys.
+ *
+ * Designed for parallel subagent execution so that concurrent agents
+ * do not hammer a single key's rate-limit quota.
+ *
+ * Keys are read from environment variables:
+ *   KIMI_API_KEY, KIMI_API_KEY_1, KIMI_API_KEY_2, … up to KIMI_API_KEY_99
+ *
+ * A pool is only created when ≥2 keys are found; otherwise `fromEnv`
+ * returns `null` and callers fall back to the root provider's key.
+ */
+
+interface KeyState {
+  consecutiveFailures: number;
+  cooldownUntil: number | null;
+}
+
+const COOLDOWN_MS = [30_000, 300_000, 1_800_000] as const;
+
+function cooldownForFailures(failures: number): number {
+  if (failures <= 0) return 0;
+  if (failures <= COOLDOWN_MS.length) return COOLDOWN_MS.at(failures - 1)!;
+  return COOLDOWN_MS.at(-1)!;
+}
+
+export class ApiKeyPool {
+  private readonly keys: readonly string[];
+  private _index = 0;
+  private readonly states: Map<string, KeyState>;
+
+  /**
+   * Build a pool from environment variables.
+   *
+   * Collects `PREFIX`, `PREFIX_1`, `PREFIX_2`, … up to `PREFIX_99`.
+   * Returns `null` when fewer than 2 keys are found.
+   */
+  static fromEnv(prefix = 'KIMI_API_KEY'): ApiKeyPool | null {
+    const keys: string[] = [];
+    const primary = process.env[prefix];
+    if (primary !== undefined && primary.trim().length > 0) {
+      keys.push(primary.trim());
+    }
+    for (let i = 1; i < 100; i++) {
+      const val = process.env[`${prefix}_${i}`];
+      if (val !== undefined && val.trim().length > 0) {
+        keys.push(val.trim());
+      }
+    }
+    if (keys.length < 2) {
+      return null;
+    }
+    return new ApiKeyPool(keys);
+  }
+
+  constructor(keys: readonly string[]) {
+    if (keys.length === 0) {
+      throw new Error('Key pool cannot be empty');
+    }
+    this.keys = keys.slice();
+    this.states = new Map<string, KeyState>();
+    for (const key of this.keys) {
+      this.states.set(key, { consecutiveFailures: 0, cooldownUntil: null });
+    }
+  }
+
+  /** Number of keys in the pool. */
+  get keyCount(): number {
+    return this.keys.length;
+  }
+
+  /**
+   * Acquire the next key in rotation.
+   *
+   * Skips keys that are in cooldown. If every key is cooling down,
+   * falls back to round-robin across the entire pool.
+   */
+  acquire(): string {
+    const now = Date.now();
+    for (let i = 0; i < this.keys.length; i++) {
+      const key = this.keys[this._index]!;
+      this._index = (this._index + 1) % this.keys.length;
+      const state = this.states.get(key);
+      if (state === undefined) {
+        continue;
+      }
+      if (state.cooldownUntil !== null) {
+        if (now < state.cooldownUntil) {
+          continue;
+        }
+        // Cooldown expired — reset the key to healthy.
+        this.states.set(key, { consecutiveFailures: 0, cooldownUntil: null });
+      }
+      return key;
+    }
+    // All keys in cooldown — fall back to round-robin.
+    const key = this.keys[this._index]!;
+    this._index = (this._index + 1) % this.keys.length;
+    return key;
+  }
+
+  /**
+   * Record a failure for the given key.
+   *
+   * Applies exponential cooldown:
+   *   1st failure → 30s
+   *   2nd failure → 5min
+   *   3rd+ failure → 30min
+   */
+  recordFailure(key: string): void {
+    const state = this.states.get(key);
+    if (state === undefined) {
+      return;
+    }
+    const failures = state.consecutiveFailures + 1;
+    this.states.set(key, {
+      consecutiveFailures: failures,
+      cooldownUntil: Date.now() + cooldownForFailures(failures),
+    });
+  }
+
+  /** Clear the failure state for a key (e.g. after a successful call). */
+  resetKey(key: string): void {
+    const state = this.states.get(key);
+    if (state === undefined) {
+      return;
+    }
+    // Don't clear an active cooldown that may have been set by a concurrent
+    // failure while this request was still in flight.
+    if (state.cooldownUntil !== null && Date.now() < state.cooldownUntil) {
+      return;
+    }
+    this.states.set(key, { consecutiveFailures: 0, cooldownUntil: null });
+  }
+}

packages/agent-core/src/session/provider-manager.ts

@@ -1,6 +1,7 @@
 import type { Logger } from '#/logging/types';
 import type { ProviderConfig as KosongProviderConfig, ModelCapability, ProviderRequestAuth } from '@moonshot-ai/kosong';
-import { APIStatusError, createProvider, UNKNOWN_CAPABILITY } from '@moonshot-ai/kosong';
+import { ApiKeyPool } from './api-key-pool';
+import { APIStatusError, createProvider, isRetryableGenerateError, UNKNOWN_CAPABILITY } from '@moonshot-ai/kosong';
 import type { KimiConfig, ModelAlias, OAuthRef, ProviderConfig } from '../config';
 import { ErrorCodes, isKimiError, KimiError } from '../errors';
 
@@ -24,6 +25,7 @@ interface ProviderManagerOptions {
   readonly kimiRequestHeaders?: Record<string, string>;
   readonly resolveOAuthTokenProvider?: OAuthTokenProviderResolver;
   readonly promptCacheKey?: string;
+  readonly apiKeyPool?: ApiKeyPool;
 }
 
 type AuthorizedRequest = <T>(
@@ -123,68 +125,99 @@ export class ProviderManager implements ModelProvider {
   ): AuthorizedRequest | undefined {
     const { providerName } = this.resolveProviderConfig(model);
     const providerConfig = this.config.providers[providerName];
-    if (providerConfig?.oauth === undefined) return undefined;
 
-    if (providerApiKey(providerConfig) !== undefined) {
-      // oauth + apiKey on the same provider makes request auth ambiguous:
-      // provider construction would prefer apiKey while runtime auth resolves
-      // OAuth. Reject it so misconfiguration surfaces at model resolution.
-      throw new KimiError(
-        ErrorCodes.CONFIG_INVALID,
-        `Provider "${providerName}" has both apiKey and oauth set in config.toml — they are mutually exclusive. Remove one.`,
-      );
-    }
+    // OAuth path
+    if (providerConfig?.oauth !== undefined) {
+      if (providerApiKey(providerConfig) !== undefined) {
+        // oauth + apiKey on the same provider makes request auth ambiguous:
+        // provider construction would prefer apiKey while runtime auth resolves
+        // OAuth. Reject it so misconfiguration surfaces at model resolution.
+        throw new KimiError(
+          ErrorCodes.CONFIG_INVALID,
+          `Provider "${providerName}" has both apiKey and oauth set in config.toml — they are mutually exclusive. Remove one.`,
+        );
+      }
 
-    const loginRequired = (cause?: unknown): KimiError =>
-      new KimiError(
-        ErrorCodes.AUTH_LOGIN_REQUIRED,
-        `OAuth provider "${providerName}" requires login before it can be used.`,
-        cause === undefined ? undefined : { cause },
-      );
+      const loginRequired = (cause?: unknown): KimiError =>
+        new KimiError(
+          ErrorCodes.AUTH_LOGIN_REQUIRED,
+          `OAuth provider "${providerName}" requires login before it can be used.`,
+          cause === undefined ? undefined : { cause },
+        );
+
+      const tokenProvider = this.options.resolveOAuthTokenProvider?.(providerName, providerConfig.oauth);
+      if (tokenProvider === undefined) {
+        return async () => {
+          throw loginRequired();
+        };
+      }
 
-    const tokenProvider = this.options.resolveOAuthTokenProvider?.(providerName, providerConfig.oauth);
-    if (tokenProvider === undefined) {
-      return async () => {
-        throw loginRequired();
+      const log = options?.log;
+      const fetchAuth = async (force: boolean): Promise<ProviderRequestAuth> => {
+        let apiKey: string;
+        try {
+          apiKey = await tokenProvider.getAccessToken(force ? { force: true } : undefined);
+        } catch (error) {
+          if (!isKimiError(error) || error.code !== ErrorCodes.AUTH_LOGIN_REQUIRED) {
+            log?.warn('oauth token fetch failed', { providerName, error });
+          }
+          throw loginRequired(error);
+        }
+        if (apiKey.trim().length === 0) throw loginRequired();
+        return { apiKey };
       };
-    }
 
-    const log = options?.log;
-    const fetchAuth = async (force: boolean): Promise<ProviderRequestAuth> => {
-      let apiKey: string;
-      try {
-        apiKey = await tokenProvider.getAccessToken(force ? { force: true } : undefined);
-      } catch (error) {
-        if (!isKimiError(error) || error.code !== ErrorCodes.AUTH_LOGIN_REQUIRED) {
-          log?.warn('oauth token fetch failed', { providerName, error });
+      return async (request) => {
+        let auth = await fetchAuth(false);
+        for (let refreshed = false; ; refreshed = true) {
+          try {
+            return await request(auth);
+          } catch (error) {
+            if (!(error instanceof APIStatusError) || error.statusCode !== 401) throw error;
+            if (refreshed) {
+              throw new KimiError(
+                ErrorCodes.AUTH_LOGIN_REQUIRED,
+                'OAuth provider credentials were rejected. Send /login to login.',
+                {
+                  cause: error,
+                  details: { statusCode: error.statusCode, requestId: error.requestId },
+                },
+              );
+            }
+            auth = await fetchAuth(true);
+          }
         }
-        throw loginRequired(error);
-      }
-      if (apiKey.trim().length === 0) throw loginRequired();
-      return { apiKey };
-    };
+      };
+    }
 
-    return async (request) => {
-      let auth = await fetchAuth(false);
-      for (let refreshed = false; ; refreshed = true) {
+    // Key pool path — only for kimi provider when a pool is configured
+    // and the provider does not already have an explicit apiKey.
+    if (
+      providerConfig?.type === 'kimi' &&
+      this.options.apiKeyPool !== undefined &&
+      providerApiKey(providerConfig) === undefined
+    ) {
+      const pool = this.options.apiKeyPool;
+      return async (request) => {
+        const key = pool.acquire();
+        const auth: ProviderRequestAuth = { apiKey: key };
         try {
-          return await request(auth);
+          const result = await request(auth);
+          pool.resetKey(key);
+          return result;
         } catch (error) {
-          if (!(error instanceof APIStatusError) || error.statusCode !== 401) throw error;
-          if (refreshed) {
-            throw new KimiError(
-              ErrorCodes.AUTH_LOGIN_REQUIRED,
-              'OAuth provider credentials were rejected. Send /login to login.',
-              {
-                cause: error,
-                details: { statusCode: error.statusCode, requestId: error.requestId },
-              },
-            );
+          if (
+            isRetryableGenerateError(error) ||
+            (error instanceof APIStatusError && error.statusCode === 401)
+          ) {
+            pool.recordFailure(key);
           }
-          auth = await fetchAuth(true);
+          throw error;
         }
-      }
-    };
+      };
+    }
+
+    return undefined;
   }
 }