feat: normalize quota and combos API responses with shared contracts

Introduce `normalizeQuotaResponse` and `normalizeCombosResponse` helpers to handle varying API response shapes (array vs wrapped object) consistently across MCP server and advanced tools. Add optional `meta` field to checkQuotaOutput schema and update sourceEndpoints to reflect current API routes.
2026-07-26 09:52:11 +03:00 · 2026-03-04 08:18:09 -03:00
parent fe9d9a5a5c
commit 5ecef5c90c
8 changed files with 289 additions and 29 deletions
--- a/open-sse/mcp-server/schemas/tools.ts
+++ b/open-sse/mcp-server/schemas/tools.ts
@@ -225,6 +225,16 @@ export const checkQuotaOutput = z.object({
      tokenStatus: z.enum(["valid", "expiring", "expired", "refreshing"]),
    })
  ),
+  meta: z
+    .object({
+      generatedAt: z.string(),
+      filters: z.object({
+        provider: z.string().nullable(),
+        connectionId: z.string().nullable(),
+      }),
+      totalProviders: z.number(),
+    })
+    .optional(),
 });

 export const checkQuotaTool: McpToolDefinition<typeof checkQuotaInput, typeof checkQuotaOutput> = {
@@ -236,7 +246,7 @@ export const checkQuotaTool: McpToolDefinition<typeof checkQuotaInput, typeof ch
  scopes: ["read:quota"],
  auditLevel: "basic",
  phase: 1,
-  sourceEndpoints: ["/api/usage/[connectionId]", "/api/token-health"],
+  sourceEndpoints: ["/api/usage/quota", "/api/token-health", "/api/rate-limits"],
 };

 // --- Tool 6: omniroute_route_request ---
--- a/open-sse/mcp-server/server.ts
+++ b/open-sse/mcp-server/server.ts
@@ -37,6 +37,7 @@ import {
  handleExplainRoute,
  handleGetSessionSnapshot,
 } from "./tools/advancedTools.ts";
+import { normalizeQuotaResponse } from "@/shared/contracts/quota";

 // ============ Configuration ============

