diff --git a/config/platform-infra/sub2api-codex-pool.yaml b/config/platform-infra/sub2api-codex-pool.yaml index ae80161d..ec9af3f1 100644 --- a/config/platform-infra/sub2api-codex-pool.yaml +++ b/config/platform-infra/sub2api-codex-pool.yaml @@ -4,6 +4,7 @@ pool: apiKeySecretName: sub2api-codex-pool-api-key apiKeySecretKey: API_KEY minOwnerBalanceUsd: 1000 + defaultAccountCapacity: 5 profiles: entries: - profile: default diff --git a/scripts/src/platform-infra-sub2api-codex.ts b/scripts/src/platform-infra-sub2api-codex.ts index 4afef069..04e31afb 100644 --- a/scripts/src/platform-infra-sub2api-codex.ts +++ b/scripts/src/platform-infra-sub2api-codex.ts @@ -47,6 +47,7 @@ interface CodexProfile { openaiResponsesWebSocketsV2Mode: OpenAIResponsesWebSocketsV2Mode | null; upstreamUserAgent: string | null; priority: number; + capacity: number; authOpenAIKeyShape: string; ok: boolean; error: string | null; @@ -60,6 +61,7 @@ interface CodexPoolConfig { apiKeySecretName: string; apiKeySecretKey: string; minOwnerBalanceUsd: number; + defaultAccountCapacity: number; profiles: CodexPoolProfileConfig[]; publicExposure: CodexPoolPublicExposureConfig; localCodex: CodexPoolLocalCodexConfig; @@ -75,6 +77,7 @@ interface CodexPoolProfileConfig { openaiResponsesWebSocketsV2Mode: OpenAIResponsesWebSocketsV2Mode | null; upstreamUserAgent: string | null; priority: number; + capacity: number | null; } interface CodexPoolPublicExposureConfig { @@ -227,6 +230,7 @@ async function codexPoolSync(config: UniDeskConfig, options: SyncOptions): Promi apiKeySecretName: pool.apiKeySecretName, apiKeySecretKey: pool.apiKeySecretKey, minOwnerBalanceUsd: pool.minOwnerBalanceUsd, + defaultAccountCapacity: pool.defaultAccountCapacity, }, profiles: profiles.map((profile) => ({ profile: profile.profile, @@ -243,6 +247,7 @@ async function codexPoolSync(config: UniDeskConfig, options: SyncOptions): Promi openaiResponsesWebSocketsV2Mode: profile.openaiResponsesWebSocketsV2Mode, upstreamUserAgent: profile.upstreamUserAgent, priority: profile.priority, + capacity: profile.capacity, })), }; const result = await capture(config, g14K3sRoute, ["script"], syncScript(payload, pool)); @@ -426,6 +431,7 @@ function collectCodexProfiles(): CodexProfile[] { openaiResponsesWebSocketsV2Mode: entry.openaiResponsesWebSocketsV2Mode, upstreamUserAgent: entry.upstreamUserAgent, priority: entry.priority, + capacity: entry.capacity ?? pool.defaultAccountCapacity, authOpenAIKeyShape: existsSync(authPath) ? "unknown" : "missing", ok: false, error: null, @@ -491,6 +497,7 @@ function discoverCodexProfileConfigs(codexDir: string): CodexPoolProfileConfig[] openaiResponsesWebSocketsV2Mode: null, upstreamUserAgent: null, priority: 1, + capacity: null, }; }); } @@ -518,6 +525,7 @@ function readCodexPoolConfig(): CodexPoolConfig { apiKeySecretName: stringValue(pool.apiKeySecretName) ?? defaults.apiKeySecretName, apiKeySecretKey: stringValue(pool.apiKeySecretKey) ?? defaults.apiKeySecretKey, minOwnerBalanceUsd: numberValue(pool.minOwnerBalanceUsd) ?? defaults.minOwnerBalanceUsd, + defaultAccountCapacity: readAccountCapacity(pool.defaultAccountCapacity, "pool.defaultAccountCapacity"), profiles: readProfileConfig(parsed.profiles, defaults.profiles), publicExposure: readPublicExposureConfig(parsed.publicExposure, defaults.publicExposure), localCodex: readLocalCodexConfig(parsed.localCodex, defaults.localCodex), @@ -538,6 +546,7 @@ function defaultCodexPoolConfig(): CodexPoolConfig { apiKeySecretName: defaultPoolApiKeySecretName, apiKeySecretKey: defaultPoolApiKeySecretKey, minOwnerBalanceUsd: defaultMinOwnerBalanceUsd, + defaultAccountCapacity: 5, profiles: [], publicExposure: { enabled: false, @@ -588,6 +597,7 @@ function readProfileConfig(value: unknown, defaults: CodexPoolProfileConfig[]): const openaiResponsesWebSocketsV2Mode = readOpenAIResponsesWebSocketsV2Mode(entry.openaiResponsesWebSocketsV2Mode, `profiles.entries[${index}].openaiResponsesWebSocketsV2Mode`); const upstreamUserAgent = readUpstreamUserAgent(entry.upstreamUserAgent, `profiles.entries[${index}].upstreamUserAgent`); const priority = readAccountPriority(entry.priority, `profiles.entries[${index}].priority`); + const capacity = entry.capacity === undefined || entry.capacity === null ? null : readAccountCapacity(entry.capacity, `profiles.entries[${index}].capacity`); return { profile, accountName, @@ -598,6 +608,7 @@ function readProfileConfig(value: unknown, defaults: CodexPoolProfileConfig[]): openaiResponsesWebSocketsV2Mode, upstreamUserAgent, priority, + capacity, }; }); } @@ -628,6 +639,14 @@ function readAccountPriority(value: unknown, key: string): number { return priority; } +function readAccountCapacity(value: unknown, key: string): number { + const capacity = numberValue(value); + if (capacity === null || !Number.isInteger(capacity) || capacity < 1 || capacity > 1000) { + throw new Error(`${codexPoolConfigPath}.${key} must be an integer from 1 to 1000`); + } + return capacity; +} + function readPublicExposureConfig(value: unknown, defaults: CodexPoolPublicExposureConfig): CodexPoolPublicExposureConfig { if (!isRecord(value)) return defaults; const masterFrpsValue = isRecord(value.masterFrps) ? value.masterFrps : {}; @@ -730,6 +749,7 @@ function redactProfile(profile: CodexProfile): Record { openaiResponsesWebSocketsV2Mode: profile.openaiResponsesWebSocketsV2Mode, upstreamUserAgent: profile.upstreamUserAgent, priority: profile.priority, + capacity: profile.capacity, apiKeyPresent: profile.apiKey !== null && profile.apiKey.length > 0, apiKeyBytes: profile.apiKey === null ? 0 : Buffer.byteLength(profile.apiKey, "utf8"), apiKeyFingerprint: profile.apiKey === null ? null : fingerprint(profile.apiKey), @@ -750,6 +770,7 @@ function poolTarget(pool = readCodexPoolConfig()): Record { groupName: pool.groupName, apiKeyName: pool.apiKeyName, apiKeySecret: `${namespace}/${pool.apiKeySecretName}.${pool.apiKeySecretKey}`, + defaultAccountCapacity: pool.defaultAccountCapacity, valuesPrinted: false, }; } @@ -1256,6 +1277,19 @@ function validateScript(pool: CodexPoolConfig): string { return remotePythonScript("validate", "", pool); } +function desiredAccountCapacityMap(pool: CodexPoolConfig): Record { + const codexDir = join(homedir(), ".codex"); + const seenAccountNames = new Set(); + const configs = pool.profiles.length > 0 ? pool.profiles : discoverCodexProfileConfigs(codexDir); + const capacities: Record = {}; + for (const entry of configs) { + const accountName = entry.accountName ?? uniqueAccountName(entry.profile, seenAccountNames); + seenAccountNames.add(accountName); + capacities[accountName] = entry.capacity ?? pool.defaultAccountCapacity; + } + return capacities; +} + function remotePythonScript(mode: "sync" | "validate", encodedPayload: string, pool: CodexPoolConfig): string { return ` set -u @@ -1279,6 +1313,8 @@ POOL_API_KEY_NAME = "${pool.apiKeyName}" POOL_API_KEY_SECRET_NAME = "${pool.apiKeySecretName}" POOL_API_KEY_SECRET_KEY = "${pool.apiKeySecretKey}" MIN_OWNER_BALANCE_USD = ${JSON.stringify(pool.minOwnerBalanceUsd)} +POOL_DEFAULT_ACCOUNT_CAPACITY = ${JSON.stringify(pool.defaultAccountCapacity)} +EXPECTED_ACCOUNT_CAPACITIES = ${JSON.stringify(desiredAccountCapacityMap(pool))} MODE = "${mode}" PAYLOAD_B64 = "${encodedPayload}" @@ -1498,7 +1534,7 @@ def account_payload(profile, group_id): "type": "apikey", "credentials": credentials, "extra": extra, - "concurrency": 1, + "concurrency": int(profile.get("capacity", 5) or 5), "priority": int(profile.get("priority", 1) or 1), "rate_multiplier": 1, "load_factor": 1, @@ -1534,6 +1570,8 @@ def ensure_accounts(token, profiles, group_id): "apiKeyFingerprint": profile["apiKeyFingerprint"], "openaiResponsesWebSocketsV2Mode": profile.get("openaiResponsesWebSocketsV2Mode"), "priority": int(profile.get("priority", 1) or 1), + "capacity": int(profile.get("capacity", 5) or 5), + "runtimeConcurrency": data.get("concurrency") if isinstance(data, dict) else None, "upstreamUserAgentConfigured": bool(profile.get("upstreamUserAgent")), "valuesPrinted": False, }) @@ -1687,6 +1725,49 @@ def validate_gateway(api_key): "valuesPrinted": False, } +def account_capacity_status(token): + accounts = list_accounts(token) + by_name = {item.get("name"): item for item in accounts if isinstance(item.get("name"), str)} + items = [] + missing = [] + mismatched = [] + for name in sorted(EXPECTED_ACCOUNT_CAPACITIES): + expected = int(EXPECTED_ACCOUNT_CAPACITIES[name]) + account = by_name.get(name) + if account is None: + missing.append(name) + items.append({ + "accountName": name, + "accountId": None, + "expectedCapacity": expected, + "runtimeConcurrency": None, + "ok": False, + }) + continue + runtime = account.get("concurrency") + ok = runtime == expected + if not ok: + mismatched.append(name) + items.append({ + "accountName": name, + "accountId": account.get("id"), + "expectedCapacity": expected, + "runtimeConcurrency": runtime, + "priority": account.get("priority"), + "status": account.get("status"), + "schedulable": account.get("schedulable"), + "ok": ok, + }) + return { + "ok": len(missing) == 0 and len(mismatched) == 0, + "defaultAccountCapacity": POOL_DEFAULT_ACCOUNT_CAPACITY, + "desired": len(EXPECTED_ACCOUNT_CAPACITIES), + "missing": missing, + "mismatched": mismatched, + "items": items, + "valuesPrinted": False, + } + def api_key_preview(api_key): if len(api_key) <= 14: return "***" @@ -1703,12 +1784,13 @@ def run_sync(): if group_id is None: raise RuntimeError("pool group id missing after ensure") account_results, pruned_account_results = ensure_accounts(token, profiles, group_id) + capacity_status = account_capacity_status(token) api_key, secret_action, secret_apply_stdout = ensure_api_key_secret(group_id) api_key_result = ensure_sub2api_api_key(token, api_key, group_id) owner_balance = ensure_pool_owner_balance(token, api_key_result["userId"]) gateway = validate_gateway(api_key) return { - "ok": gateway["ok"] is True, + "ok": gateway["ok"] is True and capacity_status["ok"] is True, "mode": "sync", "namespace": NAMESPACE, "serviceDns": SERVICE_DNS, @@ -1724,6 +1806,7 @@ def run_sync(): "prunedItems": pruned_account_results, "valuesPrinted": False, }, + "capacity": capacity_status, "apiKey": { "name": POOL_API_KEY_NAME, "secret": f"{NAMESPACE}/{POOL_API_KEY_SECRET_NAME}.{POOL_API_KEY_SECRET_KEY}", @@ -1749,9 +1832,10 @@ def run_validate(): owner_balance = None if key_item is not None and key_item.get("user_id") is not None: owner_balance = ensure_pool_owner_balance(token, key_item["user_id"]) + capacity_status = account_capacity_status(token) gateway = validate_gateway(api_key) return { - "ok": gateway["ok"] is True, + "ok": gateway["ok"] is True and capacity_status["ok"] is True, "mode": "validate", "namespace": NAMESPACE, "serviceDns": SERVICE_DNS, @@ -1765,6 +1849,7 @@ def run_validate(): "valuesPrinted": False, }, "ownerBalance": owner_balance, + "capacity": capacity_status, "validation": {"gatewayModels": gateway}, }