codex-multi-auth 2.0.1 → 2.1.0

package/dist/index.js

@@ -37,6 +37,7 @@ import { EntitlementCache, resolveEntitlementAccountKey, } from "./lib/entitleme
 import { LiveAccountSync } from "./lib/live-account-sync.js";
 import { clearCorrelationId, initLogger, logDebug, logError, logInfo, logRequest, logWarn, setCorrelationId, } from "./lib/logger.js";
 import { PreemptiveQuotaScheduler, readQuotaSchedulerSnapshot, } from "./lib/preemptive-quota-scheduler.js";
+import { createRuntimeUsageRecorder, evaluateRuntimePolicy, loadRuntimePolicyState, } from "./lib/policy/runtime-policy.js";
 import { getModelFamily, MODEL_FAMILIES, prewarmCodexInstructions, } from "./lib/prompts/codex.js";
 import { prewarmHostCodexPrompt } from "./lib/prompts/host-codex-prompt.js";
 import { fetchCodexQuotaSnapshot, formatQuotaSnapshotLine, } from "./lib/quota-probe.js";
@@ -543,6 +544,8 @@ export const OpenAIOAuthPlugin = async ({ client }) => {
                          * @returns Response from Codex API
                          */
                         async fetch(input, init) {
+                            let usageRecorder = null;
+                            let usageCompletion = null;
                             try {
                                 if (cachedAccountManager &&
                                     cachedAccountManager !== accountManager) {
@@ -608,6 +611,59 @@ export const OpenAIOAuthPlugin = async ({ client }) => {
                                 const requestTraceId = requestCorrelationId;
                                 runtimeMetrics.lastRequestAt = Date.now();
                                 syncRuntimeObservability(requestTraceId);
+                                let runtimePolicyDecision = null;
+                                try {
+                                    const runtimePolicyState = await loadRuntimePolicyState();
+                                    runtimePolicyDecision = await evaluateRuntimePolicy({
+                                        state: runtimePolicyState,
+                                        accounts: accountManager.getAccountsSnapshot(),
+                                        model: model ?? null,
+                                        capabilityPolicy: capabilityPolicyStore,
+                                    });
+                                }
+                                catch (policyError) {
+                                    logWarn("Runtime policy evaluation failed open.", {
+                                        error: policyError instanceof Error
+                                            ? policyError.message
+                                            : String(policyError),
+                                    });
+                                }
+                                usageRecorder = createRuntimeUsageRecorder({
+                                    source: "plugin-host",
+                                    operation: "responses",
+                                    model: model ?? null,
+                                    projectKey: runtimePolicyDecision?.projectKey ?? null,
+                                    requestId: requestTraceId,
+                                    startedAt: Date.now(),
+                                });
+                                usageCompletion = {
+                                    outcome: "failure",
+                                    statusCode: null,
+                                    errorCode: "plugin_host_request_failed",
+                                };
+                                if (runtimePolicyDecision && !runtimePolicyDecision.allowed) {
+                                    runtimeMetrics.failedRequests += 1;
+                                    runtimeMetrics.lastError = "Runtime policy blocked request";
+                                    syncRuntimeObservability(requestTraceId);
+                                    usageCompletion = {
+                                        outcome: "blocked",
+                                        statusCode: runtimePolicyDecision.statusCode,
+                                        errorCode: runtimePolicyDecision.errorCode,
+                                    };
+                                    return new Response(JSON.stringify({
+                                        error: {
+                                            message: "Runtime policy blocked this local request.",
+                                            code: runtimePolicyDecision.errorCode ??
+                                                "policy_blocked",
+                                            reasons: runtimePolicyDecision.reasons,
+                                        },
+                                    }), {
+                                        status: runtimePolicyDecision.statusCode,
+                                        headers: {
+                                            "content-type": "application/json; charset=utf-8",
+                                        },
+                                    });
+                                }
                                 const abortSignal = requestInit?.signal ?? init?.signal ?? null;
                                 const sleep = createAbortableSleep(abortSignal);
                                 const poolCooldownRemainingMs = getPoolExhaustionCooldownRemaining();
@@ -616,6 +672,11 @@ export const OpenAIOAuthPlugin = async ({ client }) => {
                                     runtimeMetrics.poolExhaustionFastFails += 1;
                                     runtimeMetrics.lastError = "Pool exhaustion cooldown active";
                                     syncRuntimeObservability(requestTraceId);
+                                    usageCompletion = {
+                                        outcome: "failure",
+                                        statusCode: 429,
+                                        errorCode: "pool_exhaustion_cooldown",
+                                    };
                                     return new Response(JSON.stringify({
                                         error: {
                                             message: `The account pool is cooling down after recent rate-limit exhaustion. Try again in ${formatWaitTime(poolCooldownRemainingMs)} or inspect \`codex auth status\`.`,
@@ -633,6 +694,11 @@ export const OpenAIOAuthPlugin = async ({ client }) => {
                                     runtimeMetrics.serverBurstFastFails += 1;
                                     runtimeMetrics.lastError = "Server burst cooldown active";
                                     syncRuntimeObservability(requestTraceId);
+                                    usageCompletion = {
+                                        outcome: "failure",
+                                        statusCode: 503,
+                                        errorCode: "server_burst_cooldown",
+                                    };
                                     return new Response(JSON.stringify({
                                         error: {
                                             message: `Multiple accounts recently failed with upstream server errors. Try again in ${formatWaitTime(serverBurstCooldownRemainingMs)} or inspect \`codex auth report --json\`.`,
@@ -713,12 +779,20 @@ export const OpenAIOAuthPlugin = async ({ client }) => {
                                         capabilityBoostByAccount[candidate.index] =
                                             capabilityPolicyStore.getBoost(accountKey, model ?? modelFamily);
                                     }
+                                    for (const [rawIndex, boost] of Object.entries(runtimePolicyDecision?.scoreBoostByAccount ?? {})) {
+                                        const index = Number(rawIndex);
+                                        if (!Number.isInteger(index))
+                                            continue;
+                                        capabilityBoostByAccount[index] =
+                                            (capabilityBoostByAccount[index] ?? 0) + boost;
+                                    }
                                     accountAttemptLoop: while (attempted.size < Math.max(1, accountCount)) {
                                         let account = null;
                                         if (!usedPreferredSessionAccount &&
                                             typeof preferredSessionAccountIndex === "number") {
                                             usedPreferredSessionAccount = true;
-                                            if (accountManager.isAccountAvailableForFamily(preferredSessionAccountIndex, modelFamily, model)) {
+                                            if (accountManager.isAccountAvailableForFamily(preferredSessionAccountIndex, modelFamily, model) &&
+                                                !runtimePolicyDecision?.blockedAccountIndexes.has(preferredSessionAccountIndex)) {
                                                 account = accountManager.getAccountByIndex(preferredSessionAccountIndex);
                                                 if (account) {
                                                     account.lastUsed = Date.now();
@@ -739,6 +813,9 @@ export const OpenAIOAuthPlugin = async ({ client }) => {
                                             break;
                                         }
                                         attempted.add(account.index);
+                                        if (runtimePolicyDecision?.blockedAccountIndexes.has(account.index)) {
+                                            continue;
+                                        }
                                         // Log account selection for debugging rotation
                                         logDebug(`Using account ${account.index + 1}/${accountCount}: ${account.email ?? "unknown"} for ${modelFamily}`);
                                         let accountAuth = accountManager.toAuthDetails(account);
@@ -1587,6 +1664,16 @@ export const OpenAIOAuthPlugin = async ({ client }) => {
                                             }
                                             runtimeMetrics.successfulRequests++;
                                             runtimeMetrics.lastError = null;
+                                            usageCompletion = {
+                                                outcome: "success",
+                                                statusCode: successResponse.status,
+                                                errorCode: null,
+                                                account: {
+                                                    index: successAccountForResponse.index,
+                                                    accountId: successAccountForResponse.accountId ?? null,
+                                                    email: successAccountForResponse.email ?? null,
+                                                },
+                                            };
                                             if (lastCodexCliActiveSyncIndex !==
                                                 successAccountForResponse.index) {
                                                 void accountManager.syncCodexCliActiveSelectionForIndex(successAccountForResponse.index);
@@ -1657,6 +1744,9 @@ export const OpenAIOAuthPlugin = async ({ client }) => {
                                 }
                             }
                             finally {
+                                if (usageRecorder && usageCompletion) {
+                                    await usageRecorder.record(usageCompletion);
+                                }
                                 syncRuntimeObservability(null);
                                 clearCorrelationId();
                             }