@@ -105,7 +106,12 @@ async function handleGetHealth() {
 async function handleListCombos(args: { includeMetrics?: boolean }) {
  const start = Date.now();
  try {
-    const combos = (await omniRouteFetch("/api/combos")) as any;
+    const combosRaw = (await omniRouteFetch("/api/combos")) as any;
+    const combos = Array.isArray(combosRaw?.combos)
+      ? combosRaw.combos
+      : Array.isArray(combosRaw)
+        ? combosRaw
+        : [];
    let metrics: Record<string, unknown> = {};
    if (args.includeMetrics) {
      metrics = (await omniRouteFetch("/api/combos/metrics").catch(() => ({}))) as Record<
@@ -174,10 +180,10 @@ async function handleCheckQuota(args: { provider?: string; connectionId?: string
    if (args.connectionId) path += `?connectionId=${encodeURIComponent(args.connectionId)}`;
    else if (args.provider) path += `?provider=${encodeURIComponent(args.provider)}`;

-    const raw = (await omniRouteFetch(path)) as any;
-    const result = {
-      providers: Array.isArray(raw?.providers) ? raw.providers : Array.isArray(raw) ? raw : [],
-    };
+    const result = normalizeQuotaResponse(await omniRouteFetch(path), {
+      provider: args.provider || null,
+      connectionId: args.connectionId || null,
+    });

    await logToolCall("omniroute_check_quota", args, result, Date.now() - start, true);
    return { content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }] };
--- a/open-sse/mcp-server/tools/advancedTools.ts
+++ b/open-sse/mcp-server/tools/advancedTools.ts
@@ -14,6 +14,7 @@
 */

 import { logToolCall } from "../audit.ts";
+import { normalizeQuotaResponse } from "@/shared/contracts/quota";

 const OMNIROUTE_BASE_URL = process.env.OMNIROUTE_BASE_URL || "http://localhost:20128";
 const OMNIROUTE_API_KEY = process.env.OMNIROUTE_API_KEY || "";
@@ -33,6 +34,14 @@ async function apiFetch(path: string, options: RequestInit = {}): Promise<unknow
  return response.json();
 }

+function normalizeCombosResponse(raw: unknown): any[] {
+  if (Array.isArray(raw)) return raw;
+  if (raw && typeof raw === "object" && Array.isArray((raw as any).combos)) {
+    return (raw as any).combos;
+  }
+  return [];
+}
+
 // ============ In-Memory State ============

 interface BudgetGuardState {
@@ -77,9 +86,12 @@ export async function handleSimulateRoute(args: {
      apiFetch("/api/usage/quota"),
    ]);

-    const combos = combosRaw.status === "fulfilled" ? (combosRaw.value as any[]) : [];
+    const combos = combosRaw.status === "fulfilled" ? normalizeCombosResponse(combosRaw.value) : [];
    const health = healthRaw.status === "fulfilled" ? (healthRaw.value as any) : {};
-    const quota = quotaRaw.status === "fulfilled" ? (quotaRaw.value as any) : {};
+    const quota =
+      quotaRaw.status === "fulfilled"
+        ? normalizeQuotaResponse(quotaRaw.value)
+        : normalizeQuotaResponse({});

    // Find target combo
    const targetCombo = args.combo
@@ -97,7 +109,7 @@ export async function handleSimulateRoute(args: {

    const models = targetCombo.models || targetCombo.data?.models || [];
    const breakers = health?.circuitBreakers || [];
-    const providers = quota?.providers || (Array.isArray(quota) ? quota : []);
+    const providers = quota.providers;

    // Simulate path
    const simulatedPath = models.map((m: any, idx: number) => {
@@ -233,7 +245,7 @@ export async function handleTestCombo(args: { comboId: string; testPrompt: strin
  const start = Date.now();
  try {
    // Get combo details
-    const combos = (await apiFetch("/api/combos")) as any[];
+    const combos = normalizeCombosResponse(await apiFetch("/api/combos"));
    const combo = combos.find((c: any) => c.id === args.comboId || c.name === args.comboId);
    if (!combo) {
      return {
@@ -340,13 +352,14 @@ export async function handleGetProviderMetrics(args: { provider: string }) {
    ]);

    const health = healthRaw.status === "fulfilled" ? (healthRaw.value as any) : {};
-    const quota = quotaRaw.status === "fulfilled" ? (quotaRaw.value as any) : {};
+    const quota =
+      quotaRaw.status === "fulfilled"
+        ? normalizeQuotaResponse(quotaRaw.value, { provider: args.provider })
+        : normalizeQuotaResponse({});
    const analytics = analyticsRaw.status === "fulfilled" ? (analyticsRaw.value as any) : {};

    const cb = (health.circuitBreakers || []).find((b: any) => b.provider === args.provider);
-    const providerQuota = Array.isArray(quota?.providers)
-      ? quota.providers.find((p: any) => p.provider === args.provider)
-      : null;
+    const providerQuota = quota.providers.find((p) => p.provider === args.provider) || null;

    const result = {
      provider: args.provider,
@@ -385,7 +398,7 @@ export async function handleBestComboForTask(args: {
  const start = Date.now();
  try {
    const fitness = TASK_FITNESS[args.taskType] || TASK_FITNESS.coding;
-    const combos = (await apiFetch("/api/combos")) as any[];
+    const combos = normalizeCombosResponse(await apiFetch("/api/combos"));
    const enabledCombos = combos.filter((c: any) => c.enabled !== false);

    if (enabledCombos.length === 0) {
--- a/package.json
+++ b/package.json
@@ -62,7 +62,7 @@
    "test:vitest": "vitest run open-sse/mcp-server/__tests__/*.test.ts open-sse/services/autoCombo/__tests__/*.test.ts vscode-extension/src/__tests__/*.test.ts",
    "test:ecosystem": "node scripts/run-ecosystem-tests.mjs",
    "test:coverage": "npx c8 --exclude=open-sse --check-coverage --lines 50 --functions 50 --branches 50 node --import tsx/esm --test tests/unit/*.test.mjs",
-    "test:all": "npm run test:unit && npm run test:vitest && npm run test:e2e",
+    "test:all": "npm run test:unit && npm run test:vitest && npm run test:ecosystem && npm run test:e2e",
    "check": "npm run lint && npm run test",
    "prepublishOnly": "npm run build:cli",
    "postinstall": "node scripts/postinstall.mjs",
--- a/src/lib/a2a/skills/quotaManagement.ts
+++ b/src/lib/a2a/skills/quotaManagement.ts
@@ -6,6 +6,7 @@
 */

 import type { A2ATask, TaskArtifact } from "../taskManager";
+import { normalizeQuotaResponse } from "@/shared/contracts/quota";

 const OMNIROUTE_BASE_URL = process.env.OMNIROUTE_BASE_URL || "http://localhost:20128";
 const OMNIROUTE_API_KEY = process.env.OMNIROUTE_API_KEY || "";
@@ -34,20 +35,42 @@ export async function executeQuotaManagement(task: A2ATask): Promise<QuotaManage
    quotaFetch("/api/combos"),
  ]);

-  const quota = quotaRaw.status === "fulfilled" ? (quotaRaw.value as any) : {};
-  const combos = combosRaw.status === "fulfilled" ? (combosRaw.value as any[]) : [];
-  const providers: any[] = quota?.providers || (Array.isArray(quota) ? quota : []);
+  const quota =
+    quotaRaw.status === "fulfilled"
+      ? normalizeQuotaResponse(quotaRaw.value)
+      : normalizeQuotaResponse({});
+  const combos =
+    combosRaw.status === "fulfilled"
+      ? Array.isArray((combosRaw.value as any)?.combos)
+        ? (combosRaw.value as any).combos
+        : Array.isArray(combosRaw.value)
+          ? (combosRaw.value as any[])
+          : []
+      : [];
+  const providers = quota.providers;
+
+  const availableQuota = (provider: { quotaTotal: number | null; quotaUsed: number }) => {
+    if (provider.quotaTotal === null) return Number.POSITIVE_INFINITY;
+    return provider.quotaTotal - provider.quotaUsed;
+  };

  // Query classification
  if (query.includes("ranking") || query.includes("most quota") || query.includes("best")) {
-    const sorted = [...providers].sort(
-      (a, b) => b.quotaTotal - b.quotaUsed - (a.quotaTotal - a.quotaUsed)
-    );
+    const sorted = [...providers].sort((a, b) => availableQuota(b) - availableQuota(a));
    return {
      artifacts: [
        {
          type: "text",
-          content: `**Quota Ranking (most available first):**\n${sorted.map((p, i) => `${i + 1}. **${p.provider}** — ${(p.quotaTotal - p.quotaUsed).toLocaleString()} remaining (${Math.round(((p.quotaTotal - p.quotaUsed) / (p.quotaTotal || 1)) * 100)}%)`).join("\n")}`,
+          content: `**Quota Ranking (most available first):**\n${sorted
+            .map((p, i) => {
+              const remaining = availableQuota(p);
+              const remainingLabel =
+                remaining === Number.POSITIVE_INFINITY ? "unlimited" : remaining.toLocaleString();
+              const percentLabel =
+                p.quotaTotal === null ? "n/a" : `${Math.round(p.percentRemaining)}%`;
+              return `${i + 1}. **${p.provider}** — ${remainingLabel} remaining (${percentLabel})`;
+            })
+            .join("\n")}`,
        },
      ],
      metadata: { queryType: "ranking", providers: sorted.map((p) => p.provider) },
@@ -76,13 +99,20 @@ export async function executeQuotaManagement(task: A2ATask): Promise<QuotaManage
  // Default: general quota summary
  const totalUsed = providers.reduce((sum, p) => sum + (p.quotaUsed || 0), 0);
  const totalAvailable = providers.reduce((sum, p) => sum + (p.quotaTotal || 0), 0);
-  const warnings = providers.filter((p) => p.quotaTotal && p.quotaUsed / p.quotaTotal > 0.9);
+  const warnings = providers.filter((p) => p.percentRemaining <= 10);

  return {
    artifacts: [
      {
        type: "text",
-        content: `**Quota Summary (${providers.length} providers):**\n- Total used: ${totalUsed.toLocaleString()} / ${totalAvailable.toLocaleString()}\n${providers.map((p) => `- **${p.provider}:** ${p.quotaUsed?.toLocaleString() || 0} / ${p.quotaTotal?.toLocaleString() || "∞"} (${p.tokenStatus || "ok"})`).join("\n")}${warnings.length > 0 ? `\n\n⚠️ **Warning:** ${warnings.map((w) => w.provider).join(", ")} above 90% usage` : ""}`,
+        content: `**Quota Summary (${providers.length} providers):**\n- Total used: ${totalUsed.toLocaleString()} / ${totalAvailable.toLocaleString()}\n${providers
+          .map(
+            (p) =>
+              `- **${p.provider}:** ${p.quotaUsed.toLocaleString()} / ${p.quotaTotal?.toLocaleString() || "∞"} (${p.tokenStatus})`
+          )
+          .join(
+            "\n"
+          )}${warnings.length > 0 ? `\n\n⚠️ **Warning:** ${warnings.map((w) => w.provider).join(", ")} at or below 10% remaining` : ""}`,
      },
    ],
    metadata: { queryType: "summary", providerCount: providers.length, warnings: warnings.length },
--- a/src/shared/contracts/quota.ts
+++ b/src/shared/contracts/quota.ts
@@ -0,0 +1,138 @@
+export type QuotaTokenStatus = "valid" | "expiring" | "expired" | "refreshing";
+
+export interface QuotaProviderEntry {
+  name: string;
+  provider: string;
+  connectionId: string;
+  quotaUsed: number;
+  quotaTotal: number | null;
+  percentRemaining: number;
+  resetAt: string | null;
+  tokenStatus: QuotaTokenStatus;
+}
+
+export interface QuotaResponseMeta {
+  generatedAt: string;
+  filters: {
+    provider: string | null;
+    connectionId: string | null;
+  };
+  totalProviders: number;
+}
+
+export interface QuotaResponse {
+  providers: QuotaProviderEntry[];
+  meta: QuotaResponseMeta;
+}
+
+const TOKEN_STATUS_VALUES: QuotaTokenStatus[] = ["valid", "expiring", "expired", "refreshing"];
+
+function toNumber(value: unknown): number | null {
+  if (typeof value === "number" && Number.isFinite(value)) return value;
+  if (typeof value === "string" && value.trim() !== "") {
+    const parsed = Number(value);
+    if (Number.isFinite(parsed)) return parsed;
+  }
+  return null;
+}
+
+function clamp(value: number, min: number, max: number): number {
+  return Math.min(max, Math.max(min, value));
+}
+
+function normalizeTokenStatus(value: unknown): QuotaTokenStatus {
+  if (typeof value === "string" && TOKEN_STATUS_VALUES.includes(value as QuotaTokenStatus)) {
+    return value as QuotaTokenStatus;
+  }
+  return "valid";
+}
+
+export function sanitizeQuotaProvider(input: unknown): QuotaProviderEntry {
+  const source = input && typeof input === "object" ? (input as Record<string, unknown>) : {};
+  const provider = typeof source.provider === "string" ? source.provider : "unknown";
+  const name = typeof source.name === "string" && source.name.trim() ? source.name : provider;
+  const connectionId =
+    typeof source.connectionId === "string" && source.connectionId.trim()
+      ? source.connectionId
+      : "unknown";
+
+  const quotaTotalRaw = toNumber(source.quotaTotal);
+  const quotaTotal = quotaTotalRaw !== null && quotaTotalRaw >= 0 ? quotaTotalRaw : null;
+
+  const quotaUsedRaw = toNumber(source.quotaUsed) ?? 0;
+  const quotaUsed =
+    quotaTotal !== null ? clamp(quotaUsedRaw, 0, quotaTotal) : Math.max(0, quotaUsedRaw);
+
+  let percentRemainingRaw = toNumber(source.percentRemaining);
+  if (percentRemainingRaw === null) {
+    if (quotaTotal && quotaTotal > 0) {
+      percentRemainingRaw = ((quotaTotal - quotaUsed) / quotaTotal) * 100;
+    } else {
+      percentRemainingRaw = 100;
+    }
+  }
+  const percentRemaining = clamp(percentRemainingRaw, 0, 100);
+
+  const resetAt =
+    typeof source.resetAt === "string" && source.resetAt.trim() ? source.resetAt : null;
+
+  return {
+    name,
+    provider,
+    connectionId,
+    quotaUsed,
+    quotaTotal,
+    percentRemaining,
+    resetAt,
+    tokenStatus: normalizeTokenStatus(source.tokenStatus),
+  };
+}
+
+export function normalizeQuotaResponse(
+  raw: unknown,
+  filters: { provider?: string | null; connectionId?: string | null } = {}
+): QuotaResponse {
+  const source = raw && typeof raw === "object" ? (raw as Record<string, unknown>) : {};
+  const providersRaw = Array.isArray(source.providers)
+    ? source.providers
+    : Array.isArray(raw)
+      ? raw
+      : [];
+
+  const providers = providersRaw.map((entry) => sanitizeQuotaProvider(entry));
+
+  const sourceMeta =
+    source.meta && typeof source.meta === "object" ? (source.meta as Record<string, unknown>) : {};
+  const sourceFilters =
+    sourceMeta.filters && typeof sourceMeta.filters === "object"
+      ? (sourceMeta.filters as Record<string, unknown>)
+      : {};
+
+  const providerFilter =
+    filters.provider ??
+    (typeof sourceFilters.provider === "string" && sourceFilters.provider.trim()
+      ? sourceFilters.provider
+      : null);
+  const connectionFilter =
+    filters.connectionId ??
+    (typeof sourceFilters.connectionId === "string" && sourceFilters.connectionId.trim()
+      ? sourceFilters.connectionId
+      : null);
+
+  const generatedAt =
+    typeof sourceMeta.generatedAt === "string" && sourceMeta.generatedAt.trim()
+      ? sourceMeta.generatedAt
+      : new Date().toISOString();
+
+  return {
+    providers,
+    meta: {
+      generatedAt,
+      filters: {
+        provider: providerFilter || null,
+        connectionId: connectionFilter || null,
+      },
+      totalProviders: providers.length,
+    },
+  };
+}
--- a/tests/e2e/ecosystem.test.ts
+++ b/tests/e2e/ecosystem.test.ts
@@ -50,11 +50,11 @@ describe("E2E: MCP Server (16 tools)", () => {
  });

  itCase("should return quota data", async () => {
-    const res = await apiFetch("/api/rate-limits");
+    const res = await apiFetch("/api/usage/quota");
    expect(res.ok).toBe(true);
    const data = await res.json();
-    expect(Array.isArray(data?.connections)).toBe(true);
-    expect(data).toHaveProperty("overview");
+    expect(Array.isArray(data?.providers)).toBe(true);
+    expect(data).toHaveProperty("meta");
  });

  itCase("should return usage analytics", async () => {
@@ -68,6 +68,64 @@ describe("E2E: MCP Server (16 tools)", () => {
  });
 });

+// ─── Scenario 1B: Quota Contract ─────────────────────────────
+describe("E2E: Quota Contract (/api/usage/quota)", () => {
+  itCase("should return normalized quota response shape", async () => {
+    const res = await apiFetch("/api/usage/quota");
+    expect(res.ok).toBe(true);
+
+    const data = await res.json();
+    expect(Array.isArray(data.providers)).toBe(true);
+    expect(data).toHaveProperty("meta");
+    expect(typeof data.meta.generatedAt).toBe("string");
+    expect(typeof data.meta.totalProviders).toBe("number");
+
+    if (data.providers.length > 0) {
+      const p = data.providers[0];
+      expect(typeof p.name).toBe("string");
+      expect(typeof p.provider).toBe("string");
+      expect(typeof p.connectionId).toBe("string");
+      expect(typeof p.quotaUsed).toBe("number");
+      expect(typeof p.percentRemaining).toBe("number");
+      expect(p.percentRemaining).toBeGreaterThanOrEqual(0);
+      expect(p.percentRemaining).toBeLessThanOrEqual(100);
+      expect(["valid", "expiring", "expired", "refreshing"]).toContain(p.tokenStatus);
+    }
+  });
+
+  itCase("should filter quota by provider", async () => {
+    const allRes = await apiFetch("/api/usage/quota");
+    expect(allRes.ok).toBe(true);
+    const allData = await allRes.json();
+    if (!Array.isArray(allData.providers) || allData.providers.length === 0) return;
+
+    const provider = allData.providers[0].provider;
+    const filteredRes = await apiFetch(`/api/usage/quota?provider=${encodeURIComponent(provider)}`);
+    expect(filteredRes.ok).toBe(true);
+    const filteredData = await filteredRes.json();
+    expect(filteredData.meta.filters.provider).toBe(provider);
+    expect(Array.isArray(filteredData.providers)).toBe(true);
+    expect(filteredData.providers.every((p: any) => p.provider === provider)).toBe(true);
+  });
+
+  itCase("should filter quota by connectionId", async () => {
+    const allRes = await apiFetch("/api/usage/quota");
+    expect(allRes.ok).toBe(true);
+    const allData = await allRes.json();
+    if (!Array.isArray(allData.providers) || allData.providers.length === 0) return;
+
+    const connectionId = allData.providers[0].connectionId;
+    const filteredRes = await apiFetch(
+      `/api/usage/quota?connectionId=${encodeURIComponent(connectionId)}`
+    );
+    expect(filteredRes.ok).toBe(true);
+    const filteredData = await filteredRes.json();
+    expect(filteredData.meta.filters.connectionId).toBe(connectionId);
+    expect(Array.isArray(filteredData.providers)).toBe(true);
+    expect(filteredData.providers.every((p: any) => p.connectionId === connectionId)).toBe(true);
+  });
+});
+
 // ─── Scenario 2: A2A Server Complete ─────────────────────────────
 describe("E2E: A2A Server (lifecycle)", () => {
  itCase("should serve Agent Card", async () => {
@@ -174,7 +232,7 @@ describe("E2E: Stress (100 parallel requests)", () => {

  itStress("should handle 50 parallel quota checks", async () => {
    const promises = Array.from({ length: 50 }, () =>
-      apiFetch("/api/rate-limits").then((r) => r.ok)
+      apiFetch("/api/usage/quota").then((r) => r.ok)
    );
    const results = await Promise.allSettled(promises);
    const successful = results.filter((r) => r.status === "fulfilled" && r.value).length;
--- a/tests/integration/integration-wiring.test.mjs
+++ b/tests/integration/integration-wiring.test.mjs
@@ -111,6 +111,7 @@ describe("API Routes — existence check", () => {
    "src/app/api/models/availability/route.ts",
    "src/app/api/telemetry/summary/route.ts",
    "src/app/api/usage/budget/route.ts",
+    "src/app/api/usage/quota/route.ts",
    "src/app/api/fallback/chains/route.ts",
    "src/app/api/compliance/audit-log/route.ts",
    "src/app/api/evals/route.ts",
@@ -142,6 +143,10 @@ describe("API Routes — export HTTP methods", () => {
    assertRouteMethods("src/app/api/usage/budget/route.ts", ["GET", "POST"]);
  });

+  it("/api/usage/quota should export GET", () => {
+    assertRouteMethods("src/app/api/usage/quota/route.ts", ["GET"]);
+  });
+
  it("/api/fallback/chains should export GET, POST, DELETE", () => {
    assertRouteMethods("src/app/api/fallback/chains/route.ts", ["GET", "POST", "DELETE"]);
  });