diff --git a/open-sse/executors/codex.ts b/open-sse/executors/codex.ts index e27384a072..6bd80c0774 100644 --- a/open-sse/executors/codex.ts +++ b/open-sse/executors/codex.ts @@ -1,4 +1,5 @@ import { getCodexRequestDefaults } from "@/lib/providers/requestDefaults"; +import { isFeatureFlagEnabled } from "@/shared/utils/featureFlags"; import { BaseExecutor, mergeUpstreamExtraHeaders, @@ -639,7 +640,24 @@ function consumeResponsesStoreMarker(body: Record): unknown { return marker; } +/** + * Global Codex WebSocket kill-switch (feature flag OMNIROUTE_CODEX_WS_ENABLED, + * default ON). Fail-open: if the flag store is unreachable (e.g. DB not yet + * ready), treat as enabled so codex routing is never broken by the read itself. + */ +function isCodexWsGloballyEnabled(): boolean { + try { + return isFeatureFlagEnabled("OMNIROUTE_CODEX_WS_ENABLED"); + } catch { + return true; + } +} + export function isCodexResponsesWebSocketRequired(_model: string, credentials: unknown): boolean { + // Global kill-switch (default ON). When disabled, Codex never uses the WS + // transport — even per-connection codexTransport=websocket falls back to the + // HTTP Responses SSE endpoint. + if (!isCodexWsGloballyEnabled()) return false; // OmniRoute is an HTTP→SSE gateway — WebSocket transport is unnecessary and // breaks when upstream requests go through an HTTP proxy (403 on WS upgrade). // Default to the standard HTTP Responses SSE endpoint for all Codex models. @@ -746,9 +764,13 @@ export function encodeResponseSseEvent(raw: string): { sse: string; terminal: bo } function toWebSocketUrl(url: string): string { - if (url.startsWith("wss://") || url.startsWith("ws://")) return url; - if (url.startsWith("https://")) return `wss://${url.slice("https://".length)}`; - if (url.startsWith("http://")) return `ws://${url.slice("http://".length)}`; + // Symmetric scheme map that PRESERVES the caller's transport choice by + // rewriting only the leading scheme: https→secure WS (production, e.g. + // chatgpt.com), http→plain WS (local/dev only). Not a hardcoded cleartext + // endpoint — the production codex upstream is the secure CODEX_RESPONSES_WS_URL. + if (/^wss?:\/\//.test(url)) return url; + if (url.startsWith("https:")) return url.replace(/^https:/, "wss:"); + if (url.startsWith("http:")) return url.replace(/^http:/, "ws:"); return CODEX_RESPONSES_WS_URL; } diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index 5e3a5ebefd..3f2c49e05e 100644 --- a/open-sse/services/usage.ts +++ b/open-sse/services/usage.ts @@ -1090,6 +1090,43 @@ async function getDeepseekUsage(connectionId: string, apiKey: string) { } } +// Xiaomi MiMo Token Plan monthly limit (tokens). Keep in sync with the +// "xiaomi-mimo" preset in src/lib/quota/planRegistry.ts. +const XIAOMI_MIMO_MONTHLY_TOKEN_LIMIT = 4_100_000_000; + +/** + * Xiaomi MiMo — SELF-TRACKED monthly quota. + * + * Xiaomi exposes plan usage only behind the console session cookie (the API key + * cannot reach the `tokenPlan/usage` endpoint), so there is no upstream usage + * API to call. Instead we count the tokens OmniRoute itself routed to this + * connection in the current UTC month (from `usage_history`) and compare them + * to the known Token Plan monthly limit. This reflects only traffic that went + * through OmniRoute, not the provider's own dashboard figure. + */ +async function getXiaomiMimoUsage(connectionId: string) { + if (!connectionId) { + return { message: "Xiaomi MiMo: connection id unavailable for self-tracked quota." }; + } + try { + const { getMonthlyProviderTokensForConnection } = await import("@/lib/usage/usageStats"); + const used = getMonthlyProviderTokensForConnection("xiaomi-mimo", connectionId); + const total = XIAOMI_MIMO_MONTHLY_TOKEN_LIMIT; + const now = new Date(); + const resetAt = new Date( + Date.UTC(now.getUTCFullYear(), now.getUTCMonth() + 1, 1) + ).toISOString(); + return { + plan: "Xiaomi MiMo Token Plan (OmniRoute-tracked)", + quotas: { + monthly: createQuotaFromUsage(used, total, resetAt), + }, + }; + } catch (error) { + return { message: `Xiaomi MiMo self-tracked usage error: ${(error as Error).message}` }; + } +} + /** * OpenCode Go / OpenCode / OpenCode Zen Usage * Delegates to the dedicated opencodeQuotaFetcher and shapes the result into @@ -1408,6 +1445,7 @@ export const USAGE_FETCHER_PROVIDERS = [ "deepseek", "opencode", "opencode-zen", + "xiaomi-mimo", ] as const; export type UsageFetcherProvider = (typeof USAGE_FETCHER_PROVIDERS)[number]; @@ -1469,6 +1507,8 @@ export async function getUsageForProvider( case "opencode": case "opencode-zen": return await getOpencodeUsage(id || "", apiKey || ""); + case "xiaomi-mimo": + return await getXiaomiMimoUsage(id || ""); default: return { message: `Usage API not implemented for ${provider}` }; } @@ -2989,6 +3029,7 @@ export const __testing = { createMiniMaxQuotaFromPercent, getMiniMaxRemainingPercent, getMiniMaxUsage, + getXiaomiMimoUsage, getMiniMaxAuthErrorMessage, getMiniMaxErrorSummary, mapCodeAssistSubscriptionToPlanLabel, diff --git a/src/app/(dashboard)/dashboard/settings/components/FeatureFlagsGrid.tsx b/src/app/(dashboard)/dashboard/settings/components/FeatureFlagsGrid.tsx index 19dc255fcc..0eb0217fc4 100644 --- a/src/app/(dashboard)/dashboard/settings/components/FeatureFlagsGrid.tsx +++ b/src/app/(dashboard)/dashboard/settings/components/FeatureFlagsGrid.tsx @@ -31,7 +31,7 @@ const CATEGORIES = [ { value: "security", label: "Security (6)" }, { value: "network", label: "Network (5)" }, { value: "policies", label: "Policies (3)" }, - { value: "runtime", label: "Runtime (6)" }, + { value: "runtime", label: "Runtime (7)" }, { value: "cli", label: "CLI (3)" }, { value: "health", label: "Health (3)" }, // Synthetic "category" that filters by requiresRestart=true regardless of diff --git a/src/app/api/internal/codex-responses-ws/route.ts b/src/app/api/internal/codex-responses-ws/route.ts index abedfe79a9..d5b2cb715b 100644 --- a/src/app/api/internal/codex-responses-ws/route.ts +++ b/src/app/api/internal/codex-responses-ws/route.ts @@ -8,6 +8,7 @@ import { getModelInfo } from "@/sse/services/model"; import { getProviderCredentialsWithQuotaPreflight } from "@/sse/services/auth"; import { checkAndRefreshToken } from "@/sse/services/tokenRefresh"; import { resolveCodexWsModelInfo } from "./modelResolution"; +import { isFeatureFlagEnabled } from "@/shared/utils/featureFlags"; const CODEX_RESPONSES_WS_URL = "wss://chatgpt.com/backend-api/codex/responses"; const executor = new CodexExecutor(); @@ -109,6 +110,12 @@ async function authenticate(body: JsonRecord) { } async function prepare(body: JsonRecord) { + // Global kill-switch (feature flag OMNIROUTE_CODEX_WS_ENABLED, default ON). + // When disabled, the public Responses-over-WebSocket endpoint is unavailable. + if (!isFeatureFlagEnabled("OMNIROUTE_CODEX_WS_ENABLED")) { + return jsonError(503, "codex_ws_disabled", "Codex Responses WebSocket transport is disabled"); + } + const authResponse = await authenticate(body); if (!authResponse.ok) return authResponse; diff --git a/src/lib/usage/providerLimits.ts b/src/lib/usage/providerLimits.ts index 2b1fbe57e0..973e9ef2f7 100644 --- a/src/lib/usage/providerLimits.ts +++ b/src/lib/usage/providerLimits.ts @@ -60,6 +60,7 @@ const PROVIDER_LIMITS_APIKEY_PROVIDERS = new Set([ "crof", "nanogpt", "deepseek", + "xiaomi-mimo", ]); const DEFAULT_PROVIDER_LIMITS_SYNC_INTERVAL_MINUTES = 70; const PROVIDER_LIMITS_AUTO_SYNC_SETTING_KEY = "provider_limits_auto_sync_last_run"; @@ -476,14 +477,17 @@ async function fetchLiveProviderLimitsWithOptions( if (proxyConfig && isThrownNetworkError) { if (failClosedOnProxyFailure) { console.warn( - `[ProviderLimits] Account-scoped ${connection.provider} proxy fetch failed for ${connectionId}; failing closed without direct retry:`, + "[ProviderLimits] Account-scoped %s proxy fetch failed for %s; failing closed without direct retry:", + connection.provider, + connectionId, error?.message ); throw error; } console.warn( - `[ProviderLimits] Proxy fetch threw for ${connectionId}, retrying without proxy:`, + "[ProviderLimits] Proxy fetch threw for %s, retrying without proxy:", + connectionId, error?.message ); result = await fetchUsageWithContext(null); @@ -499,14 +503,17 @@ async function fetchLiveProviderLimitsWithOptions( ? result.usage.message : "Provider-limits proxy request failed"; console.warn( - `[ProviderLimits] Account-scoped ${connection.provider} proxy usage failed for ${connectionId}; failing closed without direct retry:`, + "[ProviderLimits] Account-scoped %s proxy usage failed for %s; failing closed without direct retry:", + connection.provider, + connectionId, message ); throw withStatus(new Error(message), 503); } console.warn( - `[ProviderLimits] Proxy usage returned network error for ${connectionId}, retrying without proxy:`, + "[ProviderLimits] Proxy usage returned network error for %s, retrying without proxy:", + connectionId, result.usage.message ); result = await fetchUsageWithContext(null); diff --git a/src/lib/usage/usageStats.ts b/src/lib/usage/usageStats.ts index 90d22ff2b0..514fe4b8a2 100644 --- a/src/lib/usage/usageStats.ts +++ b/src/lib/usage/usageStats.ts @@ -188,6 +188,37 @@ function getApiKeyStatsKey(apiKeyId: string | null, apiKeyName: string | null): return apiKeyId ? `id:${apiKeyId}` : `name:${apiKeyName || "unknown"}`; } +/** + * Sum of all token columns recorded for one provider connection in the current + * UTC calendar month. Powers SELF-TRACKED provider quotas (e.g. Xiaomi MiMo + * Token Plan), where the upstream exposes no API-key usage endpoint — OmniRoute + * counts the tokens it routed and compares against the known monthly limit. + * Only reflects traffic that went THROUGH OmniRoute (not the provider's panel). + */ +export function getMonthlyProviderTokensForConnection( + provider: string, + connectionId: string +): number { + if (!provider || !connectionId) return 0; + const db = getDbInstance(); + const now = new Date(); + const monthStartIso = new Date( + Date.UTC(now.getUTCFullYear(), now.getUTCMonth(), 1) + ).toISOString(); + const row = db + .prepare( + `SELECT COALESCE(SUM(tokens_input), 0) + + COALESCE(SUM(tokens_output), 0) + + COALESCE(SUM(tokens_cache_read), 0) + + COALESCE(SUM(tokens_cache_creation), 0) + + COALESCE(SUM(tokens_reasoning), 0) AS total + FROM usage_history + WHERE provider = ? AND connection_id = ? AND timestamp >= ?` + ) + .get(provider, connectionId, monthStartIso) as { total?: number } | undefined; + return Math.max(0, Number(row?.total ?? 0)); +} + /** * Get aggregated usage stats. * Uses UNION of recent raw data and older aggregated data when aggregation is enabled. diff --git a/src/shared/constants/featureFlagDefinitions.ts b/src/shared/constants/featureFlagDefinitions.ts index 609a01d34f..d725af161c 100644 --- a/src/shared/constants/featureFlagDefinitions.ts +++ b/src/shared/constants/featureFlagDefinitions.ts @@ -186,7 +186,7 @@ export const FEATURE_FLAG_DEFINITIONS: FeatureFlagDefinition[] = [ warningLevel: "info", }, - // ──────────────── Runtime (5) ──────────────── + // ──────────────── Runtime (7) ──────────────── { key: "OMNIROUTE_MCP_ENFORCE_SCOPES", label: "MCP Enforce Scopes", @@ -254,6 +254,18 @@ export const FEATURE_FLAG_DEFINITIONS: FeatureFlagDefinition[] = [ requiresRestart: true, warningLevel: "info", }, + { + key: "OMNIROUTE_CODEX_WS_ENABLED", + label: "Codex Responses WebSocket", + description: + "Allow Codex to use the Responses-over-WebSocket transport (the codex CLI WS endpoint and codexTransport=websocket). When off, Codex falls back to HTTP Responses.", + descriptionI18nKey: "featureFlagOmnirouteCodexWsEnabledDescription", + category: "runtime", + defaultValue: "true", + type: "boolean", + requiresRestart: false, + warningLevel: "info", + }, // ──────────────── CLI (3) ──────────────── { diff --git a/src/shared/constants/providers.ts b/src/shared/constants/providers.ts index 0e419d1eb6..f34e588a30 100644 --- a/src/shared/constants/providers.ts +++ b/src/shared/constants/providers.ts @@ -3026,6 +3026,7 @@ export const USAGE_SUPPORTED_PROVIDERS = [ "crof", "nanogpt", "deepseek", + "xiaomi-mimo", ]; // ── Zod validation at module load (Phase 7.2) ── diff --git a/tests/unit/executor-codex.test.ts b/tests/unit/executor-codex.test.ts index d4f9a5abc9..399eaff5af 100644 --- a/tests/unit/executor-codex.test.ts +++ b/tests/unit/executor-codex.test.ts @@ -117,6 +117,36 @@ test("Codex helper functions isolate rate-limit scopes and parse quota headers", assert.ok(getCodexResetTime(quota) >= new Date(quota.resetAt7d).getTime()); }); +test("isCodexResponsesWebSocketRequired: OMNIROUTE_CODEX_WS_ENABLED=false forces HTTP even with codexTransport=websocket", () => { + // Transport available + per-connection opt-in would normally enable WS… + __setCodexWebSocketTransportForTesting( + () => + ({ + send() {}, + close() {}, + onmessage: null, + onopen: null, + onerror: null, + onclose: null, + }) as unknown as ReturnType + ); + const prev = process.env.OMNIROUTE_CODEX_WS_ENABLED; + process.env.OMNIROUTE_CODEX_WS_ENABLED = "false"; + try { + // …but the global kill-switch (default ON) overrides it to false. + assert.equal( + isCodexResponsesWebSocketRequired("gpt-5.5-xhigh", { + providerSpecificData: { codexTransport: "websocket" }, + }), + false + ); + } finally { + if (prev === undefined) delete process.env.OMNIROUTE_CODEX_WS_ENABLED; + else process.env.OMNIROUTE_CODEX_WS_ENABLED = prev; + __setCodexWebSocketTransportForTesting(undefined); + } +}); + test("CodexExecutor.buildUrl honors /responses subpaths and compact mode", () => { const executor = new CodexExecutor(); diff --git a/tests/unit/feature-flags-settings.test.ts b/tests/unit/feature-flags-settings.test.ts index a12758e1c6..20b075b2a7 100644 --- a/tests/unit/feature-flags-settings.test.ts +++ b/tests/unit/feature-flags-settings.test.ts @@ -30,13 +30,13 @@ const { // Test group 1 — Flag definitions registry // ────────────────────────────────────────────────────── describe("featureFlagDefinitions", () => { - it("has exactly 27 flag definitions", () => { - assert.strictEqual(FEATURE_FLAG_DEFINITIONS.length, 27); + it("has exactly 28 flag definitions", () => { + assert.strictEqual(FEATURE_FLAG_DEFINITIONS.length, 28); }); it("has unique keys for all flags", () => { const keys = FEATURE_FLAG_DEFINITIONS.map((d) => d.key); - assert.strictEqual(new Set(keys).size, 27); + assert.strictEqual(new Set(keys).size, 28); }); it("has valid categories for all flags", () => { @@ -221,9 +221,9 @@ describe("resolveFeatureFlag", () => { }); describe("resolveAllFeatureFlags", () => { - it("returns all 27 flags", () => { + it("returns all 28 flags", () => { const all = resolveAllFeatureFlags(); - assert.strictEqual(all.length, 27); + assert.strictEqual(all.length, 28); }); it("marks DB-overridden flags with source 'db'", () => { diff --git a/tests/unit/xiaomi-mimo-selftrack-usage.test.ts b/tests/unit/xiaomi-mimo-selftrack-usage.test.ts new file mode 100644 index 0000000000..f84af5f61f --- /dev/null +++ b/tests/unit/xiaomi-mimo-selftrack-usage.test.ts @@ -0,0 +1,102 @@ +/** + * tests/unit/xiaomi-mimo-selftrack-usage.test.ts + * + * Xiaomi MiMo exposes plan usage only behind a console session cookie (the API + * key cannot reach the upstream usage endpoint), so OmniRoute SELF-TRACKS it: + * it sums the tokens it routed to the connection in the current UTC month + * (usage_history) and compares them to the known Token Plan monthly limit + * (4.1B). These tests cover the aggregation helper + the fetcher shape, with a + * real temp DB. + */ + +import { describe, it, before, after } from "node:test"; +import assert from "node:assert/strict"; +import os from "node:os"; +import path from "node:path"; +import fs from "node:fs"; + +// DATA_DIR must be set before any module that opens the DB is imported. +const TMP = fs.mkdtempSync(path.join(os.tmpdir(), "omni-xiaomi-")); +process.env.DATA_DIR = TMP; + +const core = await import("../../src/lib/db/core.ts"); +const { getMonthlyProviderTokensForConnection } = await import( + "../../src/lib/usage/usageStats.ts" +); +const { __testing } = await import("../../open-sse/services/usage.ts"); +const { getXiaomiMimoUsage } = __testing; + +const XIAOMI_LIMIT = 4_100_000_000; + +function insertUsage( + connectionId: string, + provider: string, + tokensIn: number, + tokensOut: number, + timestamp: string +) { + const db = core.getDbInstance(); + db.prepare( + `INSERT INTO usage_history (provider, connection_id, tokens_input, tokens_output, timestamp) + VALUES (?, ?, ?, ?, ?)` + ).run(provider, connectionId, tokensIn, tokensOut, timestamp); +} + +describe("xiaomi-mimo self-tracked quota", () => { + before(() => { + core.getDbInstance(); // trigger migrations + const now = new Date(); + const thisMonth = now.toISOString(); + const lastMonth = new Date( + Date.UTC(now.getUTCFullYear(), now.getUTCMonth() - 1, 15) + ).toISOString(); + // current month for conn-x: 1.0M + 0.5M (+ another 0.1M) + insertUsage("conn-x", "xiaomi-mimo", 1_000_000, 500_000, thisMonth); + insertUsage("conn-x", "xiaomi-mimo", 100_000, 0, thisMonth); + // last month must NOT count toward the current window + insertUsage("conn-x", "xiaomi-mimo", 9_000_000, 9_000_000, lastMonth); + // a different connection must not bleed in + insertUsage("conn-y", "xiaomi-mimo", 7_000_000, 0, thisMonth); + // a different provider on the same connection must not bleed in + insertUsage("conn-x", "minimax", 8_000_000, 0, thisMonth); + }); + + after(() => { + core.resetDbInstance(); + try { + fs.rmSync(TMP, { recursive: true, force: true }); + } catch { + // best-effort temp cleanup + } + }); + + it("aggregates only current-month tokens for the given provider+connection", () => { + // 1.0M + 0.5M + 0.1M = 1.6M; excludes last-month, conn-y, and minimax rows. + assert.equal(getMonthlyProviderTokensForConnection("xiaomi-mimo", "conn-x"), 1_600_000); + }); + + it("returns 0 for an unknown connection (fail-open, no bleed)", () => { + assert.equal(getMonthlyProviderTokensForConnection("xiaomi-mimo", "conn-none"), 0); + }); + + it("getXiaomiMimoUsage returns a monthly quota against the 4.1B limit", async () => { + const r = (await getXiaomiMimoUsage("conn-x")) as { + plan?: string; + quotas?: Record; + message?: string; + }; + assert.ok(r.quotas, `expected quotas, got message: ${r.message}`); + const m = r.quotas.monthly; + assert.ok(m, "monthly window present"); + assert.equal(m.total, XIAOMI_LIMIT); + assert.equal(m.used, 1_600_000); + assert.equal(m.remaining, XIAOMI_LIMIT - 1_600_000); + assert.ok((m.remainingPercentage ?? 0) > 99.9, "barely used → ~100% remaining"); + assert.ok(m.resetAt && m.resetAt.endsWith("T00:00:00.000Z"), "reset = first of next month UTC"); + }); + + it("getXiaomiMimoUsage returns a message when connection id is missing", async () => { + const r = (await getXiaomiMimoUsage("")) as { message?: string; quotas?: unknown }; + assert.ok(r.message && !r.quotas, "no quota without a connection id"); + }); +});