feat: add Codex WS flag and Xiaomi MiMo usage tracking

Add a global OMNIROUTE_CODEX_WS_ENABLED kill-switch for Codex WebSocket transport, defaulting to enabled if feature flag lookup fails.

Track Xiaomi MiMo monthly token usage from OmniRoute usage history and expose it through the usage fetcher provider list.

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<string, unknown>): 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;
 }
 

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,

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

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;
 

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);

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.

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) ────────────────
   {

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) ──

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<typeof Object>
+  );
+  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();
 

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'", () => {

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<string, { used: number; total: number; remaining?: number; remainingPercentage?: number; resetAt: string | null }>;
+      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");
+  });
+});