From 5ecef5c90c4a52658c90c53c92ee0216ee9846a3 Mon Sep 17 00:00:00 2001 From: diegosouzapw Date: Wed, 4 Mar 2026 08:18:09 -0300 Subject: [PATCH] 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. --- open-sse/mcp-server/schemas/tools.ts | 12 +- open-sse/mcp-server/server.ts | 16 +- open-sse/mcp-server/tools/advancedTools.ts | 31 ++-- package.json | 2 +- src/lib/a2a/skills/quotaManagement.ts | 48 ++++-- src/shared/contracts/quota.ts | 138 ++++++++++++++++++ tests/e2e/ecosystem.test.ts | 66 ++++++++- tests/integration/integration-wiring.test.mjs | 5 + 8 files changed, 289 insertions(+), 29 deletions(-) create mode 100644 src/shared/contracts/quota.ts diff --git a/open-sse/mcp-server/schemas/tools.ts b/open-sse/mcp-server/schemas/tools.ts index bfc6ab9644..38730de507 100644 --- 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 = { @@ -236,7 +246,7 @@ export const checkQuotaTool: McpToolDefinition = {}; 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) }] }; diff --git a/open-sse/mcp-server/tools/advancedTools.ts b/open-sse/mcp-server/tools/advancedTools.ts index 85df6ccf99..918dae90e4 100644 --- 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 { @@ -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) { diff --git a/package.json b/package.json index 8f1e3110dd..5d7450bbde 100644 --- 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", diff --git a/src/lib/a2a/skills/quotaManagement.ts b/src/lib/a2a/skills/quotaManagement.ts index 87d5054897..354071540f 100644 --- 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 { + 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 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 }, diff --git a/src/shared/contracts/quota.ts b/src/shared/contracts/quota.ts new file mode 100644 index 0000000000..7dbf07af5a --- /dev/null +++ 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) : {}; + 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) : {}; + 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) : {}; + const sourceFilters = + sourceMeta.filters && typeof sourceMeta.filters === "object" + ? (sourceMeta.filters as Record) + : {}; + + 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, + }, + }; +} diff --git a/tests/e2e/ecosystem.test.ts b/tests/e2e/ecosystem.test.ts index cf731afd3c..ee6038eaaf 100644 --- 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; diff --git a/tests/integration/integration-wiring.test.mjs b/tests/integration/integration-wiring.test.mjs index 149a917b03..981296dbb4 100644 --- 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"]); });