fix: harden sub2api codex ws routing

2026-06-09 08:17:32 +00:00
parent 7d157d7a94
commit b55320c185
6 changed files with 252 additions and 5 deletions
@@ -16,6 +16,7 @@ const syntaxFiles = [
  "scripts/playwright-cli.ts",
  "scripts/playwright-cli-contract-test.ts",
  "scripts/platform-infra-sub2api-codex-local-config-contract-test.ts",
+  "scripts/platform-infra-sub2api-codex-routing-contract-test.ts",
  "scripts/src/playwright-cli.ts",
  "scripts/src/check.ts",
  "scripts/src/artifact-registry.ts",
@@ -384,6 +385,7 @@ export function runChecks(config: UniDeskConfig, options: CheckOptions = default
      fileItem("scripts/gh-cli-pr-contract-test.ts"),
      fileItem("scripts/playwright-cli-contract-test.ts"),
      fileItem("scripts/platform-infra-sub2api-codex-local-config-contract-test.ts"),
+      fileItem("scripts/platform-infra-sub2api-codex-routing-contract-test.ts"),
      fileItem("scripts/code-queue-pr-preflight-example.ts"),
      fileItem("scripts/schedule-cli-contract-test.ts"),
      fileItem("scripts/server-cleanup-plan-contract-test.ts"),
@@ -449,6 +451,7 @@ export function runChecks(config: UniDeskConfig, options: CheckOptions = default
    items.push(commandItem("gh:pr-contract", ["bun", "scripts/gh-cli-pr-contract-test.ts"], 30_000));
    items.push(commandItem("playwright:cli-wrapper-contract", ["bun", "scripts/playwright-cli-contract-test.ts"], 30_000));
    items.push(commandItem("platform-infra:sub2api-codex-local-config-contract", ["bun", "scripts/platform-infra-sub2api-codex-local-config-contract-test.ts"], 30_000));
+    items.push(commandItem("platform-infra:sub2api-codex-routing-contract", ["bun", "scripts/platform-infra-sub2api-codex-routing-contract-test.ts"], 30_000));
    items.push(commandItem("auth-broker:p0-contract", ["bun", "scripts/auth-broker-contract-test.ts"], 30_000));
    items.push(commandItem("d601:recovery-guardrails-contract", ["bun", "scripts/d601-recovery-guardrails-contract-test.ts"], 30_000));
    items.push(commandItem("hwlab:cd-wrapper-contract", ["bun", "scripts/hwlab-cd-wrapper-contract-test.ts"], 30_000));
@@ -48,6 +48,7 @@ interface CodexProfile {
  upstreamUserAgent: string | null;
  priority: number;
  capacity: number;
+  tempUnschedulable: CodexTempUnschedulablePolicy;
  authOpenAIKeyShape: string;
  ok: boolean;
  error: string | null;
@@ -55,6 +56,18 @@ interface CodexProfile {

 type OpenAIResponsesWebSocketsV2Mode = "off" | "ctx_pool" | "passthrough";

+export interface CodexTempUnschedulableRule {
+  statusCode: number;
+  keywords: string[];
+  durationMinutes: number;
+  description: string | null;
+}
+
+export interface CodexTempUnschedulablePolicy {
+  enabled: boolean;
+  rules: CodexTempUnschedulableRule[];
+}
+
 interface CodexPoolConfig {
  groupName: string;
  apiKeyName: string;
@@ -62,6 +75,7 @@ interface CodexPoolConfig {
  apiKeySecretKey: string;
  minOwnerBalanceUsd: number;
  defaultAccountCapacity: number;
+  defaultTempUnschedulable: CodexTempUnschedulablePolicy;
  profiles: CodexPoolProfileConfig[];
  publicExposure: CodexPoolPublicExposureConfig;
  localCodex: CodexPoolLocalCodexConfig;
@@ -78,6 +92,7 @@ interface CodexPoolProfileConfig {
  upstreamUserAgent: string | null;
  priority: number;
  capacity: number | null;
+  tempUnschedulable: CodexTempUnschedulablePolicy;
 }

 interface CodexPoolPublicExposureConfig {
@@ -265,6 +280,7 @@ async function codexPoolSync(config: UniDeskConfig, options: SyncOptions): Promi
      upstreamUserAgent: profile.upstreamUserAgent,
      priority: profile.priority,
      capacity: profile.capacity,
+      tempUnschedulable: profile.tempUnschedulable,
    })),
  };
  const result = await capture(config, g14K3sRoute, ["script"], syncScript(payload, pool));
@@ -454,6 +470,7 @@ function collectCodexProfiles(): CodexProfile[] {
      upstreamUserAgent: entry.upstreamUserAgent,
      priority: entry.priority,
      capacity: entry.capacity ?? pool.defaultAccountCapacity,
+      tempUnschedulable: entry.tempUnschedulable,
      authOpenAIKeyShape: existsSync(authPath) ? "unknown" : "missing",
      ok: false,
      error: null,
@@ -498,7 +515,7 @@ function collectCodexProfiles(): CodexProfile[] {
  });
 }

-function discoverCodexProfileConfigs(codexDir: string): CodexPoolProfileConfig[] {
+function discoverCodexProfileConfigs(codexDir: string, defaultTempUnschedulable = defaultCodexTempUnschedulablePolicy()): CodexPoolProfileConfig[] {
  return readdirSync(codexDir)
    .filter((file) => file === "config.toml" || file.startsWith("config.toml."))
    .sort((a, b) => {
@@ -520,6 +537,7 @@ function discoverCodexProfileConfigs(codexDir: string): CodexPoolProfileConfig[]
        upstreamUserAgent: null,
        priority: 1,
        capacity: null,
+        tempUnschedulable: defaultTempUnschedulable,
      };
    });
 }
@@ -541,6 +559,7 @@ function readCodexPoolConfig(): CodexPoolConfig {
  if (!isRecord(parsed)) throw new Error(`${codexPoolConfigPath} must contain a YAML object`);
  const pool = parsed.pool;
  if (!isRecord(pool)) throw new Error(`${codexPoolConfigPath}.pool must be a YAML object`);
+  const defaultTempUnschedulable = readTempUnschedulablePolicy(pool.defaultTempUnschedulable, "pool.defaultTempUnschedulable", defaults.defaultTempUnschedulable);
  const config: CodexPoolConfig = {
    groupName: stringValue(pool.groupName) ?? defaults.groupName,
    apiKeyName: stringValue(pool.apiKeyName) ?? defaults.apiKeyName,
@@ -548,7 +567,8 @@ function readCodexPoolConfig(): CodexPoolConfig {
    apiKeySecretKey: stringValue(pool.apiKeySecretKey) ?? defaults.apiKeySecretKey,
    minOwnerBalanceUsd: numberValue(pool.minOwnerBalanceUsd) ?? defaults.minOwnerBalanceUsd,
    defaultAccountCapacity: readAccountCapacity(pool.defaultAccountCapacity, "pool.defaultAccountCapacity"),
-    profiles: readProfileConfig(parsed.profiles, defaults.profiles),
+    defaultTempUnschedulable,
+    profiles: readProfileConfig(parsed.profiles, defaults.profiles, defaultTempUnschedulable),
    publicExposure: readPublicExposureConfig(parsed.publicExposure, defaults.publicExposure),
    localCodex: readLocalCodexConfig(parsed.localCodex, defaults.localCodex),
  };
@@ -569,6 +589,7 @@ function defaultCodexPoolConfig(): CodexPoolConfig {
    apiKeySecretKey: defaultPoolApiKeySecretKey,
    minOwnerBalanceUsd: defaultMinOwnerBalanceUsd,
    defaultAccountCapacity: 5,
+    defaultTempUnschedulable: defaultCodexTempUnschedulablePolicy(),
    profiles: [],
    publicExposure: {
      enabled: false,
@@ -605,7 +626,57 @@ function defaultCodexPoolConfig(): CodexPoolConfig {
  };
 }

-function readProfileConfig(value: unknown, defaults: CodexPoolProfileConfig[]): CodexPoolProfileConfig[] {
+export function defaultCodexTempUnschedulablePolicy(): CodexTempUnschedulablePolicy {
+  return {
+    enabled: true,
+    rules: [
+      {
+        statusCode: 401,
+        keywords: ["unauthorized", "invalid api key", "invalid_api_key", "authentication"],
+        durationMinutes: 30,
+        description: "Credential/auth failures should leave the scheduler quickly and retry after a cooldown.",
+      },
+      {
+        statusCode: 403,
+        keywords: ["forbidden", "access denied", "quota", "billing", "capacity"],
+        durationMinutes: 30,
+        description: "Permission, quota, or account-state failures should fail over to another account.",
+      },
+      {
+        statusCode: 429,
+        keywords: ["capacity", "rate limit", "rate_limit", "quota", "too many requests", "overloaded", "resource_exhausted"],
+        durationMinutes: 10,
+        description: "Capacity and rate-limit responses should cool down this account and use another account.",
+      },
+      {
+        statusCode: 500,
+        keywords: ["capacity", "overloaded", "temporarily unavailable", "temporary", "upstream"],
+        durationMinutes: 5,
+        description: "Transient upstream server failures should prefer another account for a short period.",
+      },
+      {
+        statusCode: 502,
+        keywords: ["capacity", "overloaded", "temporarily unavailable", "temporary", "upstream"],
+        durationMinutes: 5,
+        description: "Gateway upstream failures should prefer another account for a short period.",
+      },
+      {
+        statusCode: 503,
+        keywords: ["capacity", "overloaded", "temporarily unavailable", "temporary", "upstream"],
+        durationMinutes: 5,
+        description: "Service unavailable responses should prefer another account for a short period.",
+      },
+      {
+        statusCode: 529,
+        keywords: ["capacity", "overloaded", "temporarily unavailable", "temporary"],
+        durationMinutes: 10,
+        description: "Provider overloaded responses should cool down this account and use another account.",
+      },
+    ],
+  };
+}
+
+function readProfileConfig(value: unknown, defaults: CodexPoolProfileConfig[], defaultTempUnschedulable: CodexTempUnschedulablePolicy): CodexPoolProfileConfig[] {
  if (!isRecord(value)) return defaults;
  const entries = value.entries;
  if (!Array.isArray(entries)) throw new Error(`${codexPoolConfigPath}.profiles.entries must be a YAML array`);
@@ -629,6 +700,7 @@ function readProfileConfig(value: unknown, defaults: CodexPoolProfileConfig[]):
    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`);
+    const tempUnschedulable = readTempUnschedulablePolicy(entry.tempUnschedulable, `profiles.entries[${index}].tempUnschedulable`, defaultTempUnschedulable);
    return {
      profile,
      accountName,
@@ -640,6 +712,7 @@ function readProfileConfig(value: unknown, defaults: CodexPoolProfileConfig[]):
      upstreamUserAgent,
      priority,
      capacity,
+      tempUnschedulable,
    };
  });
 }
@@ -678,6 +751,66 @@ function readAccountCapacity(value: unknown, key: string): number {
  return capacity;
 }

+function readTempUnschedulablePolicy(value: unknown, key: string, fallback: CodexTempUnschedulablePolicy): CodexTempUnschedulablePolicy {
+  if (value === undefined || value === null) return cloneTempUnschedulablePolicy(fallback);
+  if (!isRecord(value)) throw new Error(`${codexPoolConfigPath}.${key} must be a YAML object`);
+  const enabled = readBooleanConfig(value.enabled, `${key}.enabled`, fallback.enabled);
+  const rules = value.rules === undefined || value.rules === null
+    ? cloneTempUnschedulablePolicy(fallback).rules
+    : readTempUnschedulableRules(value.rules, `${key}.rules`);
+  if (enabled && rules.length === 0) throw new Error(`${codexPoolConfigPath}.${key}.rules must not be empty when enabled=true`);
+  return { enabled, rules };
+}
+
+function readTempUnschedulableRules(value: unknown, key: string): CodexTempUnschedulableRule[] {
+  if (!Array.isArray(value)) throw new Error(`${codexPoolConfigPath}.${key} must be a YAML array`);
+  return value.map((entry, index) => {
+    if (!isRecord(entry)) throw new Error(`${codexPoolConfigPath}.${key}[${index}] must be an object`);
+    const statusCode = numberValue(entry.statusCode ?? entry.errorCode);
+    if (statusCode === null || !Number.isInteger(statusCode) || statusCode < 100 || statusCode > 599) {
+      throw new Error(`${codexPoolConfigPath}.${key}[${index}].statusCode must be an HTTP status code from 100 to 599`);
+    }
+    const durationMinutes = numberValue(entry.durationMinutes);
+    if (durationMinutes === null || !Number.isInteger(durationMinutes) || durationMinutes < 1 || durationMinutes > 1440) {
+      throw new Error(`${codexPoolConfigPath}.${key}[${index}].durationMinutes must be an integer from 1 to 1440`);
+    }
+    const keywords = readTempUnschedulableKeywords(entry.keywords, `${key}[${index}].keywords`);
+    const description = stringValue(entry.description);
+    if (description !== null && description.length > 240) throw new Error(`${codexPoolConfigPath}.${key}[${index}].description must be at most 240 characters`);
+    return { statusCode, keywords, durationMinutes, description };
+  });
+}
+
+function readTempUnschedulableKeywords(value: unknown, key: string): string[] {
+  if (!Array.isArray(value)) throw new Error(`${codexPoolConfigPath}.${key} must be a YAML array`);
+  const seen = new Set<string>();
+  const keywords: string[] = [];
+  for (const item of value) {
+    const keyword = stringValue(item);
+    if (keyword === null) throw new Error(`${codexPoolConfigPath}.${key} entries must be non-empty strings`);
+    if (keyword.length > 120) throw new Error(`${codexPoolConfigPath}.${key} entries must be at most 120 characters`);
+    if (/[\r\n]/u.test(keyword)) throw new Error(`${codexPoolConfigPath}.${key} entries must not contain newlines`);
+    const normalized = keyword.toLowerCase();
+    if (seen.has(normalized)) continue;
+    seen.add(normalized);
+    keywords.push(keyword);
+  }
+  if (keywords.length === 0) throw new Error(`${codexPoolConfigPath}.${key} must not be empty`);
+  return keywords;
+}
+
+function cloneTempUnschedulablePolicy(policy: CodexTempUnschedulablePolicy): CodexTempUnschedulablePolicy {
+  return {
+    enabled: policy.enabled,
+    rules: policy.rules.map((rule) => ({
+      statusCode: rule.statusCode,
+      keywords: [...rule.keywords],
+      durationMinutes: rule.durationMinutes,
+      description: rule.description,
+    })),
+  };
+}
+
 function readBooleanConfig(value: unknown, key: string, fallback: boolean): boolean {
  if (value === undefined || value === null) return fallback;
  const parsed = booleanValue(value);
@@ -1726,6 +1859,32 @@ def list_accounts(token):
    data = ensure_success(curl_api("GET", path, bearer=token), "list accounts")
    return extract_items(data)

+def temp_unschedulable_credentials(profile):
+    policy = profile.get("tempUnschedulable") or {}
+    enabled = policy.get("enabled") is True
+    rules = []
+    if enabled:
+        for rule in policy.get("rules") or []:
+            status_code = rule.get("statusCode") if isinstance(rule, dict) else None
+            duration_minutes = rule.get("durationMinutes") if isinstance(rule, dict) else None
+            keywords = rule.get("keywords") if isinstance(rule, dict) else None
+            if not isinstance(status_code, int) or not isinstance(duration_minutes, int) or not isinstance(keywords, list):
+                continue
+            clean_keywords = [str(item) for item in keywords if isinstance(item, str) and item.strip()]
+            if not clean_keywords:
+                continue
+            description = rule.get("description") if isinstance(rule.get("description"), str) else ""
+            rules.append({
+                "error_code": status_code,
+                "keywords": clean_keywords,
+                "duration_minutes": duration_minutes,
+                "description": description,
+            })
+    return {
+        "enabled": enabled and len(rules) > 0,
+        "rules": rules,
+    }
+
 def account_payload(profile, group_id):
    extra = {
        "openai_responses_mode": "force_responses",
@@ -1743,6 +1902,9 @@ def account_payload(profile, group_id):
    upstream_user_agent = profile.get("upstreamUserAgent")
    if upstream_user_agent:
        credentials["user_agent"] = upstream_user_agent
+    temp_unschedulable = temp_unschedulable_credentials(profile)
+    credentials["temp_unschedulable_enabled"] = temp_unschedulable["enabled"]
+    credentials["temp_unschedulable_rules"] = temp_unschedulable["rules"]
    return {
        "name": profile["accountName"],
        "notes": f"UniDesk-managed Codex profile {profile['profile']} from {profile['configFile']} and {profile['authFile']}; secret source={profile['apiKeySource']}; fingerprint={profile['apiKeyFingerprint']}.",
@@ -1788,6 +1950,8 @@ def ensure_accounts(token, profiles, group_id):
            "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,
+            "tempUnschedulableConfigured": bool(payload["credentials"].get("temp_unschedulable_enabled")),
+            "tempUnschedulableRuleCount": len(payload["credentials"].get("temp_unschedulable_rules") or []),
            "upstreamUserAgentConfigured": bool(profile.get("upstreamUserAgent")),
            "valuesPrinted": False,
        })