From dc6d9e2e4b5cb05984f53978bcd101d1f2efd917 Mon Sep 17 00:00:00 2001 From: diegosouzapw Date: Fri, 17 Apr 2026 09:00:32 -0300 Subject: [PATCH] feat(core): add payload rules, tag routing, and scheduled budgets Introduce runtime-configurable payload mutation/filter rules with file reload support and a settings API so upstream request bodies can be customized per model and protocol without restarts. Expand search support with Google PSE, Linkup, SearchAPI, and SearXNG, including validation, routing, analytics costing, MCP schema updates, and search-type-aware provider selection. Update Pollinations to support anonymous access, endpoint failover, and the latest public model lineup. Add OmniRoute response metadata headers/SSE comments, per-connection model exclusion rules, combo tag-based routing, buffered spend writes, and scheduled daily/weekly/monthly budget resets. Update model catalog and dashboard UIs to surface source labels and hide models excluded by all active connections. --- .env.example | 10 + config/payloadRules.json | 6 + open-sse/config/providerRegistry.ts | 60 +- open-sse/config/searchRegistry.ts | 90 ++- open-sse/executors/pollinations.ts | 23 +- open-sse/handlers/chatCore.ts | 105 +++- open-sse/handlers/search.ts | 362 +++++++++++- open-sse/mcp-server/schemas/tools.ts | 14 +- open-sse/mcp-server/server.ts | 6 +- open-sse/services/combo.ts | 107 +++- open-sse/services/model.ts | 2 + open-sse/services/payloadRules.ts | 449 ++++++++++++++ open-sse/utils/stream.ts | 45 +- .../dashboard/providers/[id]/page.tsx | 355 +++++++++-- .../dashboard/usage/components/BudgetTab.tsx | 96 ++- src/app/api/models/catalog/route.ts | 42 ++ src/app/api/models/route.ts | 40 ++ src/app/api/providers/route.ts | 3 +- src/app/api/providers/validate/route.ts | 4 + src/app/api/settings/payload-rules/route.ts | 50 ++ src/app/api/usage/budget/route.ts | 40 +- src/app/api/v1/models/catalog.ts | 61 ++ src/app/api/v1/search/analytics/route.ts | 13 +- src/app/api/v1/search/route.ts | 30 +- src/domain/connectionModelRules.ts | 82 +++ src/domain/costRules.ts | 552 +++++++++++++++--- src/domain/omnirouteResponseMeta.ts | 91 +++ src/domain/tagRouter.ts | 64 ++ src/instrumentation-node.ts | 19 + src/lib/db/core.ts | 21 +- src/lib/db/domainState.ts | 268 ++++++++- src/lib/gracefulShutdown.ts | 9 +- src/lib/initCloudSync.ts | 2 + src/lib/jobs/budgetResetJob.ts | 40 ++ src/lib/providers/requestDefaults.ts | 37 ++ src/lib/providers/validation.ts | 58 +- src/lib/spend/batchWriter.ts | 208 +++++++ src/server-init.ts | 17 + src/shared/components/ModelSelectModal.tsx | 27 +- src/shared/components/ProviderIcon.tsx | 1 + src/shared/constants/headers.ts | 11 + src/shared/constants/providers.ts | 40 ++ src/shared/utils/modelCatalogSearch.ts | 71 +++ src/shared/validation/schemas.ts | 260 ++++++++- src/sse/handlers/chat.ts | 47 +- src/sse/services/auth.ts | 12 + tests/unit/catalog-updates-v3x.test.ts | 49 ++ tests/unit/chatcore-translation-paths.test.ts | 99 ++++ tests/unit/connection-model-rules.test.ts | 47 ++ tests/unit/domain-cost-rules.test.ts | 131 +++++ tests/unit/executor-pollinations.test.ts | 18 +- tests/unit/model-catalog-search.test.ts | 40 ++ tests/unit/models-catalog-route.test.ts | 45 ++ tests/unit/omniroute-response-meta.test.ts | 68 +++ tests/unit/payload-rules.test.ts | 142 +++++ .../provider-validation-specialty.test.ts | 67 +++ tests/unit/search-handler-extended.test.ts | 210 +++++++ tests/unit/search-registry.test.ts | 140 ++++- tests/unit/search-route.test.ts | 178 ++++++ tests/unit/spend-batch-writer.test.ts | 150 +++++ tests/unit/sse-auth.test.ts | 20 + tests/unit/tag-routing.test.ts | 165 ++++++ 62 files changed, 5277 insertions(+), 242 deletions(-) create mode 100644 config/payloadRules.json create mode 100644 open-sse/services/payloadRules.ts create mode 100644 src/app/api/settings/payload-rules/route.ts create mode 100644 src/domain/connectionModelRules.ts create mode 100644 src/domain/omnirouteResponseMeta.ts create mode 100644 src/domain/tagRouter.ts create mode 100644 src/lib/jobs/budgetResetJob.ts create mode 100644 src/lib/spend/batchWriter.ts create mode 100644 src/shared/constants/headers.ts create mode 100644 src/shared/utils/modelCatalogSearch.ts create mode 100644 tests/unit/catalog-updates-v3x.test.ts create mode 100644 tests/unit/connection-model-rules.test.ts create mode 100644 tests/unit/model-catalog-search.test.ts create mode 100644 tests/unit/omniroute-response-meta.test.ts create mode 100644 tests/unit/payload-rules.test.ts create mode 100644 tests/unit/search-route.test.ts create mode 100644 tests/unit/spend-batch-writer.test.ts create mode 100644 tests/unit/tag-routing.test.ts diff --git a/.env.example b/.env.example index 2e129f54ea..9aa731e283 100644 --- a/.env.example +++ b/.env.example @@ -159,6 +159,16 @@ ALLOW_API_KEY_REVEAL=false # Values: allowlist | denylist | disabled | Default: disabled # TOOL_POLICY_MODE=disabled +# Payload manipulation rules JSON file. +# Used by: open-sse/services/payloadRules.ts — injects/removes upstream payload fields per model/protocol. +# Default: ./config/payloadRules.json +# OMNIROUTE_PAYLOAD_RULES_PATH=./config/payloadRules.json + +# Reload interval for payloadRules.json mtime checks in milliseconds. +# Used by: open-sse/services/payloadRules.ts — keeps file-based rules hot-reloadable without restart. +# Default: 5000 | Minimum: 1000 +# OMNIROUTE_PAYLOAD_RULES_RELOAD_MS=5000 + # ═══════════════════════════════════════════════════════════════════════════════ # 7. URLS & CLOUD SYNC diff --git a/config/payloadRules.json b/config/payloadRules.json new file mode 100644 index 0000000000..ddbd2a4f11 --- /dev/null +++ b/config/payloadRules.json @@ -0,0 +1,6 @@ +{ + "default": [], + "default-raw": [], + "override": [], + "filter": [] +} diff --git a/open-sse/config/providerRegistry.ts b/open-sse/config/providerRegistry.ts index 93edb4e522..f1d5624d9a 100644 --- a/open-sse/config/providerRegistry.ts +++ b/open-sse/config/providerRegistry.ts @@ -1225,13 +1225,23 @@ export const REGISTRY: Record = { models: [ { id: "gpt-oss-120b", name: "GPT OSS 120B", toolCalling: false }, { id: "openai/gpt-oss-120b", name: "GPT OSS 120B (OpenAI Prefix)", toolCalling: false }, + { id: "openai/gpt-oss-20b", name: "GPT OSS 20B", toolCalling: false }, { id: "meta/llama-3.3-70b-instruct", name: "Llama 3.3 70B" }, { id: "nvidia/llama-3.3-70b-instruct", name: "Llama 3.3 70B (NVIDIA Prefix)" }, { id: "meta/llama-4-maverick-17b-128e-instruct", name: "Llama 4 Maverick" }, + { id: "nvidia/llama-3.1-nemotron-ultra-253b-v1", name: "Llama 3.1 Nemotron Ultra 253B" }, + { id: "nvidia/nemotron-3-super-120b-a12b", name: "Nemotron 3 Super 120B A12B" }, + { id: "nvidia/llama-3.3-nemotron-super-49b-v1.5", name: "Llama 3.3 Nemotron Super 49B" }, { id: "moonshotai/kimi-k2.5", name: "Kimi K2.5" }, { id: "z-ai/glm4.7", name: "GLM 4.7" }, { id: "deepseek-ai/deepseek-v3.2", name: "DeepSeek V3.2" }, { id: "deepseek/deepseek-r1", name: "DeepSeek R1" }, + { + id: "mistralai/mistral-large-3-675b-instruct-2512", + name: "Mistral Large 3 675B", + }, + { id: "qwen/qwen3-coder-480b-a35b-instruct", name: "Qwen3 Coder 480B" }, + { id: "mistralai/devstral-2-123b-instruct-2512", name: "Devstral 2 123B" }, { id: "nvidia/llama-3.1-70b-instruct", name: "Llama 3.1 70B" }, { id: "nvidia/llama-3.1-405b-instruct", name: "Llama 3.1 405B" }, ], @@ -1440,17 +1450,47 @@ export const REGISTRY: Record = { alias: "pol", format: "openai", executor: "pollinations", - // API key required. Free Spore tier currently grants 0.01 pollen/hour. + // Primary endpoint is text.pollinations.ai. gen.pollinations.ai is the current + // OpenAI-compatible fallback used when the primary edge is rate-limited or unavailable. baseUrl: "https://text.pollinations.ai/openai/chat/completions", + baseUrls: [ + "https://text.pollinations.ai/openai/chat/completions", + "https://gen.pollinations.ai/v1/chat/completions", + ], authType: "apikey", authHeader: "bearer", models: [ - { id: "openai", name: "GPT-5 via Pollinations (Spore)" }, - { id: "claude", name: "Claude via Pollinations (Spore)" }, - { id: "gemini", name: "Gemini via Pollinations (Spore)" }, - { id: "deepseek", name: "DeepSeek V3 via Pollinations (Spore)" }, - { id: "llama", name: "Llama 4 via Pollinations (Spore)" }, - { id: "mistral", name: "Mistral via Pollinations (Spore)" }, + { id: "openai", name: "OpenAI (Pollinations)" }, + { id: "openai-fast", name: "OpenAI Fast (Pollinations)" }, + { id: "openai-large", name: "OpenAI Large (Pollinations)" }, + { id: "qwen-coder", name: "Qwen Coder (Pollinations)" }, + { id: "mistral", name: "Mistral (Pollinations)" }, + { id: "gemini", name: "Gemini (Pollinations)" }, + { id: "gemini-flash-lite-3.1", name: "Gemini Flash Lite 3.1 (Pollinations)" }, + { id: "gemini-fast", name: "Gemini Fast (Pollinations)" }, + { id: "deepseek", name: "DeepSeek (Pollinations)" }, + { id: "grok", name: "Grok (Pollinations)" }, + { id: "grok-large", name: "Grok Large (Pollinations)" }, + { id: "gemini-search", name: "Gemini Search (Pollinations)" }, + { id: "midijourney", name: "Midijourney (Pollinations)" }, + { id: "midijourney-large", name: "Midijourney Large (Pollinations)" }, + { id: "claude-fast", name: "Claude Fast (Pollinations)" }, + { id: "claude", name: "Claude (Pollinations)" }, + { id: "claude-large", name: "Claude Large (Pollinations)" }, + { id: "perplexity-fast", name: "Perplexity Fast (Pollinations)" }, + { id: "perplexity-reasoning", name: "Perplexity Reasoning (Pollinations)" }, + { id: "kimi", name: "Kimi (Pollinations)" }, + { id: "gemini-large", name: "Gemini Large (Pollinations)" }, + { id: "nova-fast", name: "Nova Fast (Pollinations)" }, + { id: "nova", name: "Nova (Pollinations)" }, + { id: "glm", name: "GLM (Pollinations)" }, + { id: "minimax", name: "MiniMax (Pollinations)" }, + { id: "mistral-large", name: "Mistral Large (Pollinations)" }, + { id: "polly", name: "Polly (Pollinations)" }, + { id: "qwen-coder-large", name: "Qwen Coder Large (Pollinations)" }, + { id: "qwen-large", name: "Qwen Large (Pollinations)" }, + { id: "qwen-vision", name: "Qwen Vision (Pollinations)" }, + { id: "qwen-safety", name: "Qwen Safety (Pollinations)" }, ], }, @@ -1512,6 +1552,12 @@ export const REGISTRY: Record = { { id: "qwen/qwen3-235b-a22b", name: "Qwen3 235B (Puter)" }, { id: "qwen/qwen3-32b", name: "Qwen3 32B (Puter)" }, { id: "qwen/qwen3-coder", name: "Qwen3 Coder 480B (Puter)" }, + // Perplexity Sonar via OpenRouter aliases exposed by Puter + { id: "perplexity/sonar", name: "Perplexity Sonar (Puter)" }, + { id: "perplexity/sonar-pro", name: "Perplexity Sonar Pro (Puter)" }, + { id: "perplexity/sonar-pro-search", name: "Perplexity Sonar Pro Search (Puter)" }, + { id: "perplexity/sonar-reasoning-pro", name: "Perplexity Sonar Reasoning Pro (Puter)" }, + { id: "perplexity/sonar-deep-research", name: "Perplexity Sonar Deep Research (Puter)" }, ], passthroughModels: true, // 500+ models available — users can type arbitrary Puter model IDs }, diff --git a/open-sse/config/searchRegistry.ts b/open-sse/config/searchRegistry.ts index 4c152571e8..2851d2bb0c 100644 --- a/open-sse/config/searchRegistry.ts +++ b/open-sse/config/searchRegistry.ts @@ -15,7 +15,7 @@ export interface SearchProviderConfig { name: string; baseUrl: string; method: "GET" | "POST"; - authType: "apikey"; + authType: "apikey" | "none"; authHeader: string; costPerQuery: number; freeMonthlyQuota: number; @@ -106,6 +106,70 @@ export const SEARCH_PROVIDERS: Record = { timeoutMs: 10_000, cacheTTLMs: 5 * 60 * 1000, }, + + "google-pse-search": { + id: "google-pse-search", + name: "Google Programmable Search", + baseUrl: "https://www.googleapis.com/customsearch/v1", + method: "GET", + authType: "apikey", + authHeader: "key", + costPerQuery: 0.005, + freeMonthlyQuota: 3000, + searchTypes: ["web", "news"], + defaultMaxResults: 5, + maxMaxResults: 10, + timeoutMs: 10_000, + cacheTTLMs: 5 * 60 * 1000, + }, + + "linkup-search": { + id: "linkup-search", + name: "Linkup Search", + baseUrl: "https://api.linkup.so/v1/search", + method: "POST", + authType: "apikey", + authHeader: "bearer", + costPerQuery: 0.005, + freeMonthlyQuota: 1000, + searchTypes: ["web"], + defaultMaxResults: 5, + maxMaxResults: 50, + timeoutMs: 10_000, + cacheTTLMs: 5 * 60 * 1000, + }, + + "searchapi-search": { + id: "searchapi-search", + name: "SearchAPI", + baseUrl: "https://www.searchapi.io/api/v1/search", + method: "GET", + authType: "apikey", + authHeader: "api_key", + costPerQuery: 0.004, + freeMonthlyQuota: 100, + searchTypes: ["web", "news"], + defaultMaxResults: 5, + maxMaxResults: 100, + timeoutMs: 10_000, + cacheTTLMs: 5 * 60 * 1000, + }, + + "searxng-search": { + id: "searxng-search", + name: "SearXNG Search", + baseUrl: "http://localhost:8888/search", + method: "GET", + authType: "none", + authHeader: "none", + costPerQuery: 0, + freeMonthlyQuota: 999999, + searchTypes: ["web", "news"], + defaultMaxResults: 5, + maxMaxResults: 50, + timeoutMs: 10_000, + cacheTTLMs: 3 * 60 * 1000, + }, }; /** @@ -123,6 +187,16 @@ export function getSearchProvider(providerId: string): SearchProviderConfig | nu return SEARCH_PROVIDERS[providerId] || null; } +export function supportsSearchType( + providerOrId: SearchProviderConfig | string | null | undefined, + searchType: string +): boolean { + const provider = + typeof providerOrId === "string" ? getSearchProvider(providerOrId) : providerOrId || null; + if (!provider) return false; + return provider.searchTypes.includes(searchType); +} + /** * Get all search providers as a flat list */ @@ -143,12 +217,20 @@ export function getAllSearchProviders(): Array<{ * If an explicit provider is given, validate and return it. * Otherwise, return the cheapest by costPerQuery. */ -export function selectProvider(explicitProvider?: string): SearchProviderConfig | null { +export function selectProvider( + explicitProvider?: string, + searchType?: string +): SearchProviderConfig | null { if (explicitProvider) { - return SEARCH_PROVIDERS[explicitProvider] || null; + const provider = SEARCH_PROVIDERS[explicitProvider] || null; + if (!provider) return null; + if (searchType && !supportsSearchType(provider, searchType)) return null; + return provider; } - const providers = Object.values(SEARCH_PROVIDERS); + const providers = Object.values(SEARCH_PROVIDERS).filter((provider) => + searchType ? supportsSearchType(provider, searchType) : true + ); if (providers.length === 0) return null; return providers.reduce((cheapest, p) => (p.costPerQuery < cheapest.costPerQuery ? p : cheapest)); diff --git a/open-sse/executors/pollinations.ts b/open-sse/executors/pollinations.ts index 5b4e41a8b7..653026fdf3 100644 --- a/open-sse/executors/pollinations.ts +++ b/open-sse/executors/pollinations.ts @@ -2,9 +2,11 @@ import { BaseExecutor } from "./base.ts"; import { PROVIDERS } from "../config/constants.ts"; /** - * PollinationsExecutor — Pollinations now requires API key auth. - * The free Spore tier grants 0.01 pollen/hour, so keep the messaging - * aligned with a key-backed free tier instead of anonymous access. + * PollinationsExecutor — OpenAI-compatible Pollinations text endpoint. + * + * Pollinations currently exposes a public endpoint and an optional key-backed tier. + * OmniRoute sends the bearer token when configured, but no auth header is required + * for the anonymous endpoint. * * Endpoint: https://text.pollinations.ai/openai/chat/completions * Docs: https://pollinations.ai/docs @@ -14,21 +16,24 @@ export class PollinationsExecutor extends BaseExecutor { super("pollinations", PROVIDERS["pollinations"] || { format: "openai" }); } - buildUrl(_model: string, _stream: boolean, _urlIndex = 0, _credentials = null): string { - return "https://text.pollinations.ai/openai/chat/completions"; + buildUrl(_model: string, _stream: boolean, urlIndex = 0, _credentials = null): string { + const baseUrls = this.getBaseUrls(); + return ( + baseUrls[urlIndex] || baseUrls[0] || "https://text.pollinations.ai/openai/chat/completions" + ); } buildHeaders(credentials: any, stream = true): Record { const key = credentials?.apiKey || credentials?.accessToken; - if (!key) { - throw new Error("Pollinations API key is required"); - } const headers: Record = { "Content-Type": "application/json", - Authorization: `Bearer ${key}`, }; + if (key) { + headers.Authorization = `Bearer ${key}`; + } + if (stream) { headers["Accept"] = "text/event-stream"; } diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 81b3cb6e08..b0af6fab8b 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -45,6 +45,7 @@ import { } from "@/lib/usage/tokenAccounting"; import { recordCost } from "@/domain/costRules"; import { calculateCost } from "@/lib/usage/costCalculator"; +import { buildOmniRouteResponseMetaHeaders } from "@/domain/omnirouteResponseMeta"; import { CLAUDE_OAUTH_TOOL_PREFIX } from "../translator/request/openai-to-claude.ts"; import { getModelNormalizeToolCallId, @@ -54,6 +55,10 @@ import { } from "@/lib/localDb"; import { getExecutor } from "../executors/index.ts"; import { getCacheControlSettings } from "@/lib/cacheControlSettings"; +import { + applyConfiguredPayloadRules, + resolvePayloadRuleProtocols, +} from "../services/payloadRules.ts"; import { shouldPreserveCacheControl, providerSupportsCaching, @@ -127,6 +132,7 @@ import { } from "@/lib/memory/settings"; import { injectSkills } from "@/lib/skills/injection"; import { handleToolCallExecution } from "@/lib/skills/interception"; +import { OMNIROUTE_RESPONSE_HEADERS } from "@/shared/constants/headers"; import { buildClaudeCodeCompatibleRequest, isClaudeCodeCompatibleProvider, @@ -641,6 +647,15 @@ export async function handleChatCore({ const cachedIdemp = checkIdempotency(idempotencyKey); if (cachedIdemp) { log?.debug?.("IDEMPOTENCY", `Hit for key=${idempotencyKey?.slice(0, 12)}...`); + const idempotentUsage = + cachedIdemp.response && typeof cachedIdemp.response === "object" + ? ((cachedIdemp.response as Record).usage as + | Record + | undefined) + : undefined; + const idempotentCost = idempotentUsage + ? await calculateCost(provider, model, idempotentUsage) + : 0; return { success: true, response: new Response(JSON.stringify(cachedIdemp.response), { @@ -649,6 +664,14 @@ export async function handleChatCore({ "Content-Type": "application/json", "Access-Control-Allow-Origin": getCorsOrigin(), "X-OmniRoute-Idempotent": "true", + ...buildOmniRouteResponseMetaHeaders({ + provider, + model, + cacheHit: false, + latencyMs: Date.now() - startTime, + usage: idempotentUsage, + costUsd: idempotentCost, + }), }, }), }; @@ -907,6 +930,10 @@ export async function handleChatCore({ if (cached) { log?.debug?.("CACHE", `Semantic cache HIT for ${model} (stream=${stream})`); reqLogger.logConvertedResponse(cached as Record); + const cachedUsage = + extractUsageFromResponse(cached as Record, provider) || + ((cached as Record)?.usage as Record | undefined); + const cachedCost = cachedUsage ? await calculateCost(provider, model, cachedUsage) : 0; persistAttemptLogs({ status: 200, tokens: (cached as Record)?.usage, @@ -922,7 +949,15 @@ export async function handleChatCore({ headers: { "Content-Type": "application/json", "Access-Control-Allow-Origin": getCorsOrigin(), - "X-OmniRoute-Cache": "HIT", + [OMNIROUTE_RESPONSE_HEADERS.cache]: "HIT", + ...buildOmniRouteResponseMetaHeaders({ + provider, + model, + cacheHit: true, + latencyMs: Date.now() - startTime, + usage: cachedUsage, + costUsd: cachedCost, + }), }, }), }; @@ -1538,6 +1573,38 @@ export async function handleChatCore({ translatedBody.model === modelToCall ? translatedBody : { ...translatedBody, model: modelToCall }; + const payloadRuleModel = + typeof bodyToSend.model === "string" && bodyToSend.model.length > 0 + ? bodyToSend.model + : modelToCall; + const payloadRuleProtocols = resolvePayloadRuleProtocols({ + provider, + targetFormat, + }); + const payloadRuleResult = await applyConfiguredPayloadRules( + bodyToSend, + payloadRuleModel, + payloadRuleProtocols + ); + bodyToSend = payloadRuleResult.payload; + + if (payloadRuleResult.applied.length > 0) { + const appliedSummary = payloadRuleResult.applied + .map((rule) => { + if (rule.type === "filter") return `${rule.type}:${rule.path}`; + const serializedValue = JSON.stringify(rule.value); + const safeValue = + typeof serializedValue === "string" && serializedValue.length > 80 + ? `${serializedValue.slice(0, 77)}...` + : serializedValue; + return `${rule.type}:${rule.path}=${safeValue}`; + }) + .join(", "); + log?.debug?.( + "PAYLOAD_RULES", + `Applied ${payloadRuleResult.applied.length} rule(s) for ${payloadRuleModel} (${payloadRuleProtocols.join(", ")}): ${appliedSummary}` + ); + } // Qwen OAuth rejects requests without a non-empty `user` field. // Some minimal OpenAI-compatible clients omit it, so we backfill a @@ -2387,11 +2454,6 @@ export async function handleChatCore({ }); } - if (apiKeyInfo?.id && usage) { - const estimatedCost = await calculateCost(provider, model, usage); - if (estimatedCost > 0) recordCost(apiKeyInfo.id, estimatedCost); - } - // Translate response to client's expected format (usually OpenAI) // Pass toolNameMap so Claude OAuth proxy_ prefix is stripped in tool_use blocks (#605) let translatedResponse = needsTranslation(responsePayloadFormat, clientResponseFormat) @@ -2515,13 +2577,31 @@ export async function handleChatCore({ cacheSource: "upstream", }); + const responseUsage = + (usage && typeof usage === "object" ? usage : null) || + (translatedResponse?.usage && typeof translatedResponse.usage === "object" + ? translatedResponse.usage + : null); + const estimatedCost = responseUsage ? await calculateCost(provider, model, responseUsage) : 0; + if (apiKeyInfo?.id && estimatedCost > 0) { + recordCost(apiKeyInfo.id, estimatedCost); + } + return { success: true, response: new Response(JSON.stringify(translatedResponse), { headers: { "Content-Type": "application/json", "Access-Control-Allow-Origin": getCorsOrigin(), - "X-OmniRoute-Cache": "MISS", + [OMNIROUTE_RESPONSE_HEADERS.cache]: "MISS", + ...buildOmniRouteResponseMetaHeaders({ + provider, + model, + cacheHit: false, + latencyMs: Date.now() - startTime, + usage: responseUsage, + costUsd: estimatedCost, + }), }, }), }; @@ -2539,6 +2619,15 @@ export async function handleChatCore({ "Cache-Control": "no-cache", Connection: "keep-alive", "Access-Control-Allow-Origin": getCorsOrigin(), + [OMNIROUTE_RESPONSE_HEADERS.cache]: "MISS", + ...buildOmniRouteResponseMetaHeaders({ + provider, + model, + cacheHit: false, + latencyMs: 0, + usage: null, + costUsd: 0, + }), }; // Create transform stream with logger for streaming response @@ -2701,7 +2790,7 @@ export async function handleChatCore({ // Chain: provider → transform → progress → client const transformedBody = pipeWithDisconnect(providerResponse, transformStream, streamController); finalStream = transformedBody.pipeThrough(progressTransform); - responseHeaders["X-OmniRoute-Progress"] = "enabled"; + responseHeaders[OMNIROUTE_RESPONSE_HEADERS.progress] = "enabled"; } else { finalStream = pipeWithDisconnect(providerResponse, transformStream, streamController); } diff --git a/open-sse/handlers/search.ts b/open-sse/handlers/search.ts index 4a8e0605b4..42e88c3969 100644 --- a/open-sse/handlers/search.ts +++ b/open-sse/handlers/search.ts @@ -3,8 +3,9 @@ import { randomUUID } from "crypto"; * Search Handler * * Handles POST /v1/search requests. - * Routes to 5 search providers with automatic failover: - * serper-search, brave-search, perplexity-search, exa-search, tavily-search + * Routes to 9 search providers with automatic failover: + * serper-search, brave-search, perplexity-search, exa-search, tavily-search, + * google-pse-search, linkup-search, searchapi-search, searxng-search * * Request format: * { @@ -231,16 +232,47 @@ function parseDomainFilter(domainFilter?: string[]): { return { includes, excludes }; } +function getProviderSettingString( + params: Pick, + key: string +): string | undefined { + const fromOptions = params.providerOptions?.[key]; + if (typeof fromOptions === "string" && fromOptions.trim().length > 0) { + return fromOptions.trim(); + } + + const fromProviderData = params.providerSpecificData?.[key]; + if (typeof fromProviderData === "string" && fromProviderData.trim().length > 0) { + return fromProviderData.trim(); + } + + return undefined; +} + +function resolveSearchBaseUrl(config: SearchProviderConfig, params: SearchRequestParams): string { + const override = getProviderSettingString(params, "baseUrl"); + return (override || config.baseUrl).replace(/\/+$/, ""); +} + +function toSearchPageNumber(offset: number | undefined, maxResults: number): number | undefined { + if (typeof offset !== "number" || offset <= 0 || maxResults <= 0) return undefined; + return Math.floor(offset / maxResults) + 1; +} + // ── Provider Request Builders ─────────────────────────────────────────── interface SearchRequestParams { query: string; searchType: string; maxResults: number; - token: string; + token?: string; country?: string; language?: string; + timeRange?: string; + offset?: number; domainFilter?: string[]; + providerOptions?: Record; + providerSpecificData?: Record; } function buildSerperRequest( @@ -344,6 +376,155 @@ function buildTavilyRequest( }; } +function buildGooglePseRequest( + config: SearchProviderConfig, + params: SearchRequestParams +): { url: string; init: RequestInit } { + const apiKey = params.token; + const cx = getProviderSettingString(params, "cx"); + if (!apiKey || !cx) { + throw new Error("Google Programmable Search requires both apiKey and cx"); + } + + const qp = new URLSearchParams({ + key: apiKey, + cx, + q: params.query, + num: String(Math.min(params.maxResults, 10)), + }); + + if (params.country) qp.set("gl", params.country.toLowerCase()); + if (params.language) qp.set("hl", params.language); + if (params.timeRange && params.timeRange !== "any") { + const dateRestrictMap: Record = { + day: "d1", + week: "w1", + month: "m1", + year: "y1", + }; + const dateRestrict = dateRestrictMap[params.timeRange]; + if (dateRestrict) qp.set("dateRestrict", dateRestrict); + } + if (typeof params.offset === "number" && params.offset > 0) { + qp.set("start", String(Math.min(params.offset + 1, 91))); + } + + return { + url: `${resolveSearchBaseUrl(config, params)}?${qp}`, + init: { + method: "GET", + headers: { Accept: "application/json" }, + }, + }; +} + +function buildLinkupRequest( + config: SearchProviderConfig, + params: SearchRequestParams +): { url: string; init: RequestInit } { + const apiKey = params.token; + if (!apiKey) { + throw new Error("Linkup Search requires an API key"); + } + + const { includes, excludes } = parseDomainFilter(params.domainFilter); + const requestedDepth = getProviderSettingString(params, "depth"); + const depth = + requestedDepth && ["fast", "standard", "deep"].includes(requestedDepth) + ? requestedDepth + : "standard"; + + const body: Record = { + q: params.query, + depth, + outputType: "searchResults", + maxResults: params.maxResults, + }; + + if (includes.length) body.includeDomains = includes; + if (excludes.length) body.excludeDomains = excludes; + if (params.timeRange && params.timeRange !== "any") { + const today = new Date(); + const toDate = today.toISOString().slice(0, 10); + const from = new Date(today); + if (params.timeRange === "day") from.setUTCDate(from.getUTCDate() - 1); + if (params.timeRange === "week") from.setUTCDate(from.getUTCDate() - 7); + if (params.timeRange === "month") from.setUTCMonth(from.getUTCMonth() - 1); + if (params.timeRange === "year") from.setUTCFullYear(from.getUTCFullYear() - 1); + body.fromDate = from.toISOString().slice(0, 10); + body.toDate = toDate; + } + + return { + url: resolveSearchBaseUrl(config, params), + init: { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${apiKey}`, + }, + body: JSON.stringify(body), + }, + }; +} + +function buildSearchApiRequest( + config: SearchProviderConfig, + params: SearchRequestParams +): { url: string; init: RequestInit } { + const apiKey = params.token; + if (!apiKey) { + throw new Error("SearchAPI requires an API key"); + } + + const qp = new URLSearchParams({ + engine: params.searchType === "news" ? "google_news" : "google", + q: params.query, + api_key: apiKey, + }); + + if (params.country) qp.set("gl", params.country.toLowerCase()); + if (params.language) qp.set("hl", params.language); + + const page = toSearchPageNumber(params.offset, params.maxResults); + if (page) qp.set("page", String(page)); + + return { + url: `${resolveSearchBaseUrl(config, params)}?${qp}`, + init: { + method: "GET", + headers: { Accept: "application/json" }, + }, + }; +} + +function buildSearxngRequest( + config: SearchProviderConfig, + params: SearchRequestParams +): { url: string; init: RequestInit } { + const baseUrl = resolveSearchBaseUrl(config, params); + const url = baseUrl.endsWith("/search") ? baseUrl : `${baseUrl}/search`; + const qp = new URLSearchParams({ + q: params.query, + format: "json", + categories: params.searchType === "news" ? "news" : "general", + }); + + if (params.language) qp.set("language", params.language); + if (params.timeRange && params.timeRange !== "any") qp.set("time_range", params.timeRange); + + const page = toSearchPageNumber(params.offset, params.maxResults); + if (page) qp.set("pageno", String(page)); + + return { + url: `${url}?${qp}`, + init: { + method: "GET", + headers: { Accept: "application/json" }, + }, + }; +} + function buildRequest( config: SearchProviderConfig, params: SearchRequestParams @@ -353,12 +534,19 @@ function buildRequest( if (config.id === "perplexity-search") return buildPerplexityRequest(config, params); if (config.id === "exa-search") return buildExaRequest(config, params); if (config.id === "tavily-search") return buildTavilyRequest(config, params); + if (config.id === "google-pse-search") return buildGooglePseRequest(config, params); + if (config.id === "linkup-search") return buildLinkupRequest(config, params); + if (config.id === "searchapi-search") return buildSearchApiRequest(config, params); + if (config.id === "searxng-search") return buildSearxngRequest(config, params); // Fallback for future providers: POST with bearer auth return { - url: config.baseUrl, + url: resolveSearchBaseUrl(config, params), init: { method: config.method, - headers: { "Content-Type": "application/json", Authorization: `Bearer ${params.token}` }, + headers: { + "Content-Type": "application/json", + ...(params.token ? { Authorization: `Bearer ${params.token}` } : {}), + }, body: JSON.stringify({ query: params.query, max_results: params.maxResults, @@ -452,6 +640,139 @@ function normalizeTavilyResponse( return { results, totalResults: results.length }; } +function normalizeGooglePseResponse( + data: any, + _query: string, + _searchType: string +): { results: SearchResult[]; totalResults: number | null } { + const now = new Date().toISOString(); + const items = Array.isArray(data.items) ? data.items : []; + const results = items.map((item: any, idx: number) => + makeResult( + "google-pse-search", + { + title: item.title, + url: item.link, + snippet: item.snippet, + image_url: + item.pagemap?.cse_image?.[0]?.src || + item.pagemap?.cse_thumbnail?.[0]?.src || + item.pagemap?.metatags?.[0]?.["og:image"], + }, + idx, + now + ) + ); + + const totalResultsRaw = + data.searchInformation?.totalResults ?? data.queries?.request?.[0]?.totalResults ?? null; + const totalResults = + typeof totalResultsRaw === "string" ? Number(totalResultsRaw) : totalResultsRaw; + + return { + results, + totalResults: Number.isFinite(totalResults) ? totalResults : null, + }; +} + +function normalizeLinkupResponse( + data: any, + _query: string, + _searchType: string +): { results: SearchResult[]; totalResults: number | null } { + const now = new Date().toISOString(); + const items = Array.isArray(data.results) ? data.results : []; + const results = items.map((item: any, idx: number) => + makeResult( + "linkup-search", + { + title: item.name || item.title, + url: item.url, + snippet: item.content || item.snippet || "", + source_type: item.type || "web", + image_url: item.image_url || item.imageUrl || null, + full_text: item.content, + text_format: "text", + }, + idx, + now + ) + ); + + return { results, totalResults: results.length }; +} + +function normalizeSearchApiResponse( + data: any, + _query: string, + _searchType: string +): { results: SearchResult[]; totalResults: number | null } { + const now = new Date().toISOString(); + const items = Array.isArray(data.organic_results) + ? data.organic_results + : Array.isArray(data.top_stories) + ? data.top_stories + : []; + + const results = items.map((item: any, idx: number) => + makeResult( + "searchapi-search", + { + title: item.title, + url: item.link, + snippet: item.snippet || item.description || "", + published_at: item.date || item.published_at, + favicon_url: item.favicon, + author: item.source || null, + image_url: item.thumbnail || null, + }, + idx, + now + ) + ); + + const totalResults = + typeof data.search_information?.total_results === "number" + ? data.search_information.total_results + : typeof data.search_information?.total_results === "string" + ? Number(data.search_information.total_results) + : null; + + return { + results, + totalResults: Number.isFinite(totalResults) ? totalResults : results.length, + }; +} + +function normalizeSearxngResponse( + data: any, + _query: string, + _searchType: string +): { results: SearchResult[]; totalResults: number | null } { + const now = new Date().toISOString(); + const items = Array.isArray(data.results) ? data.results : []; + + const results = items.map((item: any, idx: number) => + makeResult( + "searxng-search", + { + title: item.title, + url: item.url, + snippet: item.content || item.snippet || "", + published_at: item.publishedDate || item.published_date || null, + source_type: Array.isArray(item.engines) + ? item.engines.join(", ") + : item.engine || item.category || null, + image_url: item.thumbnail || item.img_src || null, + }, + idx, + now + ) + ); + + return { results, totalResults: results.length }; +} + function normalizeResponse( providerId: string, data: any, @@ -464,6 +785,11 @@ function normalizeResponse( return normalizePerplexityResponse(data, query, searchType); if (providerId === "exa-search") return normalizeExaResponse(data, query, searchType); if (providerId === "tavily-search") return normalizeTavilyResponse(data, query, searchType); + if (providerId === "google-pse-search") + return normalizeGooglePseResponse(data, query, searchType); + if (providerId === "linkup-search") return normalizeLinkupResponse(data, query, searchType); + if (providerId === "searchapi-search") return normalizeSearchApiResponse(data, query, searchType); + if (providerId === "searxng-search") return normalizeSearxngResponse(data, query, searchType); return { results: [], totalResults: null }; } @@ -477,7 +803,10 @@ export async function handleSearch(options: SearchHandlerOptions): Promise { const startTime = Date.now(); - const token = credentials.apiKey || credentials.accessToken; + const providerSpecificData = + credentials?.providerSpecificData && typeof credentials.providerSpecificData === "object" + ? credentials.providerSpecificData + : undefined; + const token = credentials.apiKey || credentials.accessToken || undefined; - if (!token) { + if (config.authType !== "none" && !token) { return { success: false, status: 401, @@ -565,7 +901,17 @@ async function tryProvider( } const { query, searchType, maxResults } = params; - const { url, init } = buildRequest(config, { ...params, token }); + let url = ""; + let init: RequestInit = {}; + try { + ({ url, init } = buildRequest(config, { ...params, token, providerSpecificData })); + } catch (err: any) { + return { + success: false, + status: 400, + error: err?.message || `Invalid search configuration for provider: ${config.id}`, + }; + } // Timeout: min of provider timeout and remaining global timeout const remainingGlobal = GLOBAL_TIMEOUT_MS - (Date.now() - globalStartTime); diff --git a/open-sse/mcp-server/schemas/tools.ts b/open-sse/mcp-server/schemas/tools.ts index 9138b39abb..3f4c6e1dc1 100644 --- a/open-sse/mcp-server/schemas/tools.ts +++ b/open-sse/mcp-server/schemas/tools.ts @@ -417,7 +417,17 @@ export const webSearchInput = z.object({ .describe("Maximum number of search results to return"), search_type: z.enum(["web", "news"]).default("web").describe("Type of search to perform"), provider: z - .enum(["serper-search", "brave-search", "perplexity-search", "exa-search", "tavily-search"]) + .enum([ + "serper-search", + "brave-search", + "perplexity-search", + "exa-search", + "tavily-search", + "google-pse-search", + "linkup-search", + "searchapi-search", + "searxng-search", + ]) .optional() .describe("Specific search provider to use"), }); @@ -445,7 +455,7 @@ export const webSearchOutput = z.object({ export const webSearchTool: McpToolDefinition = { name: "omniroute_web_search", description: - "Performs a web search using OmniRoute's search gateway. Supports multiple providers (Serper, Brave, Perplexity, Exa, Tavily) with automatic failover. Returns search results with titles, URLs, snippets, and position data.", + "Performs a web search using OmniRoute's search gateway. Supports multiple providers (Serper, Brave, Perplexity, Exa, Tavily, Google PSE, Linkup, SearchAPI, SearXNG) with automatic failover. Returns search results with titles, URLs, snippets, and position data.", inputSchema: webSearchInput, outputSchema: webSearchOutput, scopes: ["execute:search"], diff --git a/open-sse/mcp-server/server.ts b/open-sse/mcp-server/server.ts index b158bc4277..c884204fe4 100644 --- a/open-sse/mcp-server/server.ts +++ b/open-sse/mcp-server/server.ts @@ -525,7 +525,11 @@ async function handleWebSearch(args: { | "brave-search" | "perplexity-search" | "exa-search" - | "tavily-search"; + | "tavily-search" + | "google-pse-search" + | "linkup-search" + | "searchapi-search" + | "searxng-search"; }) { const start = Date.now(); try { diff --git a/open-sse/services/combo.ts b/open-sse/services/combo.ts index d0434d494c..2582a45280 100644 --- a/open-sse/services/combo.ts +++ b/open-sse/services/combo.ts @@ -33,12 +33,19 @@ import { import { supportsToolCalling } from "./modelCapabilities.ts"; import { getSessionConnection } from "./sessionManager.ts"; import { getModelContextLimit } from "../../src/lib/modelCapabilities"; +import { getProviderConnections } from "../../src/lib/db/providers"; import { getComboModelString, getComboStepTarget, getComboStepWeight, normalizeComboStep, } from "../../src/lib/combos/steps.ts"; +import { + getConnectionRoutingTags, + matchesRoutingTags, + resolveRequestRoutingTags, + type RoutingTagMatchMode, +} from "../../src/domain/tagRouter.ts"; // Status codes that should mark semaphore + record circuit breaker failures const TRANSIENT_FOR_BREAKER = [429, 502, 503, 504]; @@ -81,6 +88,7 @@ type ResolvedComboTarget = { provider: string; providerId: string | null; connectionId: string | null; + allowedConnectionIds?: string[] | null; weight: number; label: string | null; }; @@ -840,6 +848,98 @@ function dedupeTargetsByExecutionKey(targets: ResolvedComboTarget[]) { }); } +async function applyRequestTagRouting( + targets: ResolvedComboTarget[], + body: Record | null | undefined, + log: { info?: (...args: unknown[]) => void; warn?: (...args: unknown[]) => void } +): Promise { + const { tags, matchMode } = resolveRequestRoutingTags(body); + if (tags.length === 0 || targets.length === 0) { + return targets; + } + + const providerIds = Array.from( + new Set(targets.map((target) => target.providerId || target.provider)) + ).filter( + (providerId): providerId is string => typeof providerId === "string" && providerId.length > 0 + ); + const providerConnections = new Map>>(); + + await Promise.all( + providerIds.map(async (providerId) => { + try { + const connections = await getProviderConnections({ provider: providerId, isActive: true }); + providerConnections.set( + providerId, + Array.isArray(connections) ? (connections as Array>) : [] + ); + } catch (error) { + log.warn?.( + "COMBO", + `Tag routing failed to load connections for provider=${providerId}: ${error instanceof Error ? error.message : String(error)}` + ); + providerConnections.set(providerId, []); + } + }) + ); + + const filteredTargets = targets.reduce((acc, target) => { + const providerKey = target.providerId || target.provider; + const candidateConnections = + providerConnections.get(providerKey)?.filter((connection) => { + const connectionId = + typeof connection.id === "string" && connection.id.trim().length > 0 + ? connection.id + : null; + if (!connectionId) return false; + if (target.connectionId) { + return connectionId === target.connectionId; + } + return true; + }) || []; + + const matchedConnectionIds = candidateConnections + .filter((connection) => + matchesRoutingTags( + getConnectionRoutingTags(connection.providerSpecificData), + tags, + matchMode as RoutingTagMatchMode + ) + ) + .map((connection) => connection.id) + .filter((connectionId): connectionId is string => typeof connectionId === "string"); + + if (matchedConnectionIds.length === 0) { + return acc; + } + + if (target.connectionId) { + acc.push(target); + return acc; + } + + acc.push({ + ...target, + allowedConnectionIds: Array.from(new Set(matchedConnectionIds)), + }); + return acc; + }, []); + + if (filteredTargets.length === 0) { + log.info?.( + "COMBO", + `Tag routing matched 0/${targets.length} targets for [${tags.join(", ")}] (${matchMode}); falling back to the full target set` + ); + return targets; + } + + log.info?.( + "COMBO", + `Tag routing matched ${filteredTargets.length}/${targets.length} targets for [${tags.join(", ")}] (${matchMode})` + ); + return filteredTargets; +} + export function resolveComboTargets(combo, allCombos) { return allCombos ? resolveNestedComboTargets(combo, allCombos) : getDirectComboTargets(combo); } @@ -1138,6 +1238,8 @@ export async function handleComboChat({ ? resolveWeightedTargets(combo, allCombos)?.orderedTargets || [] : resolveComboTargets(combo, allCombos); + orderedTargets = await applyRequestTagRouting(orderedTargets, body, log); + if (strategy === "weighted") { log.info( "COMBO", @@ -1658,7 +1760,8 @@ async function handleRoundRobinCombo({ const retryDelayMs = config.retryDelayMs ?? 2000; const orderedTargets = resolveComboTargets(combo, allCombos); - const modelCount = orderedTargets.length; + const filteredTargets = await applyRequestTagRouting(orderedTargets, body, log); + const modelCount = filteredTargets.length; if (modelCount === 0) { return comboModelNotFoundResponse("Round-robin combo has no executable targets"); } @@ -1679,7 +1782,7 @@ async function handleRoundRobinCombo({ // Try each model starting from the round-robin target for (let offset = 0; offset < modelCount; offset++) { const modelIndex = (startIndex + offset) % modelCount; - const target = orderedTargets[modelIndex]; + const target = filteredTargets[modelIndex]; const modelStr = target.modelStr; const provider = target.provider; const profile = await getRuntimeProviderProfile(provider); diff --git a/open-sse/services/model.ts b/open-sse/services/model.ts index e68e06f1dc..5c5e7bb470 100644 --- a/open-sse/services/model.ts +++ b/open-sse/services/model.ts @@ -36,6 +36,8 @@ const PROVIDER_MODEL_ALIASES = { nvidia: { "gpt-oss-120b": "openai/gpt-oss-120b", "nvidia/gpt-oss-120b": "openai/gpt-oss-120b", + "gpt-oss-20b": "openai/gpt-oss-20b", + "nvidia/gpt-oss-20b": "openai/gpt-oss-20b", }, antigravity: ANTIGRAVITY_MODEL_ALIASES, }; diff --git a/open-sse/services/payloadRules.ts b/open-sse/services/payloadRules.ts new file mode 100644 index 0000000000..3f3102f69d --- /dev/null +++ b/open-sse/services/payloadRules.ts @@ -0,0 +1,449 @@ +import fs from "node:fs/promises"; +import path from "node:path"; +import { wildcardMatch } from "./wildcardRouter.ts"; + +type JsonRecord = Record; + +export type PayloadRuleModelSpec = { + name: string; + protocol?: string; +}; + +export type PayloadMutationRule = { + models: PayloadRuleModelSpec[]; + params: Record; +}; + +export type PayloadFilterRule = { + models: PayloadRuleModelSpec[]; + params: string[]; +}; + +export type PayloadRulesConfig = { + default: PayloadMutationRule[]; + override: PayloadMutationRule[]; + filter: PayloadFilterRule[]; + defaultRaw: PayloadMutationRule[]; +}; + +export type AppliedPayloadRule = { + type: "default" | "override" | "filter" | "default-raw"; + path: string; + value?: unknown; +}; + +const DEFAULT_PAYLOAD_RULES_CONFIG: PayloadRulesConfig = { + default: [], + override: [], + filter: [], + defaultRaw: [], +}; + +const MIN_FILE_CHECK_INTERVAL_MS = 1_000; +const DEFAULT_FILE_CHECK_INTERVAL_MS = 5_000; + +let runtimeOverride: PayloadRulesConfig | null = null; +let cachedFileConfig = clonePayloadRulesConfig(DEFAULT_PAYLOAD_RULES_CONFIG); +let cachedFilePath = ""; +let cachedFileMtimeMs = -1; +let lastFileCheckAt = 0; +let fileLoadPromise: Promise | null = null; +let lastFileErrorSignature = ""; + +function toRecord(value: unknown): JsonRecord { + return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; +} + +function toArray(value: unknown): T[] { + return Array.isArray(value) ? (value as T[]) : []; +} + +function cloneValue(value: T): T { + return structuredClone(value); +} + +function clonePayloadRulesConfig(config: PayloadRulesConfig): PayloadRulesConfig { + return { + default: config.default.map((rule) => ({ + models: rule.models.map((model) => ({ ...model })), + params: cloneValue(rule.params), + })), + override: config.override.map((rule) => ({ + models: rule.models.map((model) => ({ ...model })), + params: cloneValue(rule.params), + })), + filter: config.filter.map((rule) => ({ + models: rule.models.map((model) => ({ ...model })), + params: [...rule.params], + })), + defaultRaw: config.defaultRaw.map((rule) => ({ + models: rule.models.map((model) => ({ ...model })), + params: cloneValue(rule.params), + })), + }; +} + +function normalizeModelSpecs(value: unknown): PayloadRuleModelSpec[] { + return toArray(value) + .map((item) => { + const name = typeof item?.name === "string" ? item.name.trim() : ""; + const protocol = typeof item?.protocol === "string" ? item.protocol.trim() : ""; + if (!name) return null; + return protocol ? { name, protocol } : { name }; + }) + .filter((item): item is PayloadRuleModelSpec => !!item); +} + +function normalizeMutationRules(value: unknown): PayloadMutationRule[] { + return toArray(value) + .map((item) => { + const models = normalizeModelSpecs(item?.models); + const params = toRecord(item?.params); + if (models.length === 0 || Object.keys(params).length === 0) return null; + return { models, params }; + }) + .filter((item): item is PayloadMutationRule => !!item); +} + +function normalizeFilterRules(value: unknown): PayloadFilterRule[] { + return toArray(value) + .map((item) => { + const models = normalizeModelSpecs(item?.models); + const params = toArray(item?.params) + .map((pathValue) => (typeof pathValue === "string" ? pathValue.trim() : "")) + .filter(Boolean); + if (models.length === 0 || params.length === 0) return null; + return { models, params }; + }) + .filter((item): item is PayloadFilterRule => !!item); +} + +export function normalizePayloadRulesConfig(value: unknown): PayloadRulesConfig { + const record = toRecord(value); + const defaultRawLegacy = toArray(record["default-raw"]); + const defaultRaw = [...toArray(record.defaultRaw), ...defaultRawLegacy]; + + return { + default: normalizeMutationRules(record.default), + override: normalizeMutationRules(record.override), + filter: normalizeFilterRules(record.filter), + defaultRaw: normalizeMutationRules(defaultRaw), + }; +} + +function getPayloadRulesPath() { + return ( + process.env.OMNIROUTE_PAYLOAD_RULES_PATH || + process.env.PAYLOAD_RULES_PATH || + path.join(process.cwd(), "config", "payloadRules.json") + ); +} + +function getPayloadRulesReloadIntervalMs() { + const parsed = Number.parseInt(process.env.OMNIROUTE_PAYLOAD_RULES_RELOAD_MS || "", 10); + if (!Number.isFinite(parsed) || parsed < MIN_FILE_CHECK_INTERVAL_MS) { + return DEFAULT_FILE_CHECK_INTERVAL_MS; + } + return parsed; +} + +function clearCachedFileConfig() { + cachedFileConfig = clonePayloadRulesConfig(DEFAULT_PAYLOAD_RULES_CONFIG); + cachedFileMtimeMs = -1; +} + +async function refreshPayloadRulesFileCache(force = false) { + const filePath = getPayloadRulesPath(); + const now = Date.now(); + + if ( + !force && + filePath === cachedFilePath && + now - lastFileCheckAt < getPayloadRulesReloadIntervalMs() + ) { + return; + } + + if (fileLoadPromise) { + await fileLoadPromise; + return; + } + + fileLoadPromise = (async () => { + lastFileCheckAt = now; + cachedFilePath = filePath; + + try { + const stat = await fs.stat(filePath); + if (!force && cachedFileMtimeMs === stat.mtimeMs) { + return; + } + + const content = await fs.readFile(filePath, "utf-8"); + const parsed = JSON.parse(content); + cachedFileConfig = normalizePayloadRulesConfig(parsed); + cachedFileMtimeMs = stat.mtimeMs; + lastFileErrorSignature = ""; + } catch (error) { + if ((error as NodeJS.ErrnoException)?.code === "ENOENT") { + clearCachedFileConfig(); + lastFileErrorSignature = ""; + return; + } + + const message = error instanceof Error ? error.message : String(error); + const errorSignature = `${filePath}:${message}`; + if (errorSignature !== lastFileErrorSignature) { + console.warn(`[PAYLOAD_RULES] Failed to load ${filePath}: ${message}`); + lastFileErrorSignature = errorSignature; + } + } + })(); + + try { + await fileLoadPromise; + } finally { + fileLoadPromise = null; + } +} + +export function setPayloadRulesConfig(config: unknown) { + runtimeOverride = normalizePayloadRulesConfig(config); +} + +export function clearPayloadRulesConfigOverride() { + runtimeOverride = null; +} + +export async function getPayloadRulesConfig(options: { forceRefresh?: boolean } = {}) { + if (runtimeOverride) { + return clonePayloadRulesConfig(runtimeOverride); + } + + await refreshPayloadRulesFileCache(options.forceRefresh === true); + return clonePayloadRulesConfig(cachedFileConfig); +} + +function getPathSegments(pathValue: string) { + return pathValue + .split(".") + .map((segment) => segment.trim()) + .filter(Boolean); +} + +function isIndexSegment(segment: string) { + return /^\d+$/.test(segment); +} + +function getValueAtPath(payload: unknown, pathValue: string) { + const segments = getPathSegments(pathValue); + let cursor: unknown = payload; + + for (const segment of segments) { + if (cursor == null) return undefined; + if (Array.isArray(cursor)) { + if (!isIndexSegment(segment)) return undefined; + cursor = cursor[Number(segment)]; + continue; + } + if (typeof cursor !== "object") return undefined; + cursor = (cursor as JsonRecord)[segment]; + } + + return cursor; +} + +function setValueAtPath(payload: JsonRecord, pathValue: string, value: unknown) { + const segments = getPathSegments(pathValue); + if (segments.length === 0) return; + + let cursor: unknown = payload; + for (let index = 0; index < segments.length - 1; index++) { + const segment = segments[index]; + const nextSegment = segments[index + 1]; + const nextIsIndex = isIndexSegment(nextSegment); + + if (Array.isArray(cursor)) { + const arrayIndex = Number(segment); + if (!Number.isInteger(arrayIndex)) return; + if (cursor[arrayIndex] == null || typeof cursor[arrayIndex] !== "object") { + cursor[arrayIndex] = nextIsIndex ? [] : {}; + } + cursor = cursor[arrayIndex]; + continue; + } + + if (!cursor || typeof cursor !== "object") return; + + const recordCursor = cursor as JsonRecord; + if ( + recordCursor[segment] == null || + typeof recordCursor[segment] !== "object" || + (Array.isArray(recordCursor[segment]) && !nextIsIndex) || + (!Array.isArray(recordCursor[segment]) && nextIsIndex) + ) { + recordCursor[segment] = nextIsIndex ? [] : {}; + } + + cursor = recordCursor[segment]; + } + + const lastSegment = segments.at(-1)!; + if (Array.isArray(cursor)) { + const arrayIndex = Number(lastSegment); + if (!Number.isInteger(arrayIndex)) return; + cursor[arrayIndex] = cloneValue(value); + return; + } + + if (!cursor || typeof cursor !== "object") return; + (cursor as JsonRecord)[lastSegment] = cloneValue(value); +} + +function unsetValueAtPath(payload: JsonRecord, pathValue: string) { + const segments = getPathSegments(pathValue); + if (segments.length === 0) return false; + + let cursor: unknown = payload; + for (let index = 0; index < segments.length - 1; index++) { + const segment = segments[index]; + if (Array.isArray(cursor)) { + if (!isIndexSegment(segment)) return false; + cursor = cursor[Number(segment)]; + continue; + } + if (!cursor || typeof cursor !== "object") return false; + cursor = (cursor as JsonRecord)[segment]; + } + + const lastSegment = segments.at(-1)!; + if (Array.isArray(cursor)) { + if (!isIndexSegment(lastSegment)) return false; + const arrayIndex = Number(lastSegment); + if (arrayIndex < 0 || arrayIndex >= cursor.length) return false; + cursor.splice(arrayIndex, 1); + return true; + } + + if (!cursor || typeof cursor !== "object") return false; + if (!Object.hasOwn(cursor, lastSegment)) return false; + delete (cursor as JsonRecord)[lastSegment]; + return true; +} + +function parseDefaultRawValue(value: unknown) { + if (typeof value !== "string") return value; + + try { + return JSON.parse(value); + } catch { + return value; + } +} + +function matchesProtocol(specProtocol: string | undefined, protocols: string[]) { + if (!specProtocol) return true; + const normalizedProtocol = specProtocol.trim().toLowerCase(); + return protocols.some((protocol) => protocol.trim().toLowerCase() === normalizedProtocol); +} + +function matchesModelSpec(model: string, protocols: string[], spec: PayloadRuleModelSpec) { + return matchesProtocol(spec.protocol, protocols) && wildcardMatch(model, spec.name); +} + +function matchesRule(model: string, protocols: string[], specs: PayloadRuleModelSpec[]) { + return specs.some((spec) => matchesModelSpec(model, protocols, spec)); +} + +function toPayloadRuleProtocols(value: string | string[]) { + const protocols = Array.isArray(value) ? value : [value]; + return [...new Set(protocols.map((protocol) => protocol.trim()).filter(Boolean))]; +} + +export function resolvePayloadRuleProtocols({ + provider, + targetFormat, +}: { + provider?: string | null; + targetFormat?: string | null; +}) { + const protocols = new Set(); + + if (provider) protocols.add(provider); + if (targetFormat) protocols.add(targetFormat); + if (targetFormat === "openai-responses" || targetFormat === "openai-response") { + protocols.add("openai"); + } + if (targetFormat === "gemini-cli" || targetFormat === "antigravity") { + protocols.add("gemini"); + } + + return [...protocols]; +} + +export function applyPayloadRules( + payload: JsonRecord, + model: string, + protocol: string | string[], + rules: PayloadRulesConfig +) { + const normalizedPayload = cloneValue(payload); + const protocols = toPayloadRuleProtocols(protocol); + const applied: AppliedPayloadRule[] = []; + + for (const rule of rules.default) { + if (!matchesRule(model, protocols, rule.models)) continue; + for (const [pathValue, rawValue] of Object.entries(rule.params)) { + if (getValueAtPath(normalizedPayload, pathValue) !== undefined) continue; + setValueAtPath(normalizedPayload, pathValue, rawValue); + applied.push({ type: "default", path: pathValue, value: cloneValue(rawValue) }); + } + } + + for (const rule of rules.defaultRaw) { + if (!matchesRule(model, protocols, rule.models)) continue; + for (const [pathValue, rawValue] of Object.entries(rule.params)) { + if (getValueAtPath(normalizedPayload, pathValue) !== undefined) continue; + const parsedValue = parseDefaultRawValue(rawValue); + setValueAtPath(normalizedPayload, pathValue, parsedValue); + applied.push({ type: "default-raw", path: pathValue, value: cloneValue(parsedValue) }); + } + } + + for (const rule of rules.override) { + if (!matchesRule(model, protocols, rule.models)) continue; + for (const [pathValue, rawValue] of Object.entries(rule.params)) { + setValueAtPath(normalizedPayload, pathValue, rawValue); + applied.push({ type: "override", path: pathValue, value: cloneValue(rawValue) }); + } + } + + for (const rule of rules.filter) { + if (!matchesRule(model, protocols, rule.models)) continue; + for (const pathValue of rule.params) { + if (!unsetValueAtPath(normalizedPayload, pathValue)) continue; + applied.push({ type: "filter", path: pathValue }); + } + } + + return { payload: normalizedPayload, applied }; +} + +export async function applyConfiguredPayloadRules( + payload: JsonRecord, + model: string, + protocol: string | string[] +) { + const rules = await getPayloadRulesConfig(); + return applyPayloadRules(payload, model, protocol, rules); +} + +export function resetPayloadRulesConfigForTests() { + runtimeOverride = null; + cachedFilePath = ""; + cachedFileMtimeMs = -1; + lastFileCheckAt = 0; + fileLoadPromise = null; + lastFileErrorSignature = ""; + cachedFileConfig = clonePayloadRulesConfig(DEFAULT_PAYLOAD_RULES_CONFIG); +} diff --git a/open-sse/utils/stream.ts b/open-sse/utils/stream.ts index c6398112a9..ee4c7d5576 100644 --- a/open-sse/utils/stream.ts +++ b/open-sse/utils/stream.ts @@ -17,6 +17,8 @@ import { formatSSE, unwrapGeminiChunk, } from "./streamHelpers.ts"; +import { calculateCost } from "@/lib/usage/costCalculator"; +import { buildOmniRouteSseMetadataComment } from "@/domain/omnirouteResponseMeta"; import { createStructuredSSECollector, buildStreamSummaryFromEvents, @@ -337,6 +339,7 @@ export function createSSEStream(options: StreamOptions = {}) { // Passthrough: accumulate content and reasoning separately for call log response body let passthroughAccumulatedContent = ""; let passthroughAccumulatedReasoning = ""; + const streamStartedAt = Date.now(); // Guard against duplicate [DONE] events — ensures exactly one per stream let doneSent = false; @@ -478,6 +481,24 @@ export function createSSEStream(options: StreamOptions = {}) { controller.enqueue(encoder.encode(output)); }; + const emitFinalSseMetadata = async ( + controller: TransformStreamDefaultController, + finalUsage: UsageTokenRecord | Record | null | undefined + ) => { + const costUsd = finalUsage ? await calculateCost(provider, model, finalUsage) : 0; + const comment = buildOmniRouteSseMetadataComment({ + provider, + model, + cacheHit: false, + latencyMs: Date.now() - streamStartedAt, + usage: finalUsage, + costUsd, + }); + if (!comment) return; + reqLogger?.appendConvertedChunk?.(comment); + controller.enqueue(encoder.encode(comment)); + }; + return new TransformStream( { start(controller) { @@ -571,11 +592,15 @@ export function createSSEStream(options: StreamOptions = {}) { if (providerPayload) { providerPayloadCollector.push(providerPayload); if ((providerPayload as { done?: unknown }).done === true) { - clientPayloadCollector.push(providerPayload); + continue; } } } + if (trimmed.startsWith("data:") && trimmed.slice(5).trim() === "[DONE]") { + continue; + } + if (trimmed.startsWith("data:") && trimmed.slice(5).trim() !== "[DONE]") { try { let parsed = JSON.parse(trimmed.slice(5).trim()); @@ -835,13 +860,6 @@ export function createSSEStream(options: StreamOptions = {}) { providerPayloadCollector.push(parsed); if (parsed && parsed.done) { - if (!doneSent) { - doneSent = true; - clientPayloadCollector.push({ done: true }); - const output = "data: [DONE]\n\n"; - reqLogger?.appendConvertedChunk?.(output); - controller.enqueue(encoder.encode(output)); - } continue; } @@ -956,7 +974,7 @@ export function createSSEStream(options: StreamOptions = {}) { } }, - flush(controller) { + async flush(controller) { // Clean up idle watchdog timer if (idleTimer) { clearInterval(idleTimer); @@ -1042,6 +1060,14 @@ export function createSSEStream(options: StreamOptions = {}) { status: "200 OK", }).catch(() => {}); } + if (!doneSent) { + await emitFinalSseMetadata(controller, usage); + doneSent = true; + clientPayloadCollector.push({ done: true }); + const doneOutput = "data: [DONE]\n\n"; + reqLogger?.appendConvertedChunk?.(doneOutput); + controller.enqueue(encoder.encode(doneOutput)); + } // Notify caller for call log persistence (include full response body with accumulated content) if (onComplete) { try { @@ -1216,6 +1242,7 @@ export function createSSEStream(options: StreamOptions = {}) { // Send [DONE] (only if not already sent during transform) if (!doneSent) { + await emitFinalSseMetadata(controller, state?.usage as Record | null); doneSent = true; clientPayloadCollector.push({ done: true }); const doneOutput = "data: [DONE]\n\n"; diff --git a/src/app/(dashboard)/dashboard/providers/[id]/page.tsx b/src/app/(dashboard)/dashboard/providers/[id]/page.tsx index 3ddd9a8d99..7d839277e0 100644 --- a/src/app/(dashboard)/dashboard/providers/[id]/page.tsx +++ b/src/app/(dashboard)/dashboard/providers/[id]/page.tsx @@ -37,6 +37,11 @@ import { compatibleProviderSupportsModelImport, getCompatibleFallbackModels, } from "@/lib/providers/managedAvailableModels"; +import { + getModelCatalogSourceLabel, + matchesModelCatalogQuery, + normalizeModelCatalogSource, +} from "@/shared/utils/modelCatalogSearch"; import { useCopyToClipboard } from "@/shared/hooks/useCopyToClipboard"; import { MODEL_COMPAT_PROTOCOL_KEYS, @@ -318,7 +323,7 @@ function anyUpstreamHeadersBadge( } interface ModelRowProps { - model: { id: string; isHidden?: boolean }; + model: { id: string; name?: string; source?: string; isHidden?: boolean }; fullModel: string; copied?: string; onCopy: (text: string, key: string) => void; @@ -336,6 +341,7 @@ interface ModelRowProps { interface PassthroughModelRowProps { modelId: string; fullModel: string; + source?: string; isHidden?: boolean; copied?: string; onCopy: (text: string, key: string) => void; @@ -354,6 +360,7 @@ interface PassthroughModelRowProps { interface PassthroughModelsSectionProps { providerAlias: string; modelAliases: Record; + customModels?: CompatModelRow[]; copied?: string; onCopy: (text: string, key: string) => void; onSetAlias: (modelId: string, alias: string) => Promise; @@ -390,6 +397,7 @@ interface CompatibleModelsSectionProps { providerStorageAlias: string; providerDisplayAlias: string; modelAliases: Record; + customModels?: CompatModelRow[]; fallbackModels?: CompatModelRow[]; allowImport: boolean; description: string; @@ -431,6 +439,34 @@ interface CooldownTimerProps { until: string | number | Date; } +function getModelSourceBadgeClass(source?: string): string { + switch (normalizeModelCatalogSource(source)) { + case "api-sync": + return "border-sky-500/30 bg-sky-500/10 text-sky-300"; + case "custom": + return "border-emerald-500/30 bg-emerald-500/10 text-emerald-300"; + case "fallback": + return "border-amber-500/30 bg-amber-500/10 text-amber-300"; + case "alias": + return "border-violet-500/30 bg-violet-500/10 text-violet-300"; + case "system": + default: + return "border-border bg-sidebar/70 text-text-muted"; + } +} + +function ModelSourceBadge({ source }: { source?: string }) { + return ( + + {getModelCatalogSourceLabel(source)} + + ); +} + interface ConnectionRowConnection { id?: string; name?: string; @@ -492,9 +528,10 @@ interface AddApiKeyModalProps { isCcCompatible?: boolean; onSave: (data: { name: string; - apiKey: string; + apiKey?: string; priority: number; baseUrl?: string; + providerSpecificData?: Record; }) => Promise; onClose: () => void; } @@ -988,13 +1025,29 @@ export default function ProviderDetailPage() { // For Gemini: always use synced API models (empty if no keys added yet) // For other providers: merge registry models with custom/imported models (deduped) const models = useMemo(() => { - if (providerId === "gemini") return syncedAvailableModels; - if (!modelMeta.customModels || modelMeta.customModels.length === 0) return registryModels; - const registryIds = new Set(registryModels.map((m) => m.id)); + if (providerId === "gemini") { + return syncedAvailableModels.map((model: any) => ({ + ...model, + source: model?.source === "api-sync" ? "api-sync" : "api-sync", + })); + } + + const builtInModels = registryModels.map((model) => ({ + ...model, + source: "system", + })); + + if (!modelMeta.customModels || modelMeta.customModels.length === 0) return builtInModels; + + const registryIds = new Set(builtInModels.map((m) => m.id)); const customExtras = modelMeta.customModels .filter((cm: any) => cm.id && !registryIds.has(cm.id)) - .map((cm: any) => ({ id: cm.id, name: cm.name || cm.id })); - return [...registryModels, ...customExtras]; + .map((cm: any) => ({ + id: cm.id, + name: cm.name || cm.id, + source: cm.source === "api-sync" ? "api-sync" : "custom", + })); + return [...builtInModels, ...customExtras]; }, [providerId, registryModels, syncedAvailableModels, modelMeta.customModels]); const providerAlias = getProviderAlias(providerId); const isManagedAvailableModelsProvider = isCompatible || providerId === "openrouter"; @@ -2258,6 +2311,7 @@ export default function ProviderDetailPage() { providerStorageAlias={providerStorageAlias} providerDisplayAlias={providerDisplayAlias} modelAliases={modelAliases} + customModels={modelMeta.customModels} fallbackModels={compatibleFallbackModels} description={description} inputLabel={inputLabel} @@ -2313,6 +2367,7 @@ export default function ProviderDetailPage() { m.id.toLowerCase().includes(modelFilter.toLowerCase())) - : modelsWithVisibility; + const filteredModels = modelsWithVisibility.filter((model) => + matchesModelCatalogQuery(modelFilter, { + modelId: model.id, + modelName: model.name, + source: model.source, + }) + ); const activeCount = modelsWithVisibility.filter((m) => !m.isHidden).length; const hiddenFilteredCount = filteredModels.filter((m) => m.isHidden).length; const visibleFilteredCount = filteredModels.length - hiddenFilteredCount; @@ -2921,6 +2980,29 @@ export default function ProviderDetailPage() {

)} + {providerId === "google-pse-search" && ( +
+ tune +

+ Google Programmable Search requires two values: your API key and the Search Engine + ID (cx) from the Programmable Search Engine dashboard. +

+
+ )} + {providerId === "searxng-search" && ( +
+ dns +

+ SearXNG is self-hosted. Configure the instance base URL here. API key is optional + and can be left blank for public or unauthenticated instances. Local/private URL + validation requires{" "} + + OMNIROUTE_ALLOW_PRIVATE_PROVIDER_URLS=true + + . +

+
+ )} )} @@ -3232,6 +3314,7 @@ function ModelRow({ {fullModel} + + {isGooglePse && ( + setFormData({ ...formData, cx: e.target.value })} + placeholder="012345678901234567890:abc123xyz" + hint="Required. Find this in your Programmable Search Engine overview." + /> + )} {validationResult && ( {validationResult === "success" ? t("valid") : t("invalid")} @@ -5392,6 +5600,20 @@ function AddApiKeyModal({ placeholder="my-app/1.0" hint="Optional override sent upstream as the User-Agent header for this connection" /> + setFormData({ ...formData, routingTags: e.target.value })} + placeholder="fast, cheap, eu-region" + hint="Comma-separated tags matched against request metadata.tags for tag-based routing" + /> + setFormData({ ...formData, excludedModels: e.target.value })} + placeholder="gpt-5*, claude-opus-*, gemini-*-pro*" + hint="Comma-separated wildcard patterns. This connection will never serve matching models." + /> {provider === "bailian-coding-plan" && ( { @@ -5556,6 +5785,8 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec const rawCustomUserAgent = connection.providerSpecificData?.customUserAgent; const existingCustomUserAgent = typeof rawCustomUserAgent === "string" ? rawCustomUserAgent : ""; + const rawCx = connection.providerSpecificData?.cx; + const existingCx = typeof rawCx === "string" ? rawCx : ""; const rawAccountId = connection.providerSpecificData?.accountId; const existingAccountId = typeof rawAccountId === "string" ? rawAccountId : ""; const codexRequestDefaults = getCodexRequestDefaults(connection.providerSpecificData); @@ -5567,10 +5798,16 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec apiKey: "", healthCheckInterval: connection.healthCheckInterval ?? 60, baseUrl: existingBaseUrl || defaultBaseUrl, + cx: existingCx, region: existingRegion || (isVertex ? defaultRegion : ""), apiRegion: (connection.providerSpecificData?.apiRegion as string) || "international", validationModelId: (connection.providerSpecificData?.validationModelId as string) || "", tag: (connection.providerSpecificData?.tag as string) || "", + routingTags: formatRoutingTagsInput(connection.providerSpecificData?.tags), + excludedModels: formatExcludedModelsInput( + connection.providerSpecificData?.excludedModels ?? + connection.providerSpecificData?.excluded_models + ), customUserAgent: existingCustomUserAgent, accountId: existingAccountId, codexReasoningEffort: codexRequestDefaults.reasoningEffort, @@ -5620,7 +5857,7 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec }; const handleValidate = async () => { - if (!connection?.provider || (!isCompatible && !formData.apiKey)) return; + if (!connection?.provider || (!isCompatible && !apiKeyOptional && !formData.apiKey)) return; setValidating(true); setValidationResult(null); try { @@ -5633,6 +5870,7 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec validationModelId: formData.validationModelId || undefined, customUserAgent: formData.customUserAgent.trim() || undefined, baseUrl: formData.baseUrl.trim() || undefined, + cx: formData.cx.trim() || undefined, }), }); const data = await res.json(); @@ -5654,6 +5892,11 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec healthCheckInterval: formData.healthCheckInterval, }; + if (isGooglePse && !formData.cx.trim()) { + setSaveError("Programmable Search Engine ID (cx) is required"); + return; + } + let validatedBaseUrl = null; if (usesBaseUrl) { const checked = normalizeAndValidateHttpBaseUrl(formData.baseUrl, defaultBaseUrl); @@ -5680,6 +5923,7 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec validationModelId: formData.validationModelId || undefined, customUserAgent: formData.customUserAgent.trim() || undefined, baseUrl: formData.baseUrl.trim() || undefined, + cx: formData.cx.trim() || undefined, }), }); const data = await res.json(); @@ -5707,6 +5951,8 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec ...(connection.providerSpecificData || {}), extraApiKeys: extraApiKeys.filter((k) => k.trim().length > 0), tag: formData.tag.trim() || undefined, + tags: parseRoutingTagsInput(formData.routingTags), + excludedModels: parseExcludedModelsInput(formData.excludedModels), customUserAgent: formData.customUserAgent.trim(), }; if (connection.provider === "bailian-coding-plan") { @@ -5719,6 +5965,9 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec if (formData.validationModelId) { updates.providerSpecificData.validationModelId = formData.validationModelId; } + if (isGooglePse) { + updates.providerSpecificData.cx = formData.cx.trim() || undefined; + } if (usesBaseUrl) { updates.providerSpecificData.baseUrl = validatedBaseUrl; } else if (isVertex) { @@ -5733,6 +5982,8 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec updates.providerSpecificData = { ...(connection.providerSpecificData || {}), tag: formData.tag.trim() || undefined, + tags: parseRoutingTagsInput(formData.routingTags), + excludedModels: parseExcludedModelsInput(formData.excludedModels), }; if (isCodex) { updates.providerSpecificData.requestDefaults = { @@ -5779,6 +6030,20 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec placeholder="e.g. personal, work, team-a" hint="Used to group accounts in the provider view" /> + setFormData({ ...formData, routingTags: e.target.value })} + placeholder="fast, cheap, eu-region" + hint="Comma-separated tags matched against request metadata.tags for tag-based routing" + /> + setFormData({ ...formData, excludedModels: e.target.value })} + placeholder="gpt-5*, claude-opus-*, gemini-*-pro*" + hint="Comma-separated wildcard patterns. This connection will be skipped for matching models." + /> {isCodex && (
setFormData({ ...formData, apiKey: e.target.value })} placeholder={isVertex ? "Cole o Service Account JSON aqui" : t("enterNewApiKey")} - hint={t("leaveBlankKeepCurrentApiKey")} + hint={ + isSearxng + ? "Optional. Leave blank to keep the current key or when the instance does not require authentication." + : t("leaveBlankKeepCurrentApiKey") + } className="flex-1" />
+ {isGooglePse && ( + setFormData({ ...formData, cx: e.target.value })} + placeholder="012345678901234567890:abc123xyz" + hint="Required. Find this in your Programmable Search Engine overview." + /> + )} {validationResult && ( {validationResult === "success" ? t("valid") : t("invalid")} @@ -6054,7 +6337,11 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec )}
-
{/* Current Spend */} -
+

{t("todaysSpend")}

{formatCurrency(dailyCost)}

@@ -206,6 +245,24 @@ export default function BudgetTab() { /> )}
+
+

Active Period Spend

+

{formatCurrency(activeCost)}

+ {activeLimit > 0 && ( + + )} +
+

+ Interval: {(budget?.resetInterval || form.resetInterval || "daily").toUpperCase()} +

+

Next reset (UTC): {formatDateTime(budget?.budgetResetAt)}

+
+
{/* Budget Form */} @@ -221,6 +278,15 @@ export default function BudgetTab() { value={form.dailyLimitUsd} onChange={(e) => setForm({ ...form, dailyLimitUsd: e.target.value })} /> + setForm({ ...form, weeklyLimitUsd: e.target.value })} + /> setForm({ ...form, warningThreshold: e.target.value })} />
+
+
+ + +
+ setForm({ ...form, resetTime: e.target.value })} + /> +
+
+

Weekly limit: {weeklyLimit > 0 ? formatCurrency(weeklyLimit) : "Not configured"}

+

Next reset (UTC): {formatDateTime(budget?.budgetResetAt)}

+
diff --git a/src/app/api/models/catalog/route.ts b/src/app/api/models/catalog/route.ts index 83c69d2001..210a79d705 100644 --- a/src/app/api/models/catalog/route.ts +++ b/src/app/api/models/catalog/route.ts @@ -3,6 +3,7 @@ import { PROVIDER_MODELS, PROVIDER_ID_TO_ALIAS } from "@/shared/constants/models import { getAllEmbeddingModels } from "@omniroute/open-sse/config/embeddingRegistry.ts"; import { getAllImageModels } from "@omniroute/open-sse/config/imageRegistry.ts"; import { AI_PROVIDERS, ALIAS_TO_ID } from "@/shared/constants/providers"; +import { hasEligibleConnectionForModel } from "@/domain/connectionModelRules"; /** * GET /api/models/catalog @@ -12,6 +13,41 @@ export async function GET() { try { const connections = await getProviderConnections(); const activeProviders = new Set(connections.map((c) => c.provider)); + const connectionsByProvider = new Map(); + const registerConnectionKey = ( + key: string | null | undefined, + connection: (typeof connections)[number] + ) => { + if (!key) return; + const existing = connectionsByProvider.get(key) || []; + existing.push(connection); + connectionsByProvider.set(key, existing); + }; + for (const connection of connections) { + registerConnectionKey(connection.provider, connection); + registerConnectionKey( + PROVIDER_ID_TO_ALIAS[connection.provider] || connection.provider, + connection + ); + } + const getConnectionsForProvider = (...keys: Array) => { + const seen = new Set(); + const collected: typeof connections = []; + for (const key of keys) { + if (!key) continue; + for (const connection of connectionsByProvider.get(key) || []) { + if (!connection?.id || seen.has(connection.id)) continue; + seen.add(connection.id); + collected.push(connection); + } + } + return collected; + }; + const providerSupportsModel = (providerId: string, modelId: string) => + hasEligibleConnectionForModel( + getConnectionsForProvider(providerId, PROVIDER_ID_TO_ALIAS[providerId]), + modelId + ); const customModelsMap = await getAllCustomModels().catch(() => ({})); const catalog: Record = {}; @@ -28,6 +64,7 @@ export async function GET() { } for (const model of models) { + if (!providerSupportsModel(providerId, model.id)) continue; catalog[alias].models.push({ id: `${alias}/${model.id}`, name: model.name, @@ -39,6 +76,8 @@ export async function GET() { // Embedding models for (const emb of getAllEmbeddingModels()) { + const rawModelId = emb.id.split("/").pop() || emb.id; + if (!providerSupportsModel(emb.provider, rawModelId)) continue; const parts = emb.id.split("/"); const provAlias = parts[0]; if (!catalog[provAlias]) { @@ -58,6 +97,8 @@ export async function GET() { // Image models for (const img of getAllImageModels()) { + const rawModelId = img.id.split("/").pop() || img.id; + if (!providerSupportsModel(img.provider, rawModelId)) continue; const provAlias = img.provider; if (!catalog[provAlias]) { catalog[provAlias] = { @@ -86,6 +127,7 @@ export async function GET() { } for (const model of models as any[]) { + if (!providerSupportsModel(providerId, model.id)) continue; const fullId = `${alias}/${model.id}`; // Skip duplicates if (catalog[alias].models.some((m) => m.id === fullId)) continue; diff --git a/src/app/api/models/route.ts b/src/app/api/models/route.ts index 6c317100b9..f53c8da3ef 100644 --- a/src/app/api/models/route.ts +++ b/src/app/api/models/route.ts @@ -3,6 +3,7 @@ import { getModelAliases, setModelAlias, getProviderConnections } from "@/models import { AI_MODELS, PROVIDER_ID_TO_ALIAS } from "@/shared/constants/models"; import { updateModelAliasSchema } from "@/shared/validation/schemas"; import { isValidationFailure, validateBody } from "@/shared/validation/helpers"; +import { hasEligibleConnectionForModel } from "@/domain/connectionModelRules"; // GET /api/models - Get models with aliases (only from active providers by default) export async function GET(request: Request) { @@ -29,6 +30,45 @@ export async function GET(request: Request) { const alias = PROVIDER_ID_TO_ALIAS[pId]; if (alias) activeProviders.add(alias); } + const connectionsByProvider = new Map(); + const registerConnectionKey = ( + key: string | null | undefined, + connection: (typeof active)[number] + ) => { + if (!key) return; + const existing = connectionsByProvider.get(key) || []; + existing.push(connection); + connectionsByProvider.set(key, existing); + }; + for (const connection of active) { + registerConnectionKey(connection.provider, connection); + registerConnectionKey(PROVIDER_ID_TO_ALIAS[connection.provider], connection); + } + const getConnectionsForProvider = (...keys: Array) => { + const seen = new Set(); + const collected: typeof active = []; + for (const key of keys) { + if (!key) continue; + for (const connection of connectionsByProvider.get(key) || []) { + if (!connection?.id || seen.has(connection.id)) continue; + seen.add(connection.id); + collected.push(connection); + } + } + return collected; + }; + + activeProviders = new Set( + AI_MODELS.flatMap((model: any) => { + const providerKeys = [model.provider, PROVIDER_ID_TO_ALIAS[model.provider]]; + return hasEligibleConnectionForModel( + getConnectionsForProvider(...providerKeys), + model.model + ) + ? providerKeys.filter(Boolean) + : []; + }) + ); } catch { // If DB unavailable, show all models } diff --git a/src/app/api/providers/route.ts b/src/app/api/providers/route.ts index f3cfc76ded..1faeb25ab6 100644 --- a/src/app/api/providers/route.ts +++ b/src/app/api/providers/route.ts @@ -10,7 +10,7 @@ import { getProviderNodeById, isCloudEnabled, } from "@/models"; -import { APIKEY_PROVIDERS } from "@/shared/constants/config"; +import { APIKEY_PROVIDERS, SEARCH_PROVIDERS } from "@/shared/constants/config"; import { isClaudeCodeCompatibleProvider, isOpenAICompatibleProvider, @@ -76,6 +76,7 @@ export async function POST(request: Request) { // Business validation const isValidProvider = APIKEY_PROVIDERS[provider] || + SEARCH_PROVIDERS[provider] || provider === "qoder" || isOpenAICompatibleProvider(provider) || isAnthropicCompatibleProvider(provider); diff --git a/src/app/api/providers/validate/route.ts b/src/app/api/providers/validate/route.ts index 98ec8ebf77..41e12e7f7c 100644 --- a/src/app/api/providers/validate/route.ts +++ b/src/app/api/providers/validate/route.ts @@ -51,6 +51,7 @@ export async function POST(request) { validationModelId, customUserAgent, baseUrl: bodyBaseUrl, + cx, } = validation.data; let providerSpecificData: any = { validationModelId }; @@ -60,6 +61,9 @@ export async function POST(request) { if (bodyBaseUrl) { providerSpecificData.baseUrl = bodyBaseUrl; } + if (cx) { + providerSpecificData.cx = cx; + } if (isOpenAICompatibleProvider(provider) || isAnthropicCompatibleProvider(provider)) { const node: any = await getProviderNodeById(provider); diff --git a/src/app/api/settings/payload-rules/route.ts b/src/app/api/settings/payload-rules/route.ts new file mode 100644 index 0000000000..77d8b7413d --- /dev/null +++ b/src/app/api/settings/payload-rules/route.ts @@ -0,0 +1,50 @@ +import { NextResponse } from "next/server"; +import { updateSettings } from "@/lib/localDb"; +import { + getPayloadRulesConfig, + setPayloadRulesConfig, +} from "@omniroute/open-sse/services/payloadRules.ts"; +import { updatePayloadRulesSchema } from "@/shared/validation/schemas"; +import { isValidationFailure, validateBody } from "@/shared/validation/helpers"; + +export async function GET() { + try { + return NextResponse.json(await getPayloadRulesConfig()); + } catch (error) { + console.error("Error reading payload rules config:", error); + return NextResponse.json({ error: "Failed to read payload rules config" }, { status: 500 }); + } +} + +export async function PUT(request: Request) { + let rawBody: unknown; + try { + rawBody = await request.json(); + } catch { + return NextResponse.json( + { + error: { + message: "Invalid request", + details: [{ field: "body", message: "Invalid JSON body" }], + }, + }, + { status: 400 } + ); + } + + try { + const validation = validateBody(updatePayloadRulesSchema, rawBody); + if (isValidationFailure(validation)) { + return NextResponse.json({ error: validation.error }, { status: 400 }); + } + + setPayloadRulesConfig(validation.data); + const currentConfig = await getPayloadRulesConfig(); + await updateSettings({ payloadRules: currentConfig }); + + return NextResponse.json(currentConfig); + } catch (error) { + console.error("Error updating payload rules config:", error); + return NextResponse.json({ error: "Failed to update payload rules config" }, { status: 500 }); + } +} diff --git a/src/app/api/usage/budget/route.ts b/src/app/api/usage/budget/route.ts index 12b2638a02..ef817c50f9 100644 --- a/src/app/api/usage/budget/route.ts +++ b/src/app/api/usage/budget/route.ts @@ -12,7 +12,24 @@ export async function GET(request) { } const summary = getCostSummary(apiKeyId); const budgetCheck = checkBudget(apiKeyId); - return NextResponse.json({ ...summary, budgetCheck }); + return NextResponse.json({ + ...summary, + budgetCheck, + dailyLimitUsd: summary.dailyLimitUsd, + weeklyLimitUsd: summary.weeklyLimitUsd, + monthlyLimitUsd: summary.monthlyLimitUsd, + warningThreshold: summary.warningThreshold, + resetInterval: summary.resetInterval, + resetTime: summary.resetTime, + budgetResetAt: summary.budgetResetAt, + lastBudgetResetAt: summary.lastBudgetResetAt, + totalCostToday: summary.totalCostToday, + totalCostMonth: summary.totalCostMonth, + totalCostPeriod: summary.totalCostPeriod, + activeLimitUsd: summary.activeLimitUsd, + nextResetAt: summary.nextResetAt, + periodStartAt: summary.periodStartAt, + }); } catch (error) { console.error("Error fetching budget summary:", error); return NextResponse.json({ error: "Failed to fetch budget summary" }, { status: 500 }); @@ -40,10 +57,25 @@ export async function POST(request) { if (isValidationFailure(validation)) { return NextResponse.json({ error: validation.error }, { status: 400 }); } - const { apiKeyId, dailyLimitUsd, monthlyLimitUsd, warningThreshold } = validation.data; + const { + apiKeyId, + dailyLimitUsd, + weeklyLimitUsd, + monthlyLimitUsd, + warningThreshold, + resetInterval, + resetTime, + } = validation.data; - setBudget(apiKeyId, { dailyLimitUsd, monthlyLimitUsd, warningThreshold }); - return NextResponse.json({ success: true, apiKeyId, dailyLimitUsd }); + const budget = setBudget(apiKeyId, { + dailyLimitUsd, + weeklyLimitUsd, + monthlyLimitUsd, + warningThreshold, + resetInterval, + resetTime, + }); + return NextResponse.json({ success: true, apiKeyId, budget }); } catch (error) { console.error("Error setting budget:", error); return NextResponse.json({ error: "Failed to set budget" }, { status: 500 }); diff --git a/src/app/api/v1/models/catalog.ts b/src/app/api/v1/models/catalog.ts index 4f34fca5c7..0552134d32 100644 --- a/src/app/api/v1/models/catalog.ts +++ b/src/app/api/v1/models/catalog.ts @@ -20,6 +20,7 @@ import { getAllMusicModels } from "@omniroute/open-sse/config/musicRegistry.ts"; import { REGISTRY } from "@omniroute/open-sse/config/providerRegistry.ts"; import { getSyncedAvailableModels } from "@/lib/db/models"; import { getCompatibleFallbackModels } from "@/lib/providers/managedAvailableModels"; +import { hasEligibleConnectionForModel } from "@/domain/connectionModelRules"; const FALLBACK_ALIAS_TO_PROVIDER = { ag: "antigravity", @@ -213,12 +214,47 @@ export async function getUnifiedModelsResponse( // Build set of active provider aliases const activeAliases = new Set(); + const connectionsByProvider = new Map(); + const registerConnectionKey = ( + key: string | null | undefined, + connection: (typeof connections)[number] + ) => { + if (!key) return; + const existing = connectionsByProvider.get(key) || []; + existing.push(connection); + connectionsByProvider.set(key, existing); + }; for (const conn of connections) { const alias = providerIdToAlias[conn.provider] || conn.provider; activeAliases.add(alias); activeAliases.add(conn.provider); + registerConnectionKey(alias, conn); + registerConnectionKey(conn.provider, conn); } + const getConnectionsForProvider = (...keys: Array) => { + const seen = new Set(); + const collected: typeof connections = []; + for (const key of keys) { + if (!key) continue; + for (const connection of connectionsByProvider.get(key) || []) { + if (!connection?.id || seen.has(connection.id)) continue; + seen.add(connection.id); + collected.push(connection); + } + } + return collected; + }; + + const providerSupportsModel = (providerKey: string, modelId: string) => { + const providerId = aliasToProviderId[providerKey] || providerKey; + const alias = providerIdToAlias[providerId] || providerKey; + return hasEligibleConnectionForModel( + getConnectionsForProvider(providerKey, providerId, alias), + modelId + ); + }; + // Collect models from active providers (or all if none active) const models = []; const timestamp = Math.floor(Date.now() / 1000); @@ -256,6 +292,7 @@ export async function getUnifiedModelsResponse( const defaultContextLength = registryEntry?.defaultContextLength; for (const model of providerModels) { + if (!providerSupportsModel(canonicalProviderId, model.id)) continue; const aliasId = `${alias}/${model.id}`; if (getModelIsHidden(canonicalProviderId, model.id)) continue; @@ -302,6 +339,7 @@ export async function getUnifiedModelsResponse( try { const syncedModels = await getSyncedAvailableModels("gemini"); for (const sm of syncedModels) { + if (!providerSupportsModel("gemini", sm.id)) continue; const aliasId = `gemini/${sm.id}`; if (getModelIsHidden("gemini", sm.id)) continue; @@ -362,6 +400,8 @@ export async function getUnifiedModelsResponse( // Add embedding models (filtered by active providers) for (const embModel of getAllEmbeddingModels()) { if (!isProviderActive(embModel.provider)) continue; + const rawModelId = embModel.id.split("/").pop() || embModel.id; + if (!providerSupportsModel(embModel.provider, rawModelId)) continue; models.push({ id: embModel.id, object: "model", @@ -375,6 +415,8 @@ export async function getUnifiedModelsResponse( // Add image models (filtered by active providers) for (const imgModel of getAllImageModels()) { if (!isProviderActive(imgModel.provider)) continue; + const rawModelId = imgModel.id.split("/").pop() || imgModel.id; + if (!providerSupportsModel(imgModel.provider, rawModelId)) continue; models.push({ id: imgModel.id, object: "model", @@ -391,6 +433,8 @@ export async function getUnifiedModelsResponse( // Add rerank models (filtered by active providers) for (const rerankModel of getAllRerankModels()) { if (!isProviderActive(rerankModel.provider)) continue; + const rawModelId = rerankModel.id.split("/").pop() || rerankModel.id; + if (!providerSupportsModel(rerankModel.provider, rawModelId)) continue; models.push({ id: rerankModel.id, object: "model", @@ -403,6 +447,8 @@ export async function getUnifiedModelsResponse( // Add audio models (filtered by active providers) for (const audioModel of getAllAudioModels()) { if (!isProviderActive(audioModel.provider)) continue; + const rawModelId = audioModel.id.split("/").pop() || audioModel.id; + if (!providerSupportsModel(audioModel.provider, rawModelId)) continue; models.push({ id: audioModel.id, object: "model", @@ -416,6 +462,8 @@ export async function getUnifiedModelsResponse( // Add moderation models (filtered by active providers) for (const modModel of getAllModerationModels()) { if (!isProviderActive(modModel.provider)) continue; + const rawModelId = modModel.id.split("/").pop() || modModel.id; + if (!providerSupportsModel(modModel.provider, rawModelId)) continue; models.push({ id: modModel.id, object: "model", @@ -428,6 +476,8 @@ export async function getUnifiedModelsResponse( // Add video models (filtered by active providers) for (const videoModel of getAllVideoModels()) { if (!isProviderActive(videoModel.provider)) continue; + const rawModelId = videoModel.id.split("/").pop() || videoModel.id; + if (!providerSupportsModel(videoModel.provider, rawModelId)) continue; models.push({ id: videoModel.id, object: "model", @@ -440,6 +490,8 @@ export async function getUnifiedModelsResponse( // Add music models (filtered by active providers) for (const musicModel of getAllMusicModels()) { if (!isProviderActive(musicModel.provider)) continue; + const rawModelId = musicModel.id.split("/").pop() || musicModel.id; + if (!providerSupportsModel(musicModel.provider, rawModelId)) continue; models.push({ id: musicModel.id, object: "model", @@ -481,6 +533,14 @@ export async function getUnifiedModelsResponse( const modelId = typeof model.id === "string" ? model.id : null; if (!modelId) continue; if (model.isHidden === true) continue; + if ( + !hasEligibleConnectionForModel( + getConnectionsForProvider(alias, canonicalProviderId, providerId, parentProviderType), + modelId + ) + ) { + continue; + } // Skip if already added as built-in const aliasId = `${alias}/${modelId}`; @@ -568,6 +628,7 @@ export async function getUnifiedModelsResponse( const modelId = typeof model.id === "string" ? model.id : null; if (!modelId) continue; if (getModelIsHidden(providerId, modelId)) continue; + if (!hasEligibleConnectionForModel([conn], modelId)) continue; const aliasId = `${alias}/${modelId}`; if (models.some((m) => m.id === aliasId)) continue; diff --git a/src/app/api/v1/search/analytics/route.ts b/src/app/api/v1/search/analytics/route.ts index f11545ef78..68341baeac 100644 --- a/src/app/api/v1/search/analytics/route.ts +++ b/src/app/api/v1/search/analytics/route.ts @@ -6,6 +6,7 @@ */ import { NextResponse } from "next/server"; +import { SEARCH_PROVIDERS } from "@omniroute/open-sse/config/searchRegistry.ts"; import { getDbInstance } from "@/lib/db/core"; import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; @@ -56,19 +57,11 @@ export async function GET(req: Request) { ) .all() as Array<{ provider: string; cnt: number }>; - // Cost per search provider (matching searchRegistry.ts rates) - const COST_PER_QUERY: Record = { - "serper-search": 0.001, - "brave-search": 0.003, - "perplexity-search": 0.005, - "exa-search": 0.01, - "tavily-search": 0.004, - }; - const byProvider: Record = {}; let totalCostUsd = 0; for (const row of provRows) { - const cost = (COST_PER_QUERY[row.provider] ?? 0.001) * row.cnt; + const costPerQuery = SEARCH_PROVIDERS[row.provider]?.costPerQuery ?? 0; + const cost = costPerQuery * row.cnt; byProvider[row.provider] = { count: row.cnt, costUsd: cost }; totalCostUsd += cost; } diff --git a/src/app/api/v1/search/route.ts b/src/app/api/v1/search/route.ts index 8a1ecc9190..f8820e11e2 100644 --- a/src/app/api/v1/search/route.ts +++ b/src/app/api/v1/search/route.ts @@ -5,6 +5,7 @@ import { getAllSearchProviders, getSearchProvider, selectProvider, + supportsSearchType, SEARCH_PROVIDERS, SEARCH_CREDENTIAL_FALLBACKS, } from "@omniroute/open-sse/config/searchRegistry.ts"; @@ -111,7 +112,20 @@ export async function POST(request: Request) { if (policy.rejection) return policy.rejection; // Resolve provider and credentials - let providerConfig = selectProvider(body.provider); + if (body.provider) { + const explicitProvider = getSearchProvider(body.provider); + if (!explicitProvider) { + return errorResponse(HTTP_STATUS.BAD_REQUEST, `Unknown search provider: ${body.provider}`); + } + if (!supportsSearchType(explicitProvider, body.search_type)) { + return errorResponse( + HTTP_STATUS.BAD_REQUEST, + `Search provider ${body.provider} does not support search_type: ${body.search_type}` + ); + } + } + + let providerConfig = selectProvider(body.provider, body.search_type); if (!providerConfig) { return errorResponse( HTTP_STATUS.BAD_REQUEST, @@ -126,10 +140,20 @@ export async function POST(request: Request) { if (body.provider) { // Explicit provider — single credential lookup (with fallback) credentials = await resolveSearchCredentials(providerConfig.id); + if ( + !credentials && + providerConfig.authType === "none" && + typeof body.provider_options?.baseUrl === "string" && + body.provider_options.baseUrl.trim().length > 0 + ) { + credentials = { providerSpecificData: { baseUrl: body.provider_options.baseUrl.trim() } }; + } if (!credentials) { return errorResponse( HTTP_STATUS.BAD_REQUEST, - `No credentials configured for search provider: ${providerConfig.id}. Add an API key for "${providerConfig.id}" in the dashboard.` + providerConfig.authType === "none" + ? `Search provider ${providerConfig.id} is not configured. Set its base URL in the dashboard or pass provider_options.baseUrl.` + : `No credentials configured for search provider: ${providerConfig.id}. Add an API key for "${providerConfig.id}" in the dashboard.` ); } } else { @@ -139,6 +163,7 @@ export async function POST(request: Request) { if (!credentials) { // Sort by cost to find cheapest with credentials const sortedIds = Object.values(SEARCH_PROVIDERS) + .filter((provider) => supportsSearchType(provider, body.search_type)) .sort((a, b) => a.costPerQuery - b.costPerQuery) .map((p) => p.id); @@ -163,6 +188,7 @@ export async function POST(request: Request) { // Find alternate for failover — must bind credentials to the matched provider const otherIds = Object.values(SEARCH_PROVIDERS) + .filter((provider) => supportsSearchType(provider, body.search_type)) .sort((a, b) => a.costPerQuery - b.costPerQuery) .map((p) => p.id) .filter((id) => id !== providerConfig.id); diff --git a/src/domain/connectionModelRules.ts b/src/domain/connectionModelRules.ts new file mode 100644 index 0000000000..316ade72d8 --- /dev/null +++ b/src/domain/connectionModelRules.ts @@ -0,0 +1,82 @@ +import { wildcardMatch } from "@omniroute/open-sse/services/wildcardRouter.ts"; + +type JsonRecord = Record; + +interface ConnectionLike { + providerSpecificData?: unknown; +} + +function asRecord(value: unknown): JsonRecord { + return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; +} + +function normalizePattern(value: unknown): string | null { + if (typeof value !== "string") return null; + const normalized = value.trim(); + if (!normalized || normalized === "**") return null; + return normalized; +} + +function toPatternList(value: unknown): string[] { + if (Array.isArray(value)) { + return value.map(normalizePattern).filter((pattern): pattern is string => Boolean(pattern)); + } + if (typeof value === "string") { + return value + .split(",") + .map(normalizePattern) + .filter((pattern): pattern is string => Boolean(pattern)); + } + return []; +} + +function uniquePatterns(patterns: string[]): string[] { + return Array.from(new Set(patterns)); +} + +function getModelMatchCandidates(modelId: string): string[] { + const normalized = modelId.trim(); + if (!normalized) return []; + + const withoutExtendedContext = normalized.endsWith("[1m]") ? normalized.slice(0, -4) : normalized; + const lastSlashIndex = withoutExtendedContext.lastIndexOf("/"); + const rawModel = + lastSlashIndex >= 0 ? withoutExtendedContext.slice(lastSlashIndex + 1) : withoutExtendedContext; + + return Array.from(new Set([normalized, withoutExtendedContext, rawModel].filter(Boolean))); +} + +export function normalizeExcludedModelPatterns(value: unknown): string[] { + return uniquePatterns(toPatternList(value)); +} + +export function getConnectionExcludedModels(providerSpecificData: unknown): string[] { + const data = asRecord(providerSpecificData); + return normalizeExcludedModelPatterns(data.excludedModels ?? data.excluded_models); +} + +export function isModelExcludedByConnection( + modelId: unknown, + providerSpecificData: unknown +): boolean { + if (typeof modelId !== "string" || modelId.trim().length === 0) return false; + + const candidates = getModelMatchCandidates(modelId); + const excludedModels = getConnectionExcludedModels(providerSpecificData); + if (candidates.length === 0 || excludedModels.length === 0) return false; + + return excludedModels.some((pattern) => + candidates.some((candidate) => wildcardMatch(candidate, pattern)) + ); +} + +export function hasEligibleConnectionForModel( + connections: ConnectionLike[] | null | undefined, + modelId: unknown +): boolean { + if (!Array.isArray(connections) || connections.length === 0) return false; + + return connections.some( + (connection) => !isModelExcludedByConnection(modelId, connection?.providerSpecificData) + ); +} diff --git a/src/domain/costRules.ts b/src/domain/costRules.ts index 239a11fd83..71a6db96d4 100644 --- a/src/domain/costRules.ts +++ b/src/domain/costRules.ts @@ -2,27 +2,56 @@ * Cost Rules — Domain Layer (T-19) * * Business rules for cost management: budget thresholds, - * quota checking, and cost summaries per API key. + * scheduled reset windows, and cost summaries per API key. * - * State is persisted in SQLite via domainState.js. + * State is persisted in SQLite via domainState.ts. * * @module domain/costRules */ import { - saveBudget, - loadBudget, - saveCostEntry, - loadCostEntries, deleteAllCostData, deleteBudget as dbDeleteBudget, deleteCostEntries, + loadAllBudgets, + loadBudget, + loadCostEntries, + loadCostEntriesInRange, + saveBudget, + saveBudgetResetLog, } from "../lib/db/domainState"; +import { + discardSpendBatchEntries, + resetSpendBatchWriterForTests, + spendBatchWriter, +} from "@/lib/spend/batchWriter"; + +export type BudgetResetInterval = "daily" | "weekly" | "monthly"; interface BudgetConfig { - dailyLimitUsd: number; + dailyLimitUsd?: number; + weeklyLimitUsd?: number; monthlyLimitUsd?: number; warningThreshold?: number; + resetInterval?: BudgetResetInterval; + resetTime?: string; + budgetResetAt?: number | null; + lastBudgetResetAt?: number | null; + warningEmittedAt?: number | null; + warningPeriodStart?: number | null; +} + +interface NormalizedBudgetConfig { + dailyLimitUsd: number; + weeklyLimitUsd: number; + monthlyLimitUsd: number; + warningThreshold: number; + resetInterval: BudgetResetInterval; + resetTime: string; + budgetResetAt: number | null; + lastBudgetResetAt: number | null; + warningEmittedAt: number | null; + warningPeriodStart: number | null; } interface CostEntry { @@ -30,39 +59,285 @@ interface CostEntry { timestamp: number; } +interface BudgetWindow { + periodStartAt: number; + nextResetAt: number; +} + +interface SyncBudgetScheduleOptions { + logReset?: boolean; + persist?: boolean; +} + +interface BudgetSummary { + dailyTotal: number; + monthlyTotal: number; + totalEntries: number; + budget: NormalizedBudgetConfig | null; + totalCostToday: number; + totalCostMonth: number; + totalCostPeriod: number; + activeLimitUsd: number; + resetInterval: BudgetResetInterval | null; + resetTime: string | null; + budgetResetAt: number | null; + lastBudgetResetAt: number | null; + periodStartAt: number | null; + nextResetAt: number | null; + dailyLimitUsd: number; + weeklyLimitUsd: number; + monthlyLimitUsd: number; + warningThreshold: number | null; +} + +const VALID_RESET_INTERVALS = new Set(["daily", "weekly", "monthly"]); +const RESET_TIME_REGEX = /^(\d{2}):(\d{2})$/; + +/** @type {Map} In-memory cache for budgets */ +const budgets = new Map(); + +function toNumber(value: unknown, fallback = 0): number { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim().length > 0) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : fallback; + } + return fallback; +} + function toCostEntries(value: unknown): CostEntry[] { if (!Array.isArray(value)) return []; const entries: CostEntry[] = []; for (const item of value) { if (!item || typeof item !== "object" || Array.isArray(item)) continue; const record = item as Record; - const cost = typeof record.cost === "number" ? record.cost : Number(record.cost ?? 0); - const timestamp = - typeof record.timestamp === "number" ? record.timestamp : Number(record.timestamp ?? 0); + const cost = toNumber(record.cost, Number.NaN); + const timestamp = toNumber(record.timestamp, Number.NaN); if (!Number.isFinite(cost) || !Number.isFinite(timestamp)) continue; entries.push({ cost, timestamp }); } return entries; } -/** - * @typedef {Object} BudgetConfig - * @property {number} dailyLimitUsd - Max daily spend in USD - * @property {number} [monthlyLimitUsd] - Max monthly spend in USD - * @property {number} [warningThreshold=0.8] - Alert when usage reaches this fraction - */ +function sumEntries(entries: CostEntry[]): number { + return entries.reduce((sum, entry) => sum + entry.cost, 0); +} -/** - * @typedef {Object} CostEntry - * @property {number} cost - Cost in USD - * @property {number} timestamp - Unix timestamp - */ +function normalizeResetInterval(value: unknown): BudgetResetInterval { + if (typeof value === "string") { + const normalized = value.trim().toLowerCase() as BudgetResetInterval; + if (VALID_RESET_INTERVALS.has(normalized)) { + return normalized; + } + } + return "daily"; +} -/** @type {Map} In-memory cache for budgets */ -const budgets = new Map(); +function normalizeResetTime(value: unknown): string { + if (typeof value === "string") { + const match = value.trim().match(RESET_TIME_REGEX); + if (match) { + const hours = Math.min(Math.max(parseInt(match[1], 10), 0), 23); + const minutes = Math.min(Math.max(parseInt(match[2], 10), 0), 59); + return `${String(hours).padStart(2, "0")}:${String(minutes).padStart(2, "0")}`; + } + } + return "00:00"; +} -/** @type {boolean} */ -let _budgetsLoaded = false; +function normalizeTimestamp(value: unknown): number | null { + const numeric = toNumber(value, Number.NaN); + return Number.isFinite(numeric) && numeric > 0 ? numeric : null; +} + +function normalizeBudgetConfig(config: BudgetConfig): NormalizedBudgetConfig { + return { + dailyLimitUsd: Math.max(0, toNumber(config.dailyLimitUsd)), + weeklyLimitUsd: Math.max(0, toNumber(config.weeklyLimitUsd)), + monthlyLimitUsd: Math.max(0, toNumber(config.monthlyLimitUsd)), + warningThreshold: Math.min(Math.max(toNumber(config.warningThreshold, 0.8), 0), 1), + resetInterval: normalizeResetInterval(config.resetInterval), + resetTime: normalizeResetTime(config.resetTime), + budgetResetAt: normalizeTimestamp(config.budgetResetAt), + lastBudgetResetAt: normalizeTimestamp(config.lastBudgetResetAt), + warningEmittedAt: normalizeTimestamp(config.warningEmittedAt), + warningPeriodStart: normalizeTimestamp(config.warningPeriodStart), + }; +} + +function getResetTimeParts(resetTime: string): [number, number] { + const match = resetTime.match(RESET_TIME_REGEX); + if (!match) return [0, 0]; + return [parseInt(match[1], 10), parseInt(match[2], 10)]; +} + +function getUtcDateMs(year: number, month: number, day: number, hours: number, minutes: number) { + return Date.UTC(year, month, day, hours, minutes, 0, 0); +} + +export function getBudgetWindow( + resetInterval: BudgetResetInterval, + resetTime = "00:00", + now = Date.now() +): BudgetWindow { + const current = new Date(now); + const [hours, minutes] = getResetTimeParts(normalizeResetTime(resetTime)); + const year = current.getUTCFullYear(); + const month = current.getUTCMonth(); + const day = current.getUTCDate(); + + if (resetInterval === "weekly") { + const daysSinceMonday = (current.getUTCDay() + 6) % 7; + const thisWeekReset = getUtcDateMs(year, month, day - daysSinceMonday, hours, minutes); + return now >= thisWeekReset + ? { + periodStartAt: thisWeekReset, + nextResetAt: getUtcDateMs(year, month, day - daysSinceMonday + 7, hours, minutes), + } + : { + periodStartAt: getUtcDateMs(year, month, day - daysSinceMonday - 7, hours, minutes), + nextResetAt: thisWeekReset, + }; + } + + if (resetInterval === "monthly") { + const thisMonthReset = getUtcDateMs(year, month, 1, hours, minutes); + return now >= thisMonthReset + ? { + periodStartAt: thisMonthReset, + nextResetAt: getUtcDateMs(year, month + 1, 1, hours, minutes), + } + : { + periodStartAt: getUtcDateMs(year, month - 1, 1, hours, minutes), + nextResetAt: thisMonthReset, + }; + } + + const todayReset = getUtcDateMs(year, month, day, hours, minutes); + return now >= todayReset + ? { + periodStartAt: todayReset, + nextResetAt: getUtcDateMs(year, month, day + 1, hours, minutes), + } + : { + periodStartAt: getUtcDateMs(year, month, day - 1, hours, minutes), + nextResetAt: todayReset, + }; +} + +function getActiveBudgetLimit(budget: NormalizedBudgetConfig): number { + if (budget.resetInterval === "monthly") { + return budget.monthlyLimitUsd > 0 ? budget.monthlyLimitUsd : budget.dailyLimitUsd; + } + if (budget.resetInterval === "weekly") { + return budget.weeklyLimitUsd > 0 ? budget.weeklyLimitUsd : budget.dailyLimitUsd; + } + return budget.dailyLimitUsd; +} + +function getBudgetWindowTotal(apiKeyId: string, periodStartAt: number): number { + try { + return ( + sumEntries(toCostEntries(loadCostEntries(apiKeyId, periodStartAt))) + + spendBatchWriter.getPendingCostTotal(apiKeyId, periodStartAt) + ); + } catch { + return 0; + } +} + +function getBudgetWindowRangeTotal( + apiKeyId: string, + periodStartAt: number, + periodEndAt: number +): number { + try { + return ( + sumEntries(toCostEntries(loadCostEntriesInRange(apiKeyId, periodStartAt, periodEndAt))) + + spendBatchWriter.getPendingCostTotal(apiKeyId, periodStartAt, periodEndAt) + ); + } catch { + return 0; + } +} + +function emitBudgetWarning( + apiKeyId: string, + budget: NormalizedBudgetConfig, + projectedTotal: number, + activeLimitUsd: number, + nextResetAt: number +) { + const percentage = + activeLimitUsd > 0 ? ((projectedTotal / activeLimitUsd) * 100).toFixed(1) : "0.0"; + console.warn( + `[BudgetWarning] ${apiKeyId} reached ${percentage}% of ${budget.resetInterval} budget ($${projectedTotal.toFixed(4)} / $${activeLimitUsd.toFixed(2)}) — next reset ${new Date(nextResetAt).toISOString()}` + ); +} + +function syncBudgetSchedule( + apiKeyId: string, + config: BudgetConfig, + now = Date.now(), + options: SyncBudgetScheduleOptions = {} +): NormalizedBudgetConfig { + const normalized = normalizeBudgetConfig(config); + const window = getBudgetWindow(normalized.resetInterval, normalized.resetTime, now); + const resetRolled = + normalized.lastBudgetResetAt !== null && window.periodStartAt > normalized.lastBudgetResetAt; + + if (resetRolled && options.logReset !== false) { + const previousSpend = getBudgetWindowRangeTotal( + apiKeyId, + normalized.lastBudgetResetAt as number, + window.periodStartAt + ); + try { + saveBudgetResetLog({ + apiKeyId, + resetInterval: normalized.resetInterval, + previousSpend, + resetAt: window.periodStartAt, + nextResetAt: window.nextResetAt, + periodStart: window.periodStartAt, + periodEnd: window.nextResetAt, + }); + } catch { + // Non-fatal: budget logic still proceeds even if logging fails. + } + } + + const updated: NormalizedBudgetConfig = { + ...normalized, + budgetResetAt: window.nextResetAt, + lastBudgetResetAt: window.periodStartAt, + warningEmittedAt: resetRolled ? null : normalized.warningEmittedAt, + warningPeriodStart: resetRolled ? null : normalized.warningPeriodStart, + }; + + const changed = + normalized.budgetResetAt !== updated.budgetResetAt || + normalized.lastBudgetResetAt !== updated.lastBudgetResetAt || + normalized.warningEmittedAt !== updated.warningEmittedAt || + normalized.warningPeriodStart !== updated.warningPeriodStart || + normalized.dailyLimitUsd !== updated.dailyLimitUsd || + normalized.weeklyLimitUsd !== updated.weeklyLimitUsd || + normalized.monthlyLimitUsd !== updated.monthlyLimitUsd || + normalized.warningThreshold !== updated.warningThreshold || + normalized.resetInterval !== updated.resetInterval || + normalized.resetTime !== updated.resetTime; + + if (changed && options.persist !== false) { + try { + saveBudget(apiKeyId, updated); + } catch { + // Non-critical: in-memory cache still works. + } + } + + budgets.set(apiKeyId, updated); + return updated; +} /** * Set budget for an API key. @@ -71,43 +346,49 @@ let _budgetsLoaded = false; * @param {BudgetConfig} config */ export function setBudget(apiKeyId: string, config: BudgetConfig) { - const normalized = { - dailyLimitUsd: config.dailyLimitUsd, - monthlyLimitUsd: config.monthlyLimitUsd || 0, - warningThreshold: config.warningThreshold ?? 0.8, - }; - budgets.set(apiKeyId, normalized); - try { - saveBudget(apiKeyId, normalized); - } catch { - // Non-critical: in-memory still works - } + return syncBudgetSchedule(apiKeyId, config, Date.now(), { logReset: false, persist: true }); } /** * Get budget config for an API key. * * @param {string} apiKeyId - * @returns {BudgetConfig | null} + * @returns {NormalizedBudgetConfig | null} */ -export function getBudget(apiKeyId: string): BudgetConfig | null { - // Check in-memory cache first - if (budgets.has(apiKeyId)) { - return budgets.get(apiKeyId); +export function getBudget(apiKeyId: string): NormalizedBudgetConfig | null { + const cached = budgets.get(apiKeyId); + if (cached) { + return syncBudgetSchedule(apiKeyId, cached); } - // Try loading from DB + try { const fromDb = loadBudget(apiKeyId) as BudgetConfig | null; if (fromDb) { - budgets.set(apiKeyId, fromDb); - return fromDb; + return syncBudgetSchedule(apiKeyId, fromDb); } } catch { - // DB may not be ready + // DB may not be ready. } + return null; } +/** + * Delete budget config and recorded spend for an API key. + * + * @param {string} apiKeyId + */ +export function deleteBudget(apiKeyId: string) { + budgets.delete(apiKeyId); + discardSpendBatchEntries(apiKeyId); + try { + dbDeleteBudget(apiKeyId); + deleteCostEntries(apiKeyId); + } catch { + // Non-critical. + } +} + /** * Record a cost for an API key. * @@ -115,46 +396,122 @@ export function getBudget(apiKeyId: string): BudgetConfig | null { * @param {number} cost - Cost in USD */ export function recordCost(apiKeyId: string, cost: number): void { - const timestamp = Date.now(); try { - saveCostEntry(apiKeyId, cost, timestamp); + spendBatchWriter.increment(apiKeyId, cost, Date.now()); } catch { - // Non-critical + // Non-critical. } } +/** + * Sync all budgets against the current clock so overdue resets get persisted. + */ +export function syncAllBudgetSchedules(now = Date.now()) { + let processed = 0; + let resetCount = 0; + + try { + const allBudgets = loadAllBudgets(); + for (const [apiKeyId, budget] of Object.entries(allBudgets)) { + processed += 1; + const synced = syncBudgetSchedule(apiKeyId, budget, now, { logReset: true, persist: true }); + if (budget.lastBudgetResetAt !== synced.lastBudgetResetAt) { + resetCount += 1; + } + } + } catch { + // Non-critical. + } + + return { processed, resetCount }; +} + /** * Check if an API key has remaining budget. * * @param {string} apiKeyId * @param {number} [additionalCost=0] - Projected cost to check - * @returns {{ allowed: boolean, reason?: string, dailyUsed: number, dailyLimit: number, warningReached: boolean }} + * @returns {{ allowed: boolean, reason?: string, dailyUsed: number, dailyLimit: number, warningReached: boolean, remaining: number, periodUsed: number, activeLimitUsd: number, resetInterval: BudgetResetInterval | null, resetTime: string | null, budgetResetAt: number | null, lastBudgetResetAt: number | null, periodStartAt: number | null }} */ export function checkBudget(apiKeyId: string, additionalCost = 0) { const budget = getBudget(apiKeyId); if (!budget) { - return { allowed: true, dailyUsed: 0, dailyLimit: 0, warningReached: false }; + return { + allowed: true, + dailyUsed: 0, + dailyLimit: 0, + warningReached: false, + remaining: 0, + periodUsed: 0, + activeLimitUsd: 0, + resetInterval: null, + resetTime: null, + budgetResetAt: null, + lastBudgetResetAt: null, + periodStartAt: null, + }; } - const dailyUsed = getDailyTotal(apiKeyId); - const projectedTotal = dailyUsed + additionalCost; - const warningReached = projectedTotal >= budget.dailyLimitUsd * budget.warningThreshold; + const window = getBudgetWindow(budget.resetInterval, budget.resetTime); + const periodUsed = getBudgetWindowTotal(apiKeyId, window.periodStartAt); + const projectedTotal = periodUsed + additionalCost; + const activeLimitUsd = getActiveBudgetLimit(budget); + const warningReached = + activeLimitUsd > 0 && projectedTotal >= activeLimitUsd * budget.warningThreshold; + const remaining = Math.max(activeLimitUsd - projectedTotal, 0); - if (projectedTotal > budget.dailyLimitUsd) { + if (warningReached && budget.warningPeriodStart !== window.periodStartAt) { + const updatedBudget = { + ...budget, + warningEmittedAt: Date.now(), + warningPeriodStart: window.periodStartAt, + }; + budgets.set(apiKeyId, updatedBudget); + try { + saveBudget(apiKeyId, updatedBudget); + emitBudgetWarning( + apiKeyId, + updatedBudget, + projectedTotal, + activeLimitUsd, + window.nextResetAt + ); + } catch { + // Non-critical. + } + } + + if (activeLimitUsd > 0 && projectedTotal > activeLimitUsd) { return { allowed: false, - reason: `Daily budget exceeded: $${projectedTotal.toFixed(4)} / $${budget.dailyLimitUsd.toFixed(2)}`, - dailyUsed, - dailyLimit: budget.dailyLimitUsd, + reason: `${budget.resetInterval[0].toUpperCase()}${budget.resetInterval.slice(1)} budget exceeded: $${projectedTotal.toFixed(4)} / $${activeLimitUsd.toFixed(2)}`, + dailyUsed: periodUsed, + dailyLimit: activeLimitUsd, warningReached: true, + remaining, + periodUsed, + activeLimitUsd, + resetInterval: budget.resetInterval, + resetTime: budget.resetTime, + budgetResetAt: window.nextResetAt, + lastBudgetResetAt: window.periodStartAt, + periodStartAt: window.periodStartAt, }; } return { allowed: true, - dailyUsed, - dailyLimit: budget.dailyLimitUsd, + dailyUsed: periodUsed, + dailyLimit: activeLimitUsd, warningReached, + remaining, + periodUsed, + activeLimitUsd, + resetInterval: budget.resetInterval, + resetTime: budget.resetTime, + budgetResetAt: window.nextResetAt, + lastBudgetResetAt: window.periodStartAt, + periodStartAt: window.periodStartAt, }; } @@ -167,11 +524,12 @@ export function checkBudget(apiKeyId: string, additionalCost = 0) { export function getDailyTotal(apiKeyId: string): number { const todayStart = new Date(); todayStart.setHours(0, 0, 0, 0); - const startMs = todayStart.getTime(); try { - const entries = toCostEntries(loadCostEntries(apiKeyId, startMs)); - return entries.reduce((sum, e) => sum + e.cost, 0); + return ( + sumEntries(toCostEntries(loadCostEntries(apiKeyId, todayStart.getTime()))) + + spendBatchWriter.getPendingCostTotal(apiKeyId, todayStart.getTime()) + ); } catch { return 0; } @@ -181,35 +539,79 @@ export function getDailyTotal(apiKeyId: string): number { * Get cost summary for an API key. * * @param {string} apiKeyId - * @returns {{ dailyTotal: number, monthlyTotal: number, totalEntries: number, budget: BudgetConfig | null }} + * @returns {BudgetSummary} */ -export function getCostSummary(apiKeyId: string) { +export function getCostSummary(apiKeyId: string): BudgetSummary { const now = new Date(); - const todayStart = new Date(now); todayStart.setHours(0, 0, 0, 0); - const monthStart = new Date(now.getFullYear(), now.getMonth(), 1); - try { - const dailyEntries = toCostEntries(loadCostEntries(apiKeyId, todayStart.getTime())); - const monthlyEntries = toCostEntries(loadCostEntries(apiKeyId, monthStart.getTime())); + const budget = getBudget(apiKeyId); + const window = budget ? getBudgetWindow(budget.resetInterval, budget.resetTime) : null; - const dailyTotal = dailyEntries.reduce((sum, e) => sum + e.cost, 0); - const monthlyTotal = monthlyEntries.reduce((sum, e) => sum + e.cost, 0); + try { + const dailyEntries = [ + ...toCostEntries(loadCostEntries(apiKeyId, todayStart.getTime())), + ...spendBatchWriter.getBufferedEntries(apiKeyId, todayStart.getTime()), + ]; + const monthlyEntries = [ + ...toCostEntries(loadCostEntries(apiKeyId, monthStart.getTime())), + ...spendBatchWriter.getBufferedEntries(apiKeyId, monthStart.getTime()), + ]; + const periodEntries = + window !== null + ? [ + ...toCostEntries(loadCostEntries(apiKeyId, window.periodStartAt)), + ...spendBatchWriter.getBufferedEntries(apiKeyId, window.periodStartAt), + ] + : []; + + const dailyTotal = sumEntries(dailyEntries); + const monthlyTotal = sumEntries(monthlyEntries); + const periodTotal = sumEntries(periodEntries); + const activeLimitUsd = budget ? getActiveBudgetLimit(budget) : 0; return { dailyTotal, monthlyTotal, totalEntries: monthlyEntries.length, - budget: getBudget(apiKeyId), + budget, + totalCostToday: dailyTotal, + totalCostMonth: monthlyTotal, + totalCostPeriod: periodTotal, + activeLimitUsd, + resetInterval: budget?.resetInterval ?? null, + resetTime: budget?.resetTime ?? null, + budgetResetAt: window?.nextResetAt ?? budget?.budgetResetAt ?? null, + lastBudgetResetAt: window?.periodStartAt ?? budget?.lastBudgetResetAt ?? null, + periodStartAt: window?.periodStartAt ?? null, + nextResetAt: window?.nextResetAt ?? null, + dailyLimitUsd: budget?.dailyLimitUsd ?? 0, + weeklyLimitUsd: budget?.weeklyLimitUsd ?? 0, + monthlyLimitUsd: budget?.monthlyLimitUsd ?? 0, + warningThreshold: budget?.warningThreshold ?? null, }; } catch { return { dailyTotal: 0, monthlyTotal: 0, totalEntries: 0, - budget: getBudget(apiKeyId), + budget, + totalCostToday: 0, + totalCostMonth: 0, + totalCostPeriod: 0, + activeLimitUsd: budget ? getActiveBudgetLimit(budget) : 0, + resetInterval: budget?.resetInterval ?? null, + resetTime: budget?.resetTime ?? null, + budgetResetAt: budget?.budgetResetAt ?? null, + lastBudgetResetAt: budget?.lastBudgetResetAt ?? null, + periodStartAt: budget?.lastBudgetResetAt ?? null, + nextResetAt: budget?.budgetResetAt ?? null, + dailyLimitUsd: budget?.dailyLimitUsd ?? 0, + weeklyLimitUsd: budget?.weeklyLimitUsd ?? 0, + monthlyLimitUsd: budget?.monthlyLimitUsd ?? 0, + warningThreshold: budget?.warningThreshold ?? null, }; } } @@ -219,10 +621,10 @@ export function getCostSummary(apiKeyId: string) { */ export function resetCostData() { budgets.clear(); - _budgetsLoaded = false; + resetSpendBatchWriterForTests(); try { deleteAllCostData(); } catch { - // Non-critical + // Non-critical. } } diff --git a/src/domain/omnirouteResponseMeta.ts b/src/domain/omnirouteResponseMeta.ts new file mode 100644 index 0000000000..287ed20542 --- /dev/null +++ b/src/domain/omnirouteResponseMeta.ts @@ -0,0 +1,91 @@ +import { getProviderAlias } from "@/shared/constants/providers"; +import { OMNIROUTE_RESPONSE_HEADERS } from "@/shared/constants/headers"; + +type UsageLike = Record | null | undefined; + +function toFiniteNumber(value: unknown): number { + if (typeof value === "number" && Number.isFinite(value)) return value; + if (typeof value === "string" && value.trim().length > 0) { + const parsed = Number(value); + return Number.isFinite(parsed) ? parsed : 0; + } + return 0; +} + +function toNonNegativeInteger(value: unknown): number { + return Math.max(0, Math.round(toFiniteNumber(value))); +} + +export function getOmniRouteTokenCounts(usage: UsageLike): { input: number; output: number } { + if (!usage || typeof usage !== "object") { + return { input: 0, output: 0 }; + } + + return { + input: toNonNegativeInteger( + usage.input ?? + usage.prompt_tokens ?? + usage.input_tokens ?? + usage.promptTokens ?? + usage.inputTokens + ), + output: toNonNegativeInteger( + usage.output ?? + usage.completion_tokens ?? + usage.output_tokens ?? + usage.completionTokens ?? + usage.outputTokens + ), + }; +} + +export function formatOmniRouteCost(costUsd: unknown): string { + const normalized = toFiniteNumber(costUsd); + return normalized > 0 ? normalized.toFixed(10) : "0.0000000000"; +} + +export function buildOmniRouteResponseMetaHeaders({ + cacheHit = false, + costUsd = 0, + latencyMs = 0, + model = null, + provider = null, + usage = null, +}: { + cacheHit?: boolean; + costUsd?: unknown; + latencyMs?: unknown; + model?: string | null; + provider?: string | null; + usage?: UsageLike; +}): Record { + const tokens = getOmniRouteTokenCounts(usage); + const headers: Record = { + [OMNIROUTE_RESPONSE_HEADERS.cacheHit]: String(cacheHit), + [OMNIROUTE_RESPONSE_HEADERS.latencyMs]: String(toNonNegativeInteger(latencyMs)), + [OMNIROUTE_RESPONSE_HEADERS.responseCost]: formatOmniRouteCost(costUsd), + [OMNIROUTE_RESPONSE_HEADERS.tokensIn]: String(tokens.input), + [OMNIROUTE_RESPONSE_HEADERS.tokensOut]: String(tokens.output), + }; + + if (typeof model === "string" && model.trim().length > 0) { + headers[OMNIROUTE_RESPONSE_HEADERS.model] = model; + } + + if (typeof provider === "string" && provider.trim().length > 0) { + headers[OMNIROUTE_RESPONSE_HEADERS.provider] = getProviderAlias(provider); + } + + return headers; +} + +export function buildOmniRouteSseMetadataComment( + options: Parameters[0] +): string { + const headers = buildOmniRouteResponseMetaHeaders(options); + const lines = Object.entries(headers) + .filter(([, value]) => typeof value === "string" && value.trim().length > 0) + .map(([name, value]) => `: ${name.toLowerCase()}=${value}`); + + return lines.length > 0 ? `${lines.join("\n")}\n` : ""; +} diff --git a/src/domain/tagRouter.ts b/src/domain/tagRouter.ts new file mode 100644 index 0000000000..184d66a883 --- /dev/null +++ b/src/domain/tagRouter.ts @@ -0,0 +1,64 @@ +type JsonRecord = Record; + +export type RoutingTagMatchMode = "any" | "all"; + +function asRecord(value: unknown): JsonRecord { + return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; +} + +function normalizeSingleRoutingTag(value: unknown): string | null { + if (typeof value !== "string") return null; + const normalized = value.trim().toLowerCase(); + return normalized.length > 0 ? normalized : null; +} + +export function normalizeRoutingTags(value: unknown): string[] { + const rawValues = Array.isArray(value) + ? value + : typeof value === "string" + ? value.split(",") + : []; + + const deduped = new Set(); + for (const rawValue of rawValues) { + const normalized = normalizeSingleRoutingTag(rawValue); + if (normalized) deduped.add(normalized); + } + return Array.from(deduped); +} + +export function normalizeRoutingTagMatchMode(value: unknown): RoutingTagMatchMode { + const normalized = + typeof value === "string" ? value.trim().toLowerCase() : typeof value === "number" ? "" : ""; + return normalized === "all" ? "all" : "any"; +} + +export function getConnectionRoutingTags(providerSpecificData: unknown): string[] { + return normalizeRoutingTags(asRecord(providerSpecificData).tags); +} + +export function matchesRoutingTags( + connectionTags: string[], + requestTags: string[], + matchMode: RoutingTagMatchMode = "any" +): boolean { + if (requestTags.length === 0) return true; + if (connectionTags.length === 0) return false; + + const tagSet = new Set(connectionTags); + if (matchMode === "all") { + return requestTags.every((tag) => tagSet.has(tag)); + } + return requestTags.some((tag) => tagSet.has(tag)); +} + +export function resolveRequestRoutingTags(body: Record | null | undefined): { + tags: string[]; + matchMode: RoutingTagMatchMode; +} { + const metadata = asRecord(body?.metadata); + return { + tags: normalizeRoutingTags(metadata.tags), + matchMode: normalizeRoutingTagMatchMode(metadata.tag_match_mode ?? metadata.tagMatchMode), + }; +} diff --git a/src/instrumentation-node.ts b/src/instrumentation-node.ts index 4b945d1d51..c5479ba479 100755 --- a/src/instrumentation-node.ts +++ b/src/instrumentation-node.ts @@ -89,16 +89,20 @@ export async function registerNodejs(): Promise { { startBackgroundRefresh }, { startProviderLimitsSyncScheduler }, { getSettings }, + { startSpendBatchWriter }, ] = await Promise.all([ import("@/lib/gracefulShutdown"), import("@/lib/apiBridgeServer"), import("@/domain/quotaCache"), import("@/shared/services/providerLimitsSyncScheduler"), import("@/lib/db/settings"), + import("@/lib/spend/batchWriter"), ]); initGracefulShutdown(); initApiBridgeServer(); + startSpendBatchWriter(); + console.log("[STARTUP] Spend batch writer started"); if (!isBackgroundServicesDisabled()) { startBackgroundRefresh(); console.log("[STARTUP] Quota cache background refresh started"); @@ -151,6 +155,21 @@ export async function registerNodejs(): Promise { } } + if (settings.payloadRules) { + try { + const payloadRules = + typeof settings.payloadRules === "string" + ? JSON.parse(settings.payloadRules) + : settings.payloadRules; + const { setPayloadRulesConfig } = + await import("@omniroute/open-sse/services/payloadRules.ts"); + setPayloadRulesConfig(payloadRules); + console.log("[STARTUP] Restored payload rules config from settings"); + } catch (err: unknown) { + console.warn(`[STARTUP] Failed to parse payload rules settings:`, err); + } + } + const migration = await migrateCodexConnectionDefaultsFromLegacySettings(); if (migration.migrated) { console.log( diff --git a/src/lib/db/core.ts b/src/lib/db/core.ts index a97209e27b..d1036943ef 100644 --- a/src/lib/db/core.ts +++ b/src/lib/db/core.ts @@ -233,10 +233,29 @@ const SCHEMA_SQL = ` CREATE TABLE IF NOT EXISTS domain_budgets ( api_key_id TEXT PRIMARY KEY, daily_limit_usd REAL NOT NULL, + weekly_limit_usd REAL DEFAULT 0, monthly_limit_usd REAL DEFAULT 0, - warning_threshold REAL DEFAULT 0.8 + warning_threshold REAL DEFAULT 0.8, + reset_interval TEXT DEFAULT 'daily', + reset_time TEXT DEFAULT '00:00', + budget_reset_at INTEGER, + last_budget_reset_at INTEGER, + warning_emitted_at INTEGER, + warning_period_start INTEGER ); + CREATE TABLE IF NOT EXISTS domain_budget_reset_logs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + api_key_id TEXT NOT NULL, + reset_interval TEXT NOT NULL, + previous_spend REAL NOT NULL DEFAULT 0, + reset_at INTEGER NOT NULL, + next_reset_at INTEGER NOT NULL, + period_start INTEGER NOT NULL, + period_end INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_dbrl_key_reset ON domain_budget_reset_logs(api_key_id, reset_at DESC); + CREATE TABLE IF NOT EXISTS domain_cost_history ( id INTEGER PRIMARY KEY AUTOINCREMENT, api_key_id TEXT NOT NULL, diff --git a/src/lib/db/domainState.ts b/src/lib/db/domainState.ts index 05359ad503..680e51da04 100644 --- a/src/lib/db/domainState.ts +++ b/src/lib/db/domainState.ts @@ -10,9 +10,35 @@ * @module lib/db/domainState */ -import { getDbInstance, isBuildPhase, isCloud } from "./core"; +import { getDbInstance } from "./core"; type JsonRecord = Record; +type BudgetResetInterval = "daily" | "weekly" | "monthly"; + +interface BudgetConfigRecord { + dailyLimitUsd: number; + weeklyLimitUsd: number; + monthlyLimitUsd: number; + warningThreshold: number; + resetInterval: BudgetResetInterval; + resetTime: string; + budgetResetAt: number | null; + lastBudgetResetAt: number | null; + warningEmittedAt: number | null; + warningPeriodStart: number | null; +} + +interface BudgetResetLogRecord { + apiKeyId: string; + resetInterval: BudgetResetInterval; + previousSpend: number; + resetAt: number; + nextResetAt: number; + periodStart: number; + periodEnd: number; +} + +let _budgetSchemaChecked = false; function asRecord(value: unknown): JsonRecord { return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : {}; @@ -27,6 +53,60 @@ function toNumber(value: unknown, fallback = 0): number { return fallback; } +function ensureBudgetSchema() { + if (_budgetSchemaChecked) return; + + const db = getDbInstance(); + const columns = db.prepare("PRAGMA table_info(domain_budgets)").all(); + const columnNames = new Set( + columns + .map((column) => { + const record = asRecord(column); + return typeof record.name === "string" ? record.name : ""; + }) + .filter(Boolean) + ); + + if (!columnNames.has("weekly_limit_usd")) { + db.exec("ALTER TABLE domain_budgets ADD COLUMN weekly_limit_usd REAL DEFAULT 0"); + } + if (!columnNames.has("reset_interval")) { + db.exec("ALTER TABLE domain_budgets ADD COLUMN reset_interval TEXT DEFAULT 'daily'"); + } + if (!columnNames.has("reset_time")) { + db.exec("ALTER TABLE domain_budgets ADD COLUMN reset_time TEXT DEFAULT '00:00'"); + } + if (!columnNames.has("budget_reset_at")) { + db.exec("ALTER TABLE domain_budgets ADD COLUMN budget_reset_at INTEGER"); + } + if (!columnNames.has("last_budget_reset_at")) { + db.exec("ALTER TABLE domain_budgets ADD COLUMN last_budget_reset_at INTEGER"); + } + if (!columnNames.has("warning_emitted_at")) { + db.exec("ALTER TABLE domain_budgets ADD COLUMN warning_emitted_at INTEGER"); + } + if (!columnNames.has("warning_period_start")) { + db.exec("ALTER TABLE domain_budgets ADD COLUMN warning_period_start INTEGER"); + } + + db.exec(` + CREATE TABLE IF NOT EXISTS domain_budget_reset_logs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + api_key_id TEXT NOT NULL, + reset_interval TEXT NOT NULL, + previous_spend REAL NOT NULL DEFAULT 0, + reset_at INTEGER NOT NULL, + next_reset_at INTEGER NOT NULL, + period_start INTEGER NOT NULL, + period_end INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_dbrl_key_reset + ON domain_budget_reset_logs(api_key_id, reset_at DESC); + `); + + _budgetSchemaChecked = true; +} + // ──────────────── Fallback Chains ──────────────── /** @@ -99,15 +179,35 @@ export function deleteAllFallbackChains() { * @param {{ dailyLimitUsd: number, monthlyLimitUsd?: number, warningThreshold?: number }} config */ export function saveBudget(apiKeyId, config) { + ensureBudgetSchema(); const db = getDbInstance(); db.prepare( - `INSERT OR REPLACE INTO domain_budgets (api_key_id, daily_limit_usd, monthly_limit_usd, warning_threshold) - VALUES (?, ?, ?, ?)` + `INSERT OR REPLACE INTO domain_budgets ( + api_key_id, + daily_limit_usd, + weekly_limit_usd, + monthly_limit_usd, + warning_threshold, + reset_interval, + reset_time, + budget_reset_at, + last_budget_reset_at, + warning_emitted_at, + warning_period_start + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)` ).run( apiKeyId, - config.dailyLimitUsd, - config.monthlyLimitUsd || 0, - config.warningThreshold ?? 0.8 + toNumber(config.dailyLimitUsd), + toNumber(config.weeklyLimitUsd), + toNumber(config.monthlyLimitUsd), + toNumber(config.warningThreshold, 0.8), + typeof config.resetInterval === "string" ? config.resetInterval : "daily", + typeof config.resetTime === "string" ? config.resetTime : "00:00", + config.budgetResetAt ?? null, + config.lastBudgetResetAt ?? null, + config.warningEmittedAt ?? null, + config.warningPeriodStart ?? null ); } @@ -117,24 +217,129 @@ export function saveBudget(apiKeyId, config) { * @returns {{ dailyLimitUsd: number, monthlyLimitUsd: number, warningThreshold: number } | null} */ export function loadBudget(apiKeyId) { + ensureBudgetSchema(); const db = getDbInstance(); const row = db.prepare("SELECT * FROM domain_budgets WHERE api_key_id = ?").get(apiKeyId); const record = asRecord(row); if (!row) return null; return { dailyLimitUsd: toNumber(record.daily_limit_usd), + weeklyLimitUsd: toNumber(record.weekly_limit_usd), monthlyLimitUsd: toNumber(record.monthly_limit_usd), warningThreshold: toNumber(record.warning_threshold, 0.8), + resetInterval: + typeof record.reset_interval === "string" ? record.reset_interval : ("daily" as const), + resetTime: typeof record.reset_time === "string" ? record.reset_time : "00:00", + budgetResetAt: toNumber(record.budget_reset_at, 0) || null, + lastBudgetResetAt: toNumber(record.last_budget_reset_at, 0) || null, + warningEmittedAt: toNumber(record.warning_emitted_at, 0) || null, + warningPeriodStart: toNumber(record.warning_period_start, 0) || null, }; } +/** + * Load all budget configs. + * @returns {Record} + */ +export function loadAllBudgets() { + ensureBudgetSchema(); + const db = getDbInstance(); + const rows = db.prepare("SELECT * FROM domain_budgets ORDER BY api_key_id").all(); + const result: Record = {}; + + for (const row of rows) { + const record = asRecord(row); + const apiKeyId = typeof record.api_key_id === "string" ? record.api_key_id : ""; + if (!apiKeyId) continue; + + result[apiKeyId] = { + dailyLimitUsd: toNumber(record.daily_limit_usd), + weeklyLimitUsd: toNumber(record.weekly_limit_usd), + monthlyLimitUsd: toNumber(record.monthly_limit_usd), + warningThreshold: toNumber(record.warning_threshold, 0.8), + resetInterval: + typeof record.reset_interval === "string" + ? (record.reset_interval as BudgetResetInterval) + : "daily", + resetTime: typeof record.reset_time === "string" ? record.reset_time : "00:00", + budgetResetAt: toNumber(record.budget_reset_at, 0) || null, + lastBudgetResetAt: toNumber(record.last_budget_reset_at, 0) || null, + warningEmittedAt: toNumber(record.warning_emitted_at, 0) || null, + warningPeriodStart: toNumber(record.warning_period_start, 0) || null, + }; + } + + return result; +} + +/** + * Persist a budget reset log entry. + * @param {BudgetResetLogRecord} entry + */ +export function saveBudgetResetLog(entry: BudgetResetLogRecord) { + ensureBudgetSchema(); + const db = getDbInstance(); + db.prepare( + `INSERT INTO domain_budget_reset_logs + (api_key_id, reset_interval, previous_spend, reset_at, next_reset_at, period_start, period_end) + VALUES (?, ?, ?, ?, ?, ?, ?)` + ).run( + entry.apiKeyId, + entry.resetInterval, + entry.previousSpend, + entry.resetAt, + entry.nextResetAt, + entry.periodStart, + entry.periodEnd + ); +} + +/** + * Load recent budget reset logs for an API key. + * @param {string} apiKeyId + * @param {number} [limit=10] + * @returns {Array} + */ +export function loadBudgetResetLogs(apiKeyId: string, limit = 10) { + ensureBudgetSchema(); + const db = getDbInstance(); + return db + .prepare( + `SELECT id, api_key_id, reset_interval, previous_spend, reset_at, next_reset_at, period_start, period_end + FROM domain_budget_reset_logs + WHERE api_key_id = ? + ORDER BY reset_at DESC + LIMIT ?` + ) + .all(apiKeyId, Math.max(1, Math.floor(limit))) + .map((row) => { + const record = asRecord(row); + return { + id: toNumber(record.id), + apiKeyId: typeof record.api_key_id === "string" ? record.api_key_id : "", + resetInterval: + typeof record.reset_interval === "string" + ? (record.reset_interval as BudgetResetInterval) + : "daily", + previousSpend: toNumber(record.previous_spend), + resetAt: toNumber(record.reset_at), + nextResetAt: toNumber(record.next_reset_at), + periodStart: toNumber(record.period_start), + periodEnd: toNumber(record.period_end), + }; + }) + .filter((entry) => entry.apiKeyId.length > 0); +} + /** * Delete a budget config. * @param {string} apiKeyId */ export function deleteBudget(apiKeyId) { + ensureBudgetSchema(); const db = getDbInstance(); db.prepare("DELETE FROM domain_budgets WHERE api_key_id = ?").run(apiKeyId); + db.prepare("DELETE FROM domain_budget_reset_logs WHERE api_key_id = ?").run(apiKeyId); } // ──────────────── Cost History ──────────────── @@ -146,6 +351,7 @@ export function deleteBudget(apiKeyId) { * @param {number} [timestamp] */ export function saveCostEntry(apiKeyId, cost, timestamp = Date.now()) { + ensureBudgetSchema(); const db = getDbInstance(); db.prepare("INSERT INTO domain_cost_history (api_key_id, cost, timestamp) VALUES (?, ?, ?)").run( apiKeyId, @@ -154,6 +360,27 @@ export function saveCostEntry(apiKeyId, cost, timestamp = Date.now()) { ); } +export function batchSaveCostEntries( + entries: Array<{ apiKeyId: string; cost: number; timestamp: number }> +) { + ensureBudgetSchema(); + if (!Array.isArray(entries) || entries.length === 0) return; + + const db = getDbInstance(); + const stmt = db.prepare( + "INSERT INTO domain_cost_history (api_key_id, cost, timestamp) VALUES (?, ?, ?)" + ); + const tx = db.transaction( + (rows: Array<{ apiKeyId: string; cost: number; timestamp: number }>) => { + for (const entry of rows) { + stmt.run(entry.apiKeyId, entry.cost, entry.timestamp); + } + } + ); + + tx(entries); +} + /** * Load cost entries for an API key within a time window. * @param {string} apiKeyId @@ -161,6 +388,7 @@ export function saveCostEntry(apiKeyId, cost, timestamp = Date.now()) { * @returns {Array<{cost: number, timestamp: number}>} */ export function loadCostEntries(apiKeyId, sinceTimestamp) { + ensureBudgetSchema(); const db = getDbInstance(); return db .prepare( @@ -169,12 +397,37 @@ export function loadCostEntries(apiKeyId, sinceTimestamp) { .all(apiKeyId, sinceTimestamp); } +/** + * Load cost entries for an API key within a bounded time window. + * @param {string} apiKeyId + * @param {number} sinceTimestamp + * @param {number} untilTimestamp + * @returns {Array<{cost: number, timestamp: number}>} + */ +export function loadCostEntriesInRange( + apiKeyId: string, + sinceTimestamp: number, + untilTimestamp: number +) { + ensureBudgetSchema(); + const db = getDbInstance(); + return db + .prepare( + `SELECT cost, timestamp + FROM domain_cost_history + WHERE api_key_id = ? AND timestamp >= ? AND timestamp < ? + ORDER BY timestamp` + ) + .all(apiKeyId, sinceTimestamp, untilTimestamp); +} + /** * Delete old cost entries (cleanup). * @param {number} olderThanTimestamp * @returns {number} deleted count */ export function cleanOldCostEntries(olderThanTimestamp) { + ensureBudgetSchema(); const db = getDbInstance(); const info = db .prepare("DELETE FROM domain_cost_history WHERE timestamp < ?") @@ -187,6 +440,7 @@ export function cleanOldCostEntries(olderThanTimestamp) { * @param {string} apiKeyId */ export function deleteCostEntries(apiKeyId) { + ensureBudgetSchema(); const db = getDbInstance(); db.prepare("DELETE FROM domain_cost_history WHERE api_key_id = ?").run(apiKeyId); } @@ -195,9 +449,11 @@ export function deleteCostEntries(apiKeyId) { * Delete all cost data. */ export function deleteAllCostData() { + ensureBudgetSchema(); const db = getDbInstance(); db.prepare("DELETE FROM domain_cost_history").run(); db.prepare("DELETE FROM domain_budgets").run(); + db.prepare("DELETE FROM domain_budget_reset_logs").run(); } // ──────────────── Lockout State ──────────────── diff --git a/src/lib/gracefulShutdown.ts b/src/lib/gracefulShutdown.ts index c84c069d2c..baf6bd75fe 100644 --- a/src/lib/gracefulShutdown.ts +++ b/src/lib/gracefulShutdown.ts @@ -96,10 +96,17 @@ async function waitForDrain(): Promise { */ async function cleanup(): Promise { try { - const [{ closeAuditDb }, { closeDbInstance }] = await Promise.all([ + const [{ closeAuditDb }, { closeDbInstance }, { flushSpendBatchWriter }] = await Promise.all([ import("@omniroute/open-sse/mcp-server/audit.ts"), import("@/lib/db/core"), + import("@/lib/spend/batchWriter"), ]); + const flushResult = await flushSpendBatchWriter(); + if (flushResult.flushedEntries > 0) { + console.log( + `[Shutdown] Spend batch writer flushed ${flushResult.flushedEntries} pending entry(ies).` + ); + } if (closeAuditDb()) { console.log("[Shutdown] MCP audit database checkpointed and closed."); } diff --git a/src/lib/initCloudSync.ts b/src/lib/initCloudSync.ts index a3c188f57a..7b834b517c 100644 --- a/src/lib/initCloudSync.ts +++ b/src/lib/initCloudSync.ts @@ -1,4 +1,5 @@ import initializeCloudSync from "@/shared/services/initializeCloudSync"; +import { startBudgetResetJob } from "@/lib/jobs/budgetResetJob"; import { startModelSyncScheduler } from "@/shared/services/modelSyncScheduler"; import "@/lib/tokenHealthCheck"; // Proactive token health-check scheduler @@ -19,6 +20,7 @@ export async function ensureCloudSyncInitialized() { try { await initializeCloudSync(); startModelSyncScheduler(); + startBudgetResetJob(); initialized = true; } catch (error) { console.error("[ServerInit] Error initializing background sync services:", error); diff --git a/src/lib/jobs/budgetResetJob.ts b/src/lib/jobs/budgetResetJob.ts new file mode 100644 index 0000000000..bd640569e8 --- /dev/null +++ b/src/lib/jobs/budgetResetJob.ts @@ -0,0 +1,40 @@ +import { syncAllBudgetSchedules } from "@/domain/costRules"; + +const DEFAULT_INTERVAL_MS = 10 * 60 * 1000; + +let timer: NodeJS.Timeout | null = null; + +function getIntervalMs() { + const raw = process.env.OMNIROUTE_BUDGET_RESET_JOB_INTERVAL_MS; + const parsed = raw ? Number(raw) : Number.NaN; + return Number.isFinite(parsed) && parsed >= 10_000 ? parsed : DEFAULT_INTERVAL_MS; +} + +export function startBudgetResetJob() { + if (timer) { + return timer; + } + + const run = () => { + try { + const result = syncAllBudgetSchedules(Date.now()); + if (result.resetCount > 0) { + console.log(`[BudgetReset] processed=${result.processed} reset=${result.resetCount}`); + } + } catch (error) { + console.error("[BudgetReset] Job failed:", error); + } + }; + + run(); + timer = setInterval(run, getIntervalMs()); + timer.unref?.(); + return timer; +} + +export function stopBudgetResetJob() { + if (timer) { + clearInterval(timer); + timer = null; + } +} diff --git a/src/lib/providers/requestDefaults.ts b/src/lib/providers/requestDefaults.ts index d49d1021e7..d190bf85d8 100644 --- a/src/lib/providers/requestDefaults.ts +++ b/src/lib/providers/requestDefaults.ts @@ -1,5 +1,8 @@ type JsonRecord = Record; +import { normalizeExcludedModelPatterns } from "@/domain/connectionModelRules"; +import { normalizeRoutingTags } from "@/domain/tagRouter"; + export const CODEX_REASONING_EFFORT_VALUES = ["none", "low", "medium", "high", "xhigh"] as const; export type CodexReasoningEffort = (typeof CODEX_REASONING_EFFORT_VALUES)[number]; @@ -85,6 +88,40 @@ export function normalizeProviderSpecificData( delete normalized.openaiStoreEnabled; } + if ("tag" in normalized) { + if (typeof normalized.tag === "string") { + const trimmedTag = normalized.tag.trim(); + if (trimmedTag) { + normalized.tag = trimmedTag; + } else { + delete normalized.tag; + } + } else { + delete normalized.tag; + } + } + + if ("tags" in normalized) { + const tags = normalizeRoutingTags(normalized.tags); + if (tags.length > 0) { + normalized.tags = tags; + } else { + delete normalized.tags; + } + } + + if ("excludedModels" in normalized || "excluded_models" in normalized) { + const excludedModels = normalizeExcludedModelPatterns( + normalized.excludedModels ?? normalized.excluded_models + ); + if (excludedModels.length > 0) { + normalized.excludedModels = excludedModels; + } else { + delete normalized.excludedModels; + } + delete normalized.excluded_models; + } + return Object.keys(normalized).length > 0 ? normalized : undefined; } diff --git a/src/lib/providers/validation.ts b/src/lib/providers/validation.ts index de866585c0..55824b5235 100644 --- a/src/lib/providers/validation.ts +++ b/src/lib/providers/validation.ts @@ -985,7 +985,7 @@ async function validateSearchProvider( const SEARCH_VALIDATOR_CONFIGS: Record< string, - (apiKey: string) => { url: string; init: RequestInit } + (apiKey: string, providerSpecificData?: any) => { url: string; init: RequestInit } > = { "serper-search": (apiKey) => ({ url: "https://google.serper.dev/search", @@ -1026,6 +1026,57 @@ const SEARCH_VALIDATOR_CONFIGS: Record< body: JSON.stringify({ query: "test", max_results: 1 }), }, }), + "google-pse-search": (apiKey, providerSpecificData = {}) => { + const cx = providerSpecificData?.cx; + if (!cx || typeof cx !== "string") { + throw new Error("Programmable Search Engine ID (cx) is required"); + } + return { + url: `https://www.googleapis.com/customsearch/v1?key=${encodeURIComponent(apiKey)}&cx=${encodeURIComponent( + cx + )}&q=test&num=1`, + init: { + method: "GET", + headers: { Accept: "application/json" }, + }, + }; + }, + "linkup-search": (apiKey) => ({ + url: "https://api.linkup.so/v1/search", + init: { + method: "POST", + headers: { "Content-Type": "application/json", Authorization: `Bearer ${apiKey}` }, + body: JSON.stringify({ + q: "test", + depth: "standard", + outputType: "searchResults", + maxResults: 1, + }), + }, + }), + "searchapi-search": (apiKey) => ({ + url: `https://www.searchapi.io/api/v1/search?engine=google&q=test&api_key=${encodeURIComponent( + apiKey + )}`, + init: { + method: "GET", + headers: { Accept: "application/json" }, + }, + }), + "searxng-search": (_apiKey, providerSpecificData = {}) => { + const baseUrl = + typeof providerSpecificData?.baseUrl === "string" && providerSpecificData.baseUrl.trim() + ? providerSpecificData.baseUrl.trim().replace(/\/+$/, "") + : "http://localhost:8888/search"; + const searchUrl = baseUrl.endsWith("/search") ? baseUrl : `${baseUrl}/search`; + return { + url: `${searchUrl}?q=test&format=json`, + init: { + method: "GET", + headers: { Accept: "application/json" }, + }, + }; + }, }; async function validateGrokWebProvider({ apiKey, providerSpecificData = {} }: any) { @@ -1121,7 +1172,8 @@ async function validateGrokWebProvider({ apiKey, providerSpecificData = {} }: an } export async function validateProviderApiKey({ provider, apiKey, providerSpecificData = {} }: any) { - if (!provider || !apiKey) { + const requiresApiKey = provider !== "searxng-search"; + if (!provider || (requiresApiKey && !apiKey)) { return { valid: false, error: "Provider and API key required", unsupported: false }; } @@ -1197,7 +1249,7 @@ export async function validateProviderApiKey({ provider, apiKey, providerSpecifi Object.entries(SEARCH_VALIDATOR_CONFIGS).map(([id, configFn]) => [ id, ({ apiKey, providerSpecificData }: any) => { - const { url, init } = configFn(apiKey); + const { url, init } = configFn(apiKey, providerSpecificData); return validateSearchProvider(url, init, providerSpecificData); }, ]) diff --git a/src/lib/spend/batchWriter.ts b/src/lib/spend/batchWriter.ts new file mode 100644 index 0000000000..483c96c3a7 --- /dev/null +++ b/src/lib/spend/batchWriter.ts @@ -0,0 +1,208 @@ +import { batchSaveCostEntries } from "@/lib/db/domainState"; + +export interface BufferedCostEntry { + apiKeyId: string; + cost: number; + timestamp: number; +} + +interface SpendBatchWriterOptions { + flushIntervalMs?: number; + maxBufferSize?: number; + persistEntries?: (entries: BufferedCostEntry[]) => Promise | void; + logger?: Pick; +} + +type FlushResult = { + flushedEntries: number; + uniqueKeys: number; + requeued: boolean; +}; + +const DEFAULT_FLUSH_INTERVAL_MS = 60_000; +const DEFAULT_MAX_BUFFER_SIZE = 1_000; + +function getFlushIntervalMs() { + const parsed = Number.parseInt(process.env.OMNIROUTE_SPEND_FLUSH_INTERVAL_MS || "", 10); + return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_FLUSH_INTERVAL_MS; +} + +function getMaxBufferSize() { + const parsed = Number.parseInt(process.env.OMNIROUTE_SPEND_MAX_BUFFER_SIZE || "", 10); + return Number.isFinite(parsed) && parsed > 0 ? parsed : DEFAULT_MAX_BUFFER_SIZE; +} + +function normalizeEntry(entry: BufferedCostEntry): BufferedCostEntry | null { + if (!entry?.apiKeyId || !Number.isFinite(entry.cost) || entry.cost <= 0) return null; + return { + apiKeyId: entry.apiKeyId, + cost: entry.cost, + timestamp: Number.isFinite(entry.timestamp) ? entry.timestamp : Date.now(), + }; +} + +export class SpendBatchWriter { + private buffer: BufferedCostEntry[] = []; + private inFlightEntries: BufferedCostEntry[] = []; + private discardedApiKeyIds = new Set(); + private timer: NodeJS.Timeout | null = null; + private started = false; + private flushPromise: Promise | null = null; + private persistEntries: (entries: BufferedCostEntry[]) => Promise | void; + private logger: Pick; + private flushIntervalMs: number; + private maxBufferSize: number; + + constructor(options: SpendBatchWriterOptions = {}) { + this.persistEntries = options.persistEntries || batchSaveCostEntries; + this.logger = options.logger || console; + this.flushIntervalMs = options.flushIntervalMs ?? getFlushIntervalMs(); + this.maxBufferSize = options.maxBufferSize ?? getMaxBufferSize(); + } + + start() { + if (this.started) return; + this.started = true; + this.timer = setInterval(() => { + void this.flush(); + }, this.flushIntervalMs); + this.timer.unref?.(); + } + + increment(apiKeyId: string, cost: number, timestamp = Date.now()) { + const entry = normalizeEntry({ apiKeyId, cost, timestamp }); + if (!entry) return; + + this.start(); + this.discardedApiKeyIds.delete(entry.apiKeyId); + this.buffer.push(entry); + + if (this.buffer.length >= this.maxBufferSize) { + void this.flush(); + } + } + + getBufferedEntries( + apiKeyId: string, + sinceTimestamp = 0, + untilTimestamp = Number.POSITIVE_INFINITY + ) { + const matchesWindow = (entry: BufferedCostEntry) => + entry.apiKeyId === apiKeyId && + entry.timestamp >= sinceTimestamp && + entry.timestamp < untilTimestamp; + + return [...this.inFlightEntries, ...this.buffer].filter(matchesWindow); + } + + getPendingCostTotal( + apiKeyId: string, + sinceTimestamp = 0, + untilTimestamp = Number.POSITIVE_INFINITY + ) { + return this.getBufferedEntries(apiKeyId, sinceTimestamp, untilTimestamp).reduce( + (sum, entry) => sum + entry.cost, + 0 + ); + } + + discardEntries(apiKeyId: string) { + this.discardedApiKeyIds.add(apiKeyId); + this.buffer = this.buffer.filter((entry) => entry.apiKeyId !== apiKeyId); + this.inFlightEntries = this.inFlightEntries.filter((entry) => entry.apiKeyId !== apiKeyId); + } + + async flush(): Promise { + if (this.flushPromise) { + return this.flushPromise; + } + + if (this.buffer.length === 0) { + return { flushedEntries: 0, uniqueKeys: 0, requeued: false }; + } + + const entriesToFlush = [...this.buffer]; + this.buffer = []; + this.inFlightEntries = entriesToFlush; + + this.flushPromise = (async () => { + const entriesToPersist = entriesToFlush.filter( + (entry) => !this.discardedApiKeyIds.has(entry.apiKeyId) + ); + const uniqueKeys = new Set(entriesToPersist.map((entry) => entry.apiKeyId)).size; + + try { + if (entriesToPersist.length > 0) { + await this.persistEntries(entriesToPersist); + } + this.logger.log( + `[SpendWriter] Flushed ${entriesToPersist.length} cost entr${ + entriesToPersist.length === 1 ? "y" : "ies" + } across ${uniqueKeys} key(s)` + ); + return { + flushedEntries: entriesToPersist.length, + uniqueKeys, + requeued: false, + }; + } catch (error) { + this.buffer = [...entriesToPersist, ...this.buffer]; + const message = error instanceof Error ? error.message : String(error); + this.logger.error(`[SpendWriter] Flush error: ${message}`); + return { + flushedEntries: 0, + uniqueKeys, + requeued: true, + }; + } finally { + this.inFlightEntries = []; + this.flushPromise = null; + } + })(); + + return this.flushPromise; + } + + async stop() { + if (this.timer) { + clearInterval(this.timer); + this.timer = null; + } + this.started = false; + return this.flush(); + } + + resetForTests() { + if (this.timer) { + clearInterval(this.timer); + this.timer = null; + } + this.started = false; + this.buffer = []; + this.inFlightEntries = []; + this.discardedApiKeyIds.clear(); + this.flushPromise = null; + } +} + +export const spendBatchWriter = new SpendBatchWriter(); + +export function startSpendBatchWriter() { + spendBatchWriter.start(); +} + +export async function flushSpendBatchWriter() { + return spendBatchWriter.flush(); +} + +export async function stopSpendBatchWriter() { + return spendBatchWriter.stop(); +} + +export function resetSpendBatchWriterForTests() { + spendBatchWriter.resetForTests(); +} + +export function discardSpendBatchEntries(apiKeyId: string) { + spendBatchWriter.discardEntries(apiKeyId); +} diff --git a/src/server-init.ts b/src/server-init.ts index 42056f25e4..1c0bb818d5 100644 --- a/src/server-init.ts +++ b/src/server-init.ts @@ -4,6 +4,10 @@ import { enforceWebRuntimeEnv } from "./lib/env/runtimeEnv"; import { enforceSecrets } from "./shared/utils/secretsValidator"; import { initAuditLog, cleanupExpiredLogs, logAuditEvent } from "./lib/compliance/index"; import { initConsoleInterceptor } from "./lib/consoleInterceptor"; +import { startBudgetResetJob } from "./lib/jobs/budgetResetJob"; +import { getSettings } from "./lib/db/settings"; +import { setPayloadRulesConfig } from "@omniroute/open-sse/services/payloadRules.ts"; +import { startSpendBatchWriter } from "./lib/spend/batchWriter"; async function startServer() { // Trigger request-log layout migration during startup, before serving requests. @@ -44,8 +48,21 @@ async function startServer() { console.log("Starting server with cloud sync..."); try { + const settings = await getSettings(); + if (settings.payloadRules) { + const payloadRules = + typeof settings.payloadRules === "string" + ? JSON.parse(settings.payloadRules) + : settings.payloadRules; + setPayloadRulesConfig(payloadRules); + console.log("[STARTUP] Restored payload rules config from settings"); + } + // Initialize cloud sync + startSpendBatchWriter(); + console.log("[STARTUP] Spend batch writer started"); await initializeCloudSync(); + startBudgetResetJob(); console.log("Server started with cloud sync initialized"); // Log server start event to audit log diff --git a/src/shared/components/ModelSelectModal.tsx b/src/shared/components/ModelSelectModal.tsx index a01927f356..c9b62d4dbd 100644 --- a/src/shared/components/ModelSelectModal.tsx +++ b/src/shared/components/ModelSelectModal.tsx @@ -6,6 +6,10 @@ import { useTranslations } from "next-intl"; import Modal from "./Modal"; import { getModelsByProviderId, PROVIDER_ID_TO_ALIAS } from "@/shared/constants/models"; import { getCompatibleFallbackModels } from "@/lib/providers/managedAvailableModels"; +import { + getModelCatalogSourceLabel, + matchesModelCatalogQuery, +} from "@/shared/utils/modelCatalogSearch"; import { OAUTH_PROVIDERS, FREE_PROVIDERS, @@ -126,6 +130,7 @@ export default function ModelSelectModal({ id: fullModel.replace(`${alias}/`, ""), name: aliasName, value: fullModel, + source: "alias", })); // Merge custom models for passthrough providers @@ -136,6 +141,7 @@ export default function ModelSelectModal({ name: cm.name || cm.id, value: `${alias}/${cm.id}`, isCustom: true, + source: cm.source === "api-sync" ? "api-sync" : "custom", })); const allModels = [...aliasModels, ...customEntries]; @@ -162,6 +168,7 @@ export default function ModelSelectModal({ id: fullModel.replace(`${providerId}/`, ""), name: aliasName, value: `${nodePrefix}/${fullModel.replace(`${providerId}/`, "")}`, + source: "alias", })); const fallbackEntries = ( @@ -173,6 +180,7 @@ export default function ModelSelectModal({ name: fm.name || fm.id, value: `${nodePrefix}/${fm.id}`, isFallback: true, + source: "fallback", })); // Merge custom models for custom providers @@ -187,6 +195,7 @@ export default function ModelSelectModal({ name: cm.name || cm.id, value: `${nodePrefix}/${cm.id}`, isCustom: true, + source: cm.source === "api-sync" ? "api-sync" : "custom", })); const allModels = [...nodeModels, ...fallbackEntries, ...customEntries]; @@ -209,6 +218,7 @@ export default function ModelSelectModal({ id: m.id, name: m.name, value: `${alias}/${m.id}`, + source: "system", })); const customEntries = providerCustomModels @@ -218,6 +228,7 @@ export default function ModelSelectModal({ name: cm.name || cm.id, value: `${alias}/${cm.id}`, isCustom: true, + source: cm.source === "api-sync" ? "api-sync" : "custom", })); const allModels = [...systemEntries, ...customEntries]; @@ -251,8 +262,12 @@ export default function ModelSelectModal({ const filtered: Record = {}; Object.entries(groupedModels).forEach(([providerId, group]: [string, any]) => { - const matchedModels = group.models.filter( - (m) => m.name.toLowerCase().includes(query) || m.id.toLowerCase().includes(query) + const matchedModels = group.models.filter((model) => + matchesModelCatalogQuery(query, { + modelId: model.id, + modelName: model.name, + source: model.source, + }) ); const providerNameMatches = group.name.toLowerCase().includes(query); @@ -260,7 +275,7 @@ export default function ModelSelectModal({ if (matchedModels.length > 0 || providerNameMatches) { filtered[providerId] = { ...group, - models: matchedModels, + models: matchedModels.length > 0 ? matchedModels : group.models, }; } }); @@ -368,7 +383,11 @@ export default function ModelSelectModal({ > {isAdded && ✓} {model.name} - {model.isCustom ? " ★" : ""} + {model.source && ( + + {getModelCatalogSourceLabel(model.source)} + + )} ); })} diff --git a/src/shared/components/ProviderIcon.tsx b/src/shared/components/ProviderIcon.tsx index 87ab662aae..eddf45fcf8 100644 --- a/src/shared/components/ProviderIcon.tsx +++ b/src/shared/components/ProviderIcon.tsx @@ -25,6 +25,7 @@ const LOBEHUB_PROVIDER_MAP: Record = { gemini: "google", "gemini-cli": "gemini", google: "google", + "google-pse-search": "google", deepseek: "deepseek", groq: "groq", mistral: "mistral", diff --git a/src/shared/constants/headers.ts b/src/shared/constants/headers.ts new file mode 100644 index 0000000000..74ef3e3e0d --- /dev/null +++ b/src/shared/constants/headers.ts @@ -0,0 +1,11 @@ +export const OMNIROUTE_RESPONSE_HEADERS = { + cache: "X-OmniRoute-Cache", + cacheHit: "X-OmniRoute-Cache-Hit", + latencyMs: "X-OmniRoute-Latency-Ms", + model: "X-OmniRoute-Model", + progress: "X-OmniRoute-Progress", + provider: "X-OmniRoute-Provider", + responseCost: "X-OmniRoute-Response-Cost", + tokensIn: "X-OmniRoute-Tokens-In", + tokensOut: "X-OmniRoute-Tokens-Out", +} as const; diff --git a/src/shared/constants/providers.ts b/src/shared/constants/providers.ts index bb751b1e31..d8497fea38 100644 --- a/src/shared/constants/providers.ts +++ b/src/shared/constants/providers.ts @@ -961,6 +961,46 @@ export const SEARCH_PROVIDERS = { website: "https://tavily.com", authHint: "API key from app.tavily.com (format: tvly-...)", }, + "google-pse-search": { + id: "google-pse-search", + alias: "google-pse", + name: "Google Programmable Search", + icon: "travel_explore", + color: "#4285F4", + textIcon: "GP", + website: "https://developers.google.com/custom-search/v1/overview", + authHint: "Requires a Google API key and your Programmable Search Engine ID (cx)", + }, + "linkup-search": { + id: "linkup-search", + alias: "linkup", + name: "Linkup Search", + icon: "public", + color: "#0F766E", + textIcon: "LU", + website: "https://docs.linkup.so", + authHint: "Bearer API key from the Linkup dashboard", + }, + "searchapi-search": { + id: "searchapi-search", + alias: "searchapi", + name: "SearchAPI", + icon: "manage_search", + color: "#2563EB", + textIcon: "SA", + website: "https://www.searchapi.io/docs", + authHint: "API key from SearchAPI (query param or Bearer auth)", + }, + "searxng-search": { + id: "searxng-search", + alias: "searxng", + name: "SearXNG Search", + icon: "search", + color: "#1A237E", + textIcon: "SX", + website: "https://docs.searxng.org", + authHint: "Set your SearXNG base URL. API key is optional for public/self-hosted instances.", + }, }; // Audio Only Providers diff --git a/src/shared/utils/modelCatalogSearch.ts b/src/shared/utils/modelCatalogSearch.ts new file mode 100644 index 0000000000..bf9a0bed11 --- /dev/null +++ b/src/shared/utils/modelCatalogSearch.ts @@ -0,0 +1,71 @@ +export type ModelCatalogSource = "system" | "custom" | "api-sync" | "fallback" | "alias"; + +type ModelCatalogTarget = { + modelId?: string | null; + modelName?: string | null; + alias?: string | null; + source?: string | null; +}; + +function normalizeText(value: string | null | undefined): string { + return typeof value === "string" ? value.trim().toLowerCase() : ""; +} + +export function normalizeModelCatalogSource(source?: string | null): ModelCatalogSource { + const normalized = normalizeText(source); + + if (normalized === "api-sync" || normalized === "synced") return "api-sync"; + if (normalized === "fallback") return "fallback"; + if (normalized === "alias") return "alias"; + if (normalized === "custom" || normalized === "manual" || normalized === "imported") { + return "custom"; + } + + return "system"; +} + +export function getModelCatalogSourceLabel(source?: string | null): string { + switch (normalizeModelCatalogSource(source)) { + case "api-sync": + return "Synced"; + case "custom": + return "Custom"; + case "fallback": + return "Fallback"; + case "alias": + return "Alias"; + case "system": + default: + return "Built-in"; + } +} + +function getModelCatalogSourceSearchText(source?: string | null): string { + switch (normalizeModelCatalogSource(source)) { + case "api-sync": + return "synced api imported discovered"; + case "custom": + return "custom manual imported"; + case "fallback": + return "fallback compatible"; + case "alias": + return "alias shortcut"; + case "system": + default: + return "built-in builtin official catalog"; + } +} + +export function matchesModelCatalogQuery(query: string, target: ModelCatalogTarget): boolean { + const normalizedQuery = normalizeText(query); + if (!normalizedQuery) return true; + + const haystacks = [ + normalizeText(target.modelId), + normalizeText(target.modelName), + normalizeText(target.alias), + getModelCatalogSourceSearchText(target.source), + ].filter(Boolean); + + return haystacks.some((value) => value.includes(normalizedQuery)); +} diff --git a/src/shared/validation/schemas.ts b/src/shared/validation/schemas.ts index c9539d526a..24c3f9e595 100644 --- a/src/shared/validation/schemas.ts +++ b/src/shared/validation/schemas.ts @@ -42,6 +42,15 @@ function validateProviderSpecificData( }); } + const cx = data.cx; + if (cx !== undefined && cx !== null && (typeof cx !== "string" || cx.length > 500)) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "providerSpecificData.cx must be a string up to 500 chars", + path: ["cx"], + }); + } + const openaiStoreEnabled = data.openaiStoreEnabled; if (openaiStoreEnabled !== undefined && typeof openaiStoreEnabled !== "boolean") { ctx.addIssue({ @@ -110,6 +119,75 @@ function validateProviderSpecificData( path: ["consoleApiKey"], }); } + + const groupTag = data.tag; + if ( + groupTag !== undefined && + groupTag !== null && + (typeof groupTag !== "string" || groupTag.length > 100) + ) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "providerSpecificData.tag must be a string up to 100 chars", + path: ["tag"], + }); + } + + const routingTags = data.tags; + if (routingTags !== undefined && routingTags !== null) { + if (!Array.isArray(routingTags) || routingTags.length > 50) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "providerSpecificData.tags must be an array with at most 50 items", + path: ["tags"], + }); + } else if ( + routingTags.some( + (tag) => typeof tag !== "string" || tag.trim().length === 0 || tag.trim().length > 64 + ) + ) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: + "providerSpecificData.tags must contain non-empty strings up to 64 characters each", + path: ["tags"], + }); + } + } + + const excludedModels = data.excludedModels ?? data.excluded_models; + if (excludedModels !== undefined && excludedModels !== null) { + if (typeof excludedModels === "string") { + if (excludedModels.length > 5000) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "providerSpecificData.excludedModels string must be up to 5000 chars", + path: ["excludedModels"], + }); + } + } else if (!Array.isArray(excludedModels) || excludedModels.length > 100) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "providerSpecificData.excludedModels must be an array with at most 100 items", + path: ["excludedModels"], + }); + } else if ( + excludedModels.some( + (pattern) => + typeof pattern !== "string" || + pattern.trim().length === 0 || + pattern.trim().length > 200 || + pattern.trim() === "**" + ) + ) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: + "providerSpecificData.excludedModels must contain non-empty patterns up to 200 characters", + path: ["excludedModels"], + }); + } + } } // Re-export validation helpers from dedicated module to avoid webpack barrel-file @@ -119,21 +197,47 @@ export type { ValidationResult } from "./helpers"; // ──── Provider Schemas ──── -export const createProviderSchema = z.object({ - provider: z.string().min(1).max(100), - apiKey: z.string().min(1).max(10000), - name: z.string().min(1).max(200), - priority: z.number().int().min(1).max(100).optional(), - globalPriority: z.number().int().min(1).max(100).nullable().optional(), - defaultModel: z.string().max(200).nullable().optional(), - testStatus: z.string().max(50).optional(), - providerSpecificData: z - .record(z.string(), z.unknown()) - .optional() - .superRefine((data, ctx) => { - validateProviderSpecificData(data, ctx); - }), -}); +export const createProviderSchema = z + .object({ + provider: z.string().min(1).max(100), + apiKey: z.string().max(10000).optional(), + name: z.string().min(1).max(200), + priority: z.number().int().min(1).max(100).optional(), + globalPriority: z.number().int().min(1).max(100).nullable().optional(), + defaultModel: z.string().max(200).nullable().optional(), + testStatus: z.string().max(50).optional(), + providerSpecificData: z + .record(z.string(), z.unknown()) + .optional() + .superRefine((data, ctx) => { + validateProviderSpecificData(data, ctx); + }), + }) + .superRefine((data, ctx) => { + const apiKey = typeof data.apiKey === "string" ? data.apiKey.trim() : ""; + if (data.provider !== "searxng-search" && apiKey.length === 0) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "API key is required", + path: ["apiKey"], + }); + } + + const cx = + data.providerSpecificData && typeof data.providerSpecificData === "object" + ? (data.providerSpecificData as Record).cx + : undefined; + if ( + data.provider === "google-pse-search" && + (typeof cx !== "string" || cx.trim().length === 0) + ) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Programmable Search Engine ID (cx) is required", + path: ["providerSpecificData", "cx"], + }); + } + }); // ──── API Key Schemas ──── @@ -433,15 +537,38 @@ export const v1CountTokensSchema = z }) .catchall(z.unknown()); -export const setBudgetSchema = z.object({ - apiKeyId: z.string().trim().min(1, "apiKeyId is required"), - dailyLimitUsd: z.coerce.number().positive("dailyLimitUsd must be greater than zero"), - monthlyLimitUsd: z.coerce - .number() - .positive("monthlyLimitUsd must be greater than zero") - .optional(), - warningThreshold: z.coerce.number().min(0).max(1).optional(), -}); +export const setBudgetSchema = z + .object({ + apiKeyId: z.string().trim().min(1, "apiKeyId is required"), + dailyLimitUsd: z.coerce.number().positive("dailyLimitUsd must be greater than zero").optional(), + weeklyLimitUsd: z.coerce + .number() + .positive("weeklyLimitUsd must be greater than zero") + .optional(), + monthlyLimitUsd: z.coerce + .number() + .positive("monthlyLimitUsd must be greater than zero") + .optional(), + warningThreshold: z.coerce.number().min(0).max(1).optional(), + resetInterval: z.enum(["daily", "weekly", "monthly"]).optional(), + resetTime: z + .string() + .trim() + .regex(/^\d{2}:\d{2}$/, "resetTime must be in HH:MM format") + .optional(), + }) + .superRefine((value, ctx) => { + const hasAnyLimit = [value.dailyLimitUsd, value.weeklyLimitUsd, value.monthlyLimitUsd].some( + (entry) => typeof entry === "number" && Number.isFinite(entry) && entry > 0 + ); + if (!hasAnyLimit) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "At least one budget limit must be provided", + path: ["dailyLimitUsd"], + }); + } + }); export const policyActionSchema = z .object({ @@ -757,6 +884,54 @@ export const updateThinkingBudgetSchema = z } }); +const payloadRuleModelSpecSchema = z + .object({ + name: z.string().trim().min(1), + protocol: z.string().trim().min(1).optional(), + }) + .strict(); + +const payloadMutationRuleSchema = z + .object({ + models: z.array(payloadRuleModelSpecSchema).min(1), + params: z + .record(z.string().trim().min(1), z.unknown()) + .refine((value) => Object.keys(value).length > 0, "params must contain at least one path"), + }) + .strict(); + +const payloadFilterRuleSchema = z + .object({ + models: z.array(payloadRuleModelSpecSchema).min(1), + params: z.array(z.string().trim().min(1)).min(1), + }) + .strict(); + +export const updatePayloadRulesSchema = z + .object({ + default: z.array(payloadMutationRuleSchema).optional(), + override: z.array(payloadMutationRuleSchema).optional(), + filter: z.array(payloadFilterRuleSchema).optional(), + defaultRaw: z.array(payloadMutationRuleSchema).optional(), + "default-raw": z.array(payloadMutationRuleSchema).optional(), + }) + .strict() + .superRefine((value, ctx) => { + if ( + value.default === undefined && + value.override === undefined && + value.filter === undefined && + value.defaultRaw === undefined && + value["default-raw"] === undefined + ) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "No valid fields to update", + path: [], + }); + } + }); + const ipFilterModeSchema = z.enum(["blacklist", "whitelist"]); const tempBanSchema = z.object({ ip: z.string().trim().min(1), @@ -1268,13 +1443,24 @@ export const providersBatchTestSchema = z } }); -export const validateProviderApiKeySchema = z.object({ - provider: z.string().trim().min(1, "Provider and API key required"), - apiKey: z.string().trim().optional(), - validationModelId: z.string().trim().optional(), - customUserAgent: z.string().trim().max(500).optional(), - baseUrl: z.string().trim().url().optional(), -}); +export const validateProviderApiKeySchema = z + .object({ + provider: z.string().trim().min(1, "Provider and API key required"), + apiKey: z.string().trim().optional(), + validationModelId: z.string().trim().optional(), + customUserAgent: z.string().trim().max(500).optional(), + baseUrl: z.string().trim().url().optional(), + cx: z.string().trim().max(500).optional(), + }) + .superRefine((data, ctx) => { + if (data.provider === "google-pse-search" && !data.cx) { + ctx.addIssue({ + code: z.ZodIssueCode.custom, + message: "Programmable Search Engine ID (cx) is required", + path: ["cx"], + }); + } + }); const geminiPartSchema = z .object({ @@ -1404,7 +1590,17 @@ export const v1SearchSchema = z .min(1, "Query is required") .max(500, "Query must be 500 characters or fewer"), provider: z - .enum(["serper-search", "brave-search", "perplexity-search", "exa-search", "tavily-search"]) + .enum([ + "serper-search", + "brave-search", + "perplexity-search", + "exa-search", + "tavily-search", + "google-pse-search", + "linkup-search", + "searchapi-search", + "searxng-search", + ]) .optional(), max_results: z.coerce.number().int().min(1).max(100).default(5), search_type: z.enum(["web", "news"]).default("web"), diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 97aa950d77..3578f7c04c 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -88,6 +88,25 @@ registerCodexQuotaFetcher(); // can proactively switch accounts before quota is exhausted. registerBailianCodingPlanQuotaFetcher(); +function normalizeAllowedConnectionIds(value: unknown): string[] | null { + if (!Array.isArray(value)) return null; + const ids = value.filter( + (entry): entry is string => typeof entry === "string" && entry.trim().length > 0 + ); + return ids.length > 0 ? ids : null; +} + +function intersectAllowedConnectionIds(primary: unknown, secondary: unknown): string[] | null { + const first = normalizeAllowedConnectionIds(primary); + const second = normalizeAllowedConnectionIds(secondary); + + if (first && second) { + return first.filter((id) => second.includes(id)); + } + + return first || second || null; +} + /** * Handle chat completion request * Supports: OpenAI, Claude, Gemini, OpenAI Responses API formats @@ -269,7 +288,7 @@ export async function handleChat(request: any, clientRawRequest: any = null) { // avoid pre-skipping so each model gets a real execution attempt. const checkModelAvailable = async ( modelString: string, - target?: { connectionId?: string | null } + target?: { connectionId?: string | null; allowedConnectionIds?: string[] | null } ) => { if (isComboLiveTest) return true; @@ -281,6 +300,14 @@ export async function handleChat(request: any, clientRawRequest: any = null) { const resolvedModel = modelInfo.model || modelString; const hasForcedConnection = typeof target?.connectionId === "string" && target.connectionId.trim().length > 0; + const allowedConnections = intersectAllowedConnectionIds( + apiKeyInfo?.allowedConnections ?? null, + target?.allowedConnectionIds ?? null + ); + + if (Array.isArray(allowedConnections) && allowedConnections.length === 0) { + return false; + } // Fixed-account combo steps must bypass the provider/model cooldown gate here. // A previous account failure can quarantine the model globally, but the next @@ -293,7 +320,7 @@ export async function handleChat(request: any, clientRawRequest: any = null) { const creds = await getProviderCredentialsWithQuotaPreflight( provider, null, - apiKeyInfo?.allowedConnections ?? null, + allowedConnections, resolvedModel, { ...(target?.connectionId ? { forcedConnectionId: target.connectionId } : {}), @@ -339,6 +366,7 @@ export async function handleChat(request: any, clientRawRequest: any = null) { sessionId, forceLiveComboTest: isComboLiveTest, forcedConnectionId: target?.connectionId ?? null, + allowedConnectionIds: target?.allowedConnectionIds ?? null, comboStepId: target?.stepId || null, comboExecutionKey: target?.executionKey || target?.stepId || null, }, @@ -455,6 +483,7 @@ async function handleSingleModelChat( forceLiveComboTest?: boolean; sessionId?: string | null; forcedConnectionId?: string | null; + allowedConnectionIds?: string[] | null; comboStepId?: string | null; comboExecutionKey?: string | null; } = {}, @@ -470,6 +499,10 @@ async function handleSingleModelChat( const hasForcedConnection = typeof runtimeOptions.forcedConnectionId === "string" && runtimeOptions.forcedConnectionId.trim().length > 0; + const effectiveAllowedConnections = intersectAllowedConnectionIds( + apiKeyInfo?.allowedConnections ?? null, + runtimeOptions.allowedConnectionIds ?? null + ); const bypassReason = forceLiveComboTest ? "combo live test" : hasForcedConnection @@ -508,6 +541,14 @@ async function handleSingleModelChat( : baseRetrySettings; const requestSignal = request?.signal ?? null; + if (Array.isArray(effectiveAllowedConnections) && effectiveAllowedConnections.length === 0) { + log.debug("AUTH", `${provider}/${model} filtered out by connection-level routing constraints`); + return errorResponse( + HTTP_STATUS.SERVICE_UNAVAILABLE, + "No eligible connections matched the requested routing constraints" + ); + } + // 3. Credential retry loop let requestRetryAttempt = 0; let requestRetryLastError = null; @@ -524,7 +565,7 @@ async function handleSingleModelChat( const credentials = await getProviderCredentialsWithQuotaPreflight( provider, null, - apiKeyInfo?.allowedConnections ?? null, + effectiveAllowedConnections, model, { excludeConnectionIds: Array.from(excludedConnectionIds), diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index f2d966bccc..906fd4dbd7 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -25,6 +25,7 @@ import { COOLDOWN_MS } from "@omniroute/open-sse/config/constants.ts"; import { preflightQuota } from "@omniroute/open-sse/services/quotaPreflight.ts"; import { getCodexModelScope } from "@omniroute/open-sse/executors/codex.ts"; import { getProviderAlias, resolveProviderId } from "@/shared/constants/providers"; +import { isModelExcludedByConnection } from "@/domain/connectionModelRules"; import * as log from "../utils/logger"; import { fisherYatesShuffle, getNextFromDeckSync } from "@/shared/utils/shuffleDeck"; @@ -668,6 +669,9 @@ export async function getProviderCredentials( // Filter out unavailable accounts and excluded connection const availableConnections = connections.filter((c) => { if (excludedConnectionIds.has(c.id)) return false; + if (requestedModel && isModelExcludedByConnection(requestedModel, c.providerSpecificData)) { + return false; + } if (!allowSuppressedConnections) { if (isAccountUnavailable(c.rateLimitedUntil)) return false; if (isTerminalConnectionStatus(c)) return false; @@ -689,11 +693,19 @@ export async function getProviderCredentials( const codexScopeLimited = provider === "codex" && isCodexScopeUnavailable(c, requestedModel); const modelLocked = Boolean(requestedModel) && isModelLocked(provider, c.id, requestedModel as string); + const modelExcluded = + Boolean(requestedModel) && + isModelExcludedByConnection(requestedModel as string, c.providerSpecificData); if (excluded || rateLimited) { log.debug( "AUTH", ` → ${c.id?.slice(0, 8)} | ${excluded ? "excluded" : ""} ${rateLimited ? `rateLimited until ${c.rateLimitedUntil}` : ""}${allowSuppressedConnections && rateLimited ? " (retained for combo live test)" : ""}` ); + } else if (modelExcluded) { + log.debug( + "AUTH", + ` → ${c.id?.slice(0, 8)} | excluded by per-account model rule for ${requestedModel}` + ); } else if (terminalStatus) { log.debug( "AUTH", diff --git a/tests/unit/catalog-updates-v3x.test.ts b/tests/unit/catalog-updates-v3x.test.ts new file mode 100644 index 0000000000..606578fbad --- /dev/null +++ b/tests/unit/catalog-updates-v3x.test.ts @@ -0,0 +1,49 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { getModelsByProviderId } from "../../open-sse/config/providerModels.ts"; +import { resolveCanonicalProviderModel } from "../../open-sse/services/model.ts"; + +test("Pollinations catalog mirrors the current public text model lineup", () => { + const models = getModelsByProviderId("pollinations"); + const ids = new Set(models.map((model) => model.id)); + const names = models.map((model) => model.name); + + assert.ok(ids.has("openai-fast")); + assert.ok(ids.has("openai-large")); + assert.ok(ids.has("perplexity-fast")); + assert.ok(ids.has("qwen-coder-large")); + assert.ok(ids.has("claude-large")); + assert.equal(ids.has("llama"), false); + assert.equal( + names.some((name) => /GPT-5 via Pollinations/i.test(name)), + false + ); +}); + +test("Puter catalog exposes the currently documented Sonar models", () => { + const ids = new Set(getModelsByProviderId("puter").map((model) => model.id)); + + assert.ok(ids.has("perplexity/sonar")); + assert.ok(ids.has("perplexity/sonar-pro")); + assert.ok(ids.has("perplexity/sonar-pro-search")); + assert.ok(ids.has("perplexity/sonar-reasoning-pro")); + assert.ok(ids.has("perplexity/sonar-deep-research")); +}); + +test("NVIDIA catalog includes the verified 2026 additions and GPT OSS 20B alias resolution", () => { + const ids = new Set(getModelsByProviderId("nvidia").map((model) => model.id)); + + assert.ok(ids.has("openai/gpt-oss-20b")); + assert.ok(ids.has("nvidia/llama-3.1-nemotron-ultra-253b-v1")); + assert.ok(ids.has("nvidia/nemotron-3-super-120b-a12b")); + assert.ok(ids.has("nvidia/llama-3.3-nemotron-super-49b-v1.5")); + assert.ok(ids.has("mistralai/mistral-large-3-675b-instruct-2512")); + assert.ok(ids.has("qwen/qwen3-coder-480b-a35b-instruct")); + assert.ok(ids.has("mistralai/devstral-2-123b-instruct-2512")); + + assert.deepEqual(resolveCanonicalProviderModel("nvidia", "gpt-oss-20b"), { + provider: "nvidia", + model: "openai/gpt-oss-20b", + }); +}); diff --git a/tests/unit/chatcore-translation-paths.test.ts b/tests/unit/chatcore-translation-paths.test.ts index 8f23a87516..5f60ff4acb 100644 --- a/tests/unit/chatcore-translation-paths.test.ts +++ b/tests/unit/chatcore-translation-paths.test.ts @@ -31,6 +31,8 @@ const { isTokenExpiringSoon, clearUpstreamProxyConfigCache, } = await import("../../open-sse/handlers/chatCore.ts"); +const { resetPayloadRulesConfigForTests, setPayloadRulesConfig } = + await import("../../open-sse/services/payloadRules.ts"); const { FORMATS } = await import("../../open-sse/translator/formats.ts"); const { register, getRequestTranslator } = await import("../../open-sse/translator/registry.ts"); @@ -232,6 +234,7 @@ function collectTextBlocks(messages) { async function resetStorage() { clearUpstreamProxyConfigCache(); + resetPayloadRulesConfigForTests(); register(FORMATS.OPENAI_RESPONSES, FORMATS.OPENAI, originalResponsesToOpenAI, null); invalidateCacheControlSettingsCache(); clearCache(); @@ -437,6 +440,52 @@ test("chatCore helper exports detect responses passthrough paths and token expir assert.equal(isTokenExpiringSoon(null), false); }); +test("chatCore applies payload rules after translating Responses input into Chat payloads", async () => { + setPayloadRulesConfig({ + default: [ + { + models: [{ name: "gpt-*", protocol: "openai" }], + params: { + "messages.0.metadata.routeTag": "feature-110", + }, + }, + ], + override: [ + { + models: [{ name: "gpt-*", protocol: "openai" }], + params: { + temperature: 0.25, + }, + }, + ], + filter: [], + defaultRaw: [], + }); + + const { call, result } = await invokeChatCore({ + provider: "openai", + model: "gpt-4o-mini", + endpoint: "/v1/responses", + body: { + model: "gpt-4o-mini", + stream: false, + input: [ + { + type: "message", + role: "user", + content: [{ type: "input_text", text: "hello" }], + }, + ], + }, + responseFormat: "openai", + }); + + assert.equal(result.success, true); + assert.equal(call.body.temperature, 0.25); + assert.equal(call.body.messages[0].metadata.routeTag, "feature-110"); + assert.equal(call.body.messages[0].role, "user"); +}); + test("chatCore builds Claude Code-compatible upstream requests for CC providers", async () => { const { call, result } = await invokeChatCore({ provider: "anthropic-compatible-cc-test", @@ -1172,6 +1221,28 @@ test("chatCore skips semantic cache when disabled in settings", async () => { assert.equal(payload.choices[0].message.content, "fresh-2"); }); +test("chatCore attaches OmniRoute response metadata headers to non-stream responses", async () => { + const { result } = await invokeChatCore({ + provider: "claude", + model: "claude-sonnet-4-6", + body: { + model: "claude-sonnet-4-6", + stream: false, + messages: [{ role: "user", content: "header metadata" }], + }, + responseFormat: "claude", + }); + + assert.equal(result.success, true); + assert.equal(result.response.headers.get("X-OmniRoute-Provider"), "cc"); + assert.equal(result.response.headers.get("X-OmniRoute-Model"), "claude-sonnet-4-6"); + assert.equal(result.response.headers.get("X-OmniRoute-Cache-Hit"), "false"); + assert.equal(result.response.headers.get("X-OmniRoute-Tokens-In"), "12"); + assert.equal(result.response.headers.get("X-OmniRoute-Tokens-Out"), "3"); + assert.ok(Number(result.response.headers.get("X-OmniRoute-Latency-Ms")) >= 0); + assert.match(String(result.response.headers.get("X-OmniRoute-Response-Cost")), /^\d+\.\d{10}$/); +}); + test("chatCore normalizes tool finish reasons and estimates usage when upstream omits it", async () => { const { result } = await invokeChatCore({ provider: "openai", @@ -1760,6 +1831,34 @@ test("chatCore injects progress events into streaming responses when requested", assert.match(streamText, /event: progress/); }); +test("chatCore emits final SSE metadata comments before [DONE] on streaming responses", async () => { + const { result } = await invokeChatCore({ + provider: "openai", + model: "gpt-4o-mini", + accept: "text/event-stream", + body: { + model: "gpt-4o-mini", + stream: true, + messages: [{ role: "user", content: "stream metadata" }], + }, + responseFactory() { + return buildOpenAIResponse(true, "streamed"); + }, + }); + + const streamText = await result.response.text(); + + assert.equal(result.success, true); + assert.equal(result.response.headers.get("X-OmniRoute-Provider"), "openai"); + assert.equal(result.response.headers.get("X-OmniRoute-Model"), "gpt-4o-mini"); + assert.match(streamText, /: x-omniroute-response-cost=\d+\.\d{10}/); + assert.match(streamText, /: x-omniroute-tokens-in=\d+/); + assert.match(streamText, /: x-omniroute-tokens-out=\d+/); + assert.ok( + streamText.indexOf(": x-omniroute-response-cost=") < streamText.indexOf("data: [DONE]") + ); +}); + test("chatCore maps upstream aborts to request-aborted errors", async () => { const { result } = await invokeChatCore({ provider: "openai", diff --git a/tests/unit/connection-model-rules.test.ts b/tests/unit/connection-model-rules.test.ts new file mode 100644 index 0000000000..d6bec2ac3f --- /dev/null +++ b/tests/unit/connection-model-rules.test.ts @@ -0,0 +1,47 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + getConnectionExcludedModels, + hasEligibleConnectionForModel, + isModelExcludedByConnection, + normalizeExcludedModelPatterns, +} from "../../src/domain/connectionModelRules.ts"; + +test("normalizeExcludedModelPatterns accepts arrays and CSV strings, trims values, and removes duplicates", () => { + assert.deepEqual(normalizeExcludedModelPatterns([" gpt-4o* ", "gpt-4o*", "", "**"]), ["gpt-4o*"]); + assert.deepEqual(normalizeExcludedModelPatterns("gpt-4.1*, gpt-4o*, gpt-4.1*"), [ + "gpt-4.1*", + "gpt-4o*", + ]); +}); + +test("isModelExcludedByConnection matches both provider-scoped and raw model ids", () => { + const providerSpecificData = { + excludedModels: ["gpt-4o*", "claude-opus-*"], + }; + + assert.equal(isModelExcludedByConnection("openai/gpt-4o-mini", providerSpecificData), true); + assert.equal(isModelExcludedByConnection("gpt-4o-mini", providerSpecificData), true); + assert.equal(isModelExcludedByConnection("claude-sonnet-4-5", providerSpecificData), false); +}); + +test("hasEligibleConnectionForModel only returns true when at least one connection can serve the model", () => { + const connections = [ + { providerSpecificData: { excludedModels: ["gpt-4o*"] } }, + { providerSpecificData: { excludedModels: ["gpt-4.1*"] } }, + ]; + + assert.equal(hasEligibleConnectionForModel(connections, "gpt-4o-mini"), true); + assert.equal( + hasEligibleConnectionForModel( + [ + { providerSpecificData: { excludedModels: ["gpt-4o*"] } }, + { providerSpecificData: { excludedModels: ["gpt-4o-mini"] } }, + ], + "gpt-4o-mini" + ), + false + ); + assert.deepEqual(getConnectionExcludedModels({ excluded_models: ["gpt-4o*"] }), ["gpt-4o*"]); +}); diff --git a/tests/unit/domain-cost-rules.test.ts b/tests/unit/domain-cost-rules.test.ts index 62997d3115..3fdb2bf148 100644 --- a/tests/unit/domain-cost-rules.test.ts +++ b/tests/unit/domain-cost-rules.test.ts @@ -48,8 +48,15 @@ test("setBudget normalizes defaults and getBudget returns the stored config", () assert.deepEqual(costRules.getBudget("key-budget"), { dailyLimitUsd: 12.5, + weeklyLimitUsd: 0, monthlyLimitUsd: 0, warningThreshold: 0.8, + resetInterval: "daily", + resetTime: "00:00", + budgetResetAt: costRules.getBudget("key-budget")?.budgetResetAt ?? null, + lastBudgetResetAt: costRules.getBudget("key-budget")?.lastBudgetResetAt ?? null, + warningEmittedAt: null, + warningPeriodStart: null, }); assert.equal(costRules.getBudget("missing-key"), null); }); @@ -69,6 +76,14 @@ test("checkBudget reports warning and blocks when projected spend exceeds the da dailyUsed: 5, dailyLimit: 10, warningReached: true, + remaining: 4, + periodUsed: 5, + activeLimitUsd: 10, + resetInterval: "daily", + resetTime: "00:00", + budgetResetAt: warning.budgetResetAt, + lastBudgetResetAt: warning.lastBudgetResetAt, + periodStartAt: warning.periodStartAt, }); assert.equal(denied.allowed, false); assert.equal(denied.warningReached, true); @@ -99,22 +114,55 @@ test("getDailyTotal and getCostSummary split daily and monthly totals correctly" totalEntries: 2, budget: { dailyLimitUsd: 50, + weeklyLimitUsd: 0, monthlyLimitUsd: 100, warningThreshold: 0.75, + resetInterval: "daily", + resetTime: "00:00", + budgetResetAt: costRules.getBudget("key-summary")?.budgetResetAt ?? null, + lastBudgetResetAt: costRules.getBudget("key-summary")?.lastBudgetResetAt ?? null, + warningEmittedAt: null, + warningPeriodStart: null, }, + totalCostToday: 2.5, + totalCostMonth: 4, + totalCostPeriod: 2.5, + activeLimitUsd: 50, + resetInterval: "daily", + resetTime: "00:00", + budgetResetAt: costRules.getBudget("key-summary")?.budgetResetAt ?? null, + lastBudgetResetAt: costRules.getBudget("key-summary")?.lastBudgetResetAt ?? null, + periodStartAt: costRules.getBudget("key-summary")?.lastBudgetResetAt ?? null, + nextResetAt: costRules.getBudget("key-summary")?.budgetResetAt ?? null, + dailyLimitUsd: 50, + weeklyLimitUsd: 0, + monthlyLimitUsd: 100, + warningThreshold: 0.75, }); }); test("costRules covers DB-loaded budgets, malformed entries and storage failure fallbacks", () => { domainState.saveBudget("db-loaded", { dailyLimitUsd: 7, + weeklyLimitUsd: 14, monthlyLimitUsd: 21, warningThreshold: 0.7, + resetInterval: "weekly", + resetTime: "06:30", + budgetResetAt: 111, + lastBudgetResetAt: 99, }); assert.deepEqual(costRules.getBudget("db-loaded"), { dailyLimitUsd: 7, + weeklyLimitUsd: 14, monthlyLimitUsd: 21, warningThreshold: 0.7, + resetInterval: "weekly", + resetTime: "06:30", + budgetResetAt: costRules.getBudget("db-loaded")?.budgetResetAt ?? null, + lastBudgetResetAt: costRules.getBudget("db-loaded")?.lastBudgetResetAt ?? null, + warningEmittedAt: null, + warningPeriodStart: null, }); assert.deepEqual(costRules.checkBudget("missing-budget"), { @@ -122,6 +170,14 @@ test("costRules covers DB-loaded budgets, malformed entries and storage failure dailyUsed: 0, dailyLimit: 0, warningReached: false, + remaining: 0, + periodUsed: 0, + activeLimitUsd: 0, + resetInterval: null, + resetTime: null, + budgetResetAt: null, + lastBudgetResetAt: null, + periodStartAt: null, }); const db = core.getDbInstance(); @@ -148,6 +204,20 @@ test("costRules covers DB-loaded budgets, malformed entries and storage failure monthlyTotal: 2.25, totalEntries: 1, budget: null, + totalCostToday: 2.25, + totalCostMonth: 2.25, + totalCostPeriod: 0, + activeLimitUsd: 0, + resetInterval: null, + resetTime: null, + budgetResetAt: null, + lastBudgetResetAt: null, + periodStartAt: null, + nextResetAt: null, + dailyLimitUsd: 0, + weeklyLimitUsd: 0, + monthlyLimitUsd: 0, + warningThreshold: null, }); db.exec("DROP TABLE domain_cost_history"); @@ -157,5 +227,66 @@ test("costRules covers DB-loaded budgets, malformed entries and storage failure monthlyTotal: 0, totalEntries: 0, budget: null, + totalCostToday: 0, + totalCostMonth: 0, + totalCostPeriod: 0, + activeLimitUsd: 0, + resetInterval: null, + resetTime: null, + budgetResetAt: null, + lastBudgetResetAt: null, + periodStartAt: null, + nextResetAt: null, + dailyLimitUsd: 0, + weeklyLimitUsd: 0, + monthlyLimitUsd: 0, + warningThreshold: null, }); }); + +test("weekly budgets use the weekly window limit and expose the next reset metadata", () => { + costRules.setBudget("key-weekly", { + dailyLimitUsd: 5, + weeklyLimitUsd: 20, + resetInterval: "weekly", + resetTime: "06:30", + }); + + const budget = costRules.getBudget("key-weekly"); + const summary = costRules.getCostSummary("key-weekly"); + const check = costRules.checkBudget("key-weekly", 0); + + assert.equal(budget?.resetInterval, "weekly"); + assert.equal(budget?.resetTime, "06:30"); + assert.equal(summary.activeLimitUsd, 20); + assert.equal(summary.resetInterval, "weekly"); + assert.equal(check.dailyLimit, 20); + assert.ok(typeof summary.budgetResetAt === "number" && summary.budgetResetAt > Date.now()); +}); + +test("syncAllBudgetSchedules advances overdue budgets and records a reset log", () => { + const now = Date.UTC(2026, 3, 17, 12, 0, 0); + const previousPeriodStart = Date.UTC(2026, 3, 15, 0, 0, 0); + const overdueResetAt = Date.UTC(2026, 3, 16, 0, 0, 0); + + domainState.saveBudget("key-reset", { + dailyLimitUsd: 10, + warningThreshold: 0.8, + resetInterval: "daily", + resetTime: "00:00", + budgetResetAt: overdueResetAt, + lastBudgetResetAt: previousPeriodStart, + }); + domainState.saveCostEntry("key-reset", 3.5, Date.UTC(2026, 3, 15, 12, 0, 0)); + + const result = costRules.syncAllBudgetSchedules(now); + const synced = costRules.getBudget("key-reset"); + const logs = domainState.loadBudgetResetLogs("key-reset", 5); + + assert.equal(result.processed, 1); + assert.equal(result.resetCount, 1); + assert.equal(synced?.lastBudgetResetAt, Date.UTC(2026, 3, 17, 0, 0, 0)); + assert.equal(logs.length, 1); + assert.equal(logs[0].previousSpend, 3.5); + assert.equal(logs[0].resetInterval, "daily"); +}); diff --git a/tests/unit/executor-pollinations.test.ts b/tests/unit/executor-pollinations.test.ts index d551b9f57c..bf21aaed98 100644 --- a/tests/unit/executor-pollinations.test.ts +++ b/tests/unit/executor-pollinations.test.ts @@ -3,24 +3,32 @@ import assert from "node:assert/strict"; import { PollinationsExecutor } from "../../open-sse/executors/pollinations.ts"; -test("PollinationsExecutor.buildUrl uses the free Pollinations endpoint", () => { +test("PollinationsExecutor.buildUrl uses the primary endpoint and the gen fallback", () => { const executor = new PollinationsExecutor(); assert.equal( executor.buildUrl("openai", true), "https://text.pollinations.ai/openai/chat/completions" ); + assert.equal( + executor.buildUrl("openai", true, 1), + "https://gen.pollinations.ai/v1/chat/completions" + ); }); -test("PollinationsExecutor.buildHeaders requires an API key", () => { +test("PollinationsExecutor.buildHeaders supports anonymous access and optional SSE accept", () => { const executor = new PollinationsExecutor(); - assert.throws(() => executor.buildHeaders({}, true), /Pollinations API key is required/); + assert.deepEqual(executor.buildHeaders({}, true), { + "Content-Type": "application/json", + Accept: "text/event-stream", + }); }); -test("PollinationsExecutor.buildHeaders sends API auth for the Pollinations key-backed tier", () => { +test("PollinationsExecutor.buildHeaders sends API auth for the key-backed tier when configured", () => { const executor = new PollinationsExecutor(); - assert.deepEqual(executor.buildHeaders({ apiKey: "poll-key" }, false), { + assert.deepEqual(executor.buildHeaders({ apiKey: "poll-key" }, true), { "Content-Type": "application/json", Authorization: "Bearer poll-key", + Accept: "text/event-stream", }); }); diff --git a/tests/unit/model-catalog-search.test.ts b/tests/unit/model-catalog-search.test.ts new file mode 100644 index 0000000000..b0bec7e9f8 --- /dev/null +++ b/tests/unit/model-catalog-search.test.ts @@ -0,0 +1,40 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + getModelCatalogSourceLabel, + matchesModelCatalogQuery, + normalizeModelCatalogSource, +} from "../../src/shared/utils/modelCatalogSearch.ts"; + +test("model catalog source normalization groups manual/imported rows as custom", () => { + assert.equal(normalizeModelCatalogSource("manual"), "custom"); + assert.equal(normalizeModelCatalogSource("imported"), "custom"); + assert.equal(normalizeModelCatalogSource("api-sync"), "api-sync"); + assert.equal(normalizeModelCatalogSource("fallback"), "fallback"); + assert.equal(normalizeModelCatalogSource("alias"), "alias"); + assert.equal(normalizeModelCatalogSource(undefined), "system"); +}); + +test("model catalog source labels stay user-facing", () => { + assert.equal(getModelCatalogSourceLabel("system"), "Built-in"); + assert.equal(getModelCatalogSourceLabel("custom"), "Custom"); + assert.equal(getModelCatalogSourceLabel("api-sync"), "Synced"); + assert.equal(getModelCatalogSourceLabel("fallback"), "Fallback"); + assert.equal(getModelCatalogSourceLabel("alias"), "Alias"); +}); + +test("model catalog query matches id, display name, alias and source label", () => { + const target = { + modelId: "qwen/qwen3-coder-480b-a35b-instruct", + modelName: "Qwen3 Coder 480B", + alias: "best-qwen", + source: "api-sync", + }; + + assert.equal(matchesModelCatalogQuery("", target), true); + assert.equal(matchesModelCatalogQuery("coder-480b", target), true); + assert.equal(matchesModelCatalogQuery("best-qwen", target), true); + assert.equal(matchesModelCatalogQuery("synced", target), true); + assert.equal(matchesModelCatalogQuery("built-in", target), false); +}); diff --git a/tests/unit/models-catalog-route.test.ts b/tests/unit/models-catalog-route.test.ts index 47583fa035..1924043f42 100644 --- a/tests/unit/models-catalog-route.test.ts +++ b/tests/unit/models-catalog-route.test.ts @@ -98,6 +98,51 @@ test("v1 models catalog accepts bearer API keys and filters the list by allowed ); }); +test("v1 models catalog hides models excluded by every active connection while keeping models served by at least one account", async () => { + const first = await seedConnection("openai", { + name: "openai-first", + providerSpecificData: { + excludedModels: ["gpt-4o*"], + }, + }); + const second = await seedConnection("openai", { + name: "openai-second", + providerSpecificData: { + excludedModels: ["gpt-4.1*"], + }, + }); + + let response = await v1ModelsCatalog.getUnifiedModelsResponse( + new Request("http://localhost/api/v1/models") + ); + let body = await response.json(); + let ids = new Set(body.data.map((item) => item.id)); + + assert.equal(response.status, 200); + assert.equal(ids.has("openai/gpt-4o-mini"), true); + + await providersDb.updateProviderConnection(second.id, { + providerSpecificData: { + excludedModels: ["gpt-4o*"], + }, + }); + + response = await v1ModelsCatalog.getUnifiedModelsResponse( + new Request("http://localhost/api/v1/models") + ); + body = await response.json(); + ids = new Set(body.data.map((item) => item.id)); + + assert.equal(response.status, 200); + assert.equal(ids.has("openai/gpt-4o-mini"), false); + + await providersDb.updateProviderConnection(first.id, { + providerSpecificData: { + excludedModels: [], + }, + }); +}); + test("v1 models catalog includes combos and custom models while excluding hidden models and blocked providers", async () => { await settingsDb.updateSettings({ blockedProviders: ["claude"], diff --git a/tests/unit/omniroute-response-meta.test.ts b/tests/unit/omniroute-response-meta.test.ts new file mode 100644 index 0000000000..8bb197df6a --- /dev/null +++ b/tests/unit/omniroute-response-meta.test.ts @@ -0,0 +1,68 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + buildOmniRouteResponseMetaHeaders, + buildOmniRouteSseMetadataComment, + formatOmniRouteCost, + getOmniRouteTokenCounts, +} from "../../src/domain/omnirouteResponseMeta.ts"; + +test("getOmniRouteTokenCounts normalizes common usage shapes", () => { + assert.deepEqual( + getOmniRouteTokenCounts({ + prompt_tokens: 12, + completion_tokens: 5, + }), + { input: 12, output: 5 } + ); + assert.deepEqual( + getOmniRouteTokenCounts({ + input_tokens: "9", + output_tokens: "4", + }), + { input: 9, output: 4 } + ); +}); + +test("buildOmniRouteResponseMetaHeaders formats provider alias, tokens, latency, and cost", () => { + const headers = buildOmniRouteResponseMetaHeaders({ + provider: "claude", + model: "claude-sonnet-4-6", + cacheHit: true, + latencyMs: 1234.6, + usage: { + prompt_tokens: 11, + completion_tokens: 7, + }, + costUsd: 0.00123456789, + }); + + assert.equal(headers["X-OmniRoute-Provider"], "cc"); + assert.equal(headers["X-OmniRoute-Model"], "claude-sonnet-4-6"); + assert.equal(headers["X-OmniRoute-Cache-Hit"], "true"); + assert.equal(headers["X-OmniRoute-Latency-Ms"], "1235"); + assert.equal(headers["X-OmniRoute-Tokens-In"], "11"); + assert.equal(headers["X-OmniRoute-Tokens-Out"], "7"); + assert.equal(headers["X-OmniRoute-Response-Cost"], "0.0012345679"); +}); + +test("buildOmniRouteSseMetadataComment emits comment lines compatible with SSE", () => { + const comment = buildOmniRouteSseMetadataComment({ + provider: "openai", + model: "gpt-4o-mini", + usage: { + prompt_tokens: 4, + completion_tokens: 2, + }, + latencyMs: 50, + costUsd: formatOmniRouteCost(0), + }); + + assert.match(comment, /^: x-omniroute-cache-hit=false/m); + assert.match(comment, /^: x-omniroute-provider=openai/m); + assert.match(comment, /^: x-omniroute-model=gpt-4o-mini/m); + assert.match(comment, /^: x-omniroute-tokens-in=4/m); + assert.match(comment, /^: x-omniroute-tokens-out=2/m); + assert.match(comment, /^: x-omniroute-response-cost=0\.0000000000/m); +}); diff --git a/tests/unit/payload-rules.test.ts b/tests/unit/payload-rules.test.ts new file mode 100644 index 0000000000..29ccc032b9 --- /dev/null +++ b/tests/unit/payload-rules.test.ts @@ -0,0 +1,142 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const { applyPayloadRules, getPayloadRulesConfig, resetPayloadRulesConfigForTests } = + await import("../../open-sse/services/payloadRules.ts"); + +const ORIGINAL_PAYLOAD_RULES_PATH = process.env.OMNIROUTE_PAYLOAD_RULES_PATH; +const ORIGINAL_PAYLOAD_RULES_RELOAD_MS = process.env.OMNIROUTE_PAYLOAD_RULES_RELOAD_MS; + +test.afterEach(() => { + resetPayloadRulesConfigForTests(); + + if (ORIGINAL_PAYLOAD_RULES_PATH === undefined) { + delete process.env.OMNIROUTE_PAYLOAD_RULES_PATH; + } else { + process.env.OMNIROUTE_PAYLOAD_RULES_PATH = ORIGINAL_PAYLOAD_RULES_PATH; + } + + if (ORIGINAL_PAYLOAD_RULES_RELOAD_MS === undefined) { + delete process.env.OMNIROUTE_PAYLOAD_RULES_RELOAD_MS; + } else { + process.env.OMNIROUTE_PAYLOAD_RULES_RELOAD_MS = ORIGINAL_PAYLOAD_RULES_RELOAD_MS; + } +}); + +test("payload rules apply default, default-raw, override, and filter operations", () => { + const { payload, applied } = applyPayloadRules( + { + temperature: 0.4, + metadata: { removeMe: true }, + dangerous: { enabled: true }, + }, + "gpt-5.4", + "openai", + { + default: [ + { + models: [{ name: "gpt-*", protocol: "openai" }], + params: { + "metadata.routeTag": "feature-110", + }, + }, + ], + defaultRaw: [ + { + models: [{ name: "gpt-*", protocol: "openai" }], + params: { + response_format: + '{"type":"json_schema","json_schema":{"name":"shape","schema":{"type":"object"}}}', + }, + }, + ], + override: [ + { + models: [{ name: "gpt-*", protocol: "openai" }], + params: { + "reasoning.effort": "high", + }, + }, + ], + filter: [ + { + models: [{ name: "gpt-*", protocol: "openai" }], + params: ["dangerous", "metadata.removeMe"], + }, + ], + } + ); + + assert.equal(payload.temperature, 0.4); + assert.equal(payload.metadata.routeTag, "feature-110"); + assert.equal("removeMe" in payload.metadata, false); + assert.equal("dangerous" in payload, false); + assert.equal(payload.reasoning.effort, "high"); + assert.deepEqual(payload.response_format, { + type: "json_schema", + json_schema: { + name: "shape", + schema: { type: "object" }, + }, + }); + assert.deepEqual( + applied.map((rule) => `${rule.type}:${rule.path}`), + [ + "default:metadata.routeTag", + "default-raw:response_format", + "override:reasoning.effort", + "filter:dangerous", + "filter:metadata.removeMe", + ] + ); +}); + +test("payload rules load from JSON file and reload changed content", async () => { + const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-payload-rules-")); + const configPath = path.join(tempDir, "payloadRules.json"); + + process.env.OMNIROUTE_PAYLOAD_RULES_PATH = configPath; + + fs.writeFileSync( + configPath, + JSON.stringify({ + override: [ + { + models: [{ name: "gpt-*", protocol: "openai" }], + params: { temperature: 0.1 }, + }, + ], + }), + "utf-8" + ); + + const first = await getPayloadRulesConfig({ forceRefresh: true }); + assert.equal(first.override.length, 1); + assert.equal(first.override[0].params.temperature, 0.1); + assert.equal(first.defaultRaw.length, 0); + + await new Promise((resolve) => setTimeout(resolve, 20)); + + fs.writeFileSync( + configPath, + JSON.stringify({ + "default-raw": [ + { + models: [{ name: "gpt-*", protocol: "openai" }], + params: { response_format: { type: "json_object" } }, + }, + ], + }), + "utf-8" + ); + + const second = await getPayloadRulesConfig({ forceRefresh: true }); + assert.equal(second.override.length, 0); + assert.equal(second.defaultRaw.length, 1); + assert.deepEqual(second.defaultRaw[0].params.response_format, { type: "json_object" }); + + fs.rmSync(tempDir, { recursive: true, force: true }); +}); diff --git a/tests/unit/provider-validation-specialty.test.ts b/tests/unit/provider-validation-specialty.test.ts index dab5f7ee6a..44a6fa7470 100644 --- a/tests/unit/provider-validation-specialty.test.ts +++ b/tests/unit/provider-validation-specialty.test.ts @@ -122,6 +122,73 @@ test("search provider validators cover success, client errors, server errors and assert.equal(calls[0].init.headers["User-Agent"], "SearchSuite/1.0"); }); +test("extended search provider validators cover Google PSE, Linkup, SearchAPI and SearXNG", async () => { + const originalAllowPrivateProviderUrls = process.env.OMNIROUTE_ALLOW_PRIVATE_PROVIDER_URLS; + process.env.OMNIROUTE_ALLOW_PRIVATE_PROVIDER_URLS = "true"; + const calls = []; + try { + globalThis.fetch = async (url, init = {}) => { + calls.push({ url: String(url), init }); + const target = String(url); + if (target.startsWith("https://www.googleapis.com/customsearch/v1")) { + return new Response(JSON.stringify({ items: [] }), { status: 200 }); + } + if (target.startsWith("https://api.linkup.so/v1/search")) { + return new Response(JSON.stringify({ error: "bad request" }), { status: 400 }); + } + if (target.startsWith("https://www.searchapi.io/api/v1/search")) { + return new Response(JSON.stringify({ organic_results: [] }), { status: 200 }); + } + if (target.startsWith("http://localhost:9999/search")) { + return new Response(JSON.stringify({ results: [] }), { status: 200 }); + } + throw new Error(`unexpected fetch: ${target}`); + }; + + const google = await validateProviderApiKey({ + provider: "google-pse-search", + apiKey: "google-key", + providerSpecificData: { cx: "engine-id" }, + }); + const linkup = await validateProviderApiKey({ + provider: "linkup-search", + apiKey: "linkup-key", + }); + const searchapi = await validateProviderApiKey({ + provider: "searchapi-search", + apiKey: "searchapi-key", + }); + const searxng = await validateProviderApiKey({ + provider: "searxng-search", + providerSpecificData: { baseUrl: "http://localhost:9999/search" }, + }); + + assert.equal(google.valid, true); + assert.equal(linkup.valid, true); + assert.equal(searchapi.valid, true); + assert.equal(searxng.valid, true); + assert.match(calls[0].url, /cx=engine-id/); + assert.equal(calls[1].init.headers.Authorization, "Bearer linkup-key"); + assert.match(calls[2].url, /api_key=searchapi-key/); + } finally { + if (originalAllowPrivateProviderUrls === undefined) { + delete process.env.OMNIROUTE_ALLOW_PRIVATE_PROVIDER_URLS; + } else { + process.env.OMNIROUTE_ALLOW_PRIVATE_PROVIDER_URLS = originalAllowPrivateProviderUrls; + } + } +}); + +test("google PSE validator requires cx", async () => { + const result = await validateProviderApiKey({ + provider: "google-pse-search", + apiKey: "google-key", + }); + + assert.equal(result.valid, false); + assert.equal(result.error, "Programmable Search Engine ID (cx) is required"); +}); + test("OpenAI-compatible validator covers /responses mode and final ping fallback", async () => { const calls = []; globalThis.fetch = async (url, init = {}) => { diff --git a/tests/unit/search-handler-extended.test.ts b/tests/unit/search-handler-extended.test.ts index ccc25fb88d..a725580595 100644 --- a/tests/unit/search-handler-extended.test.ts +++ b/tests/unit/search-handler-extended.test.ts @@ -288,6 +288,216 @@ test("handleSearch builds Tavily requests with topic and raw content normalizati } }); +test("handleSearch builds Google PSE requests with key/cx query params and normalizes items", async () => { + const originalFetch = globalThis.fetch; + let capturedUrl; + + globalThis.fetch = async (url) => { + capturedUrl = String(url); + + return new Response( + JSON.stringify({ + items: [ + { + title: "Programmable result", + link: "https://google.example.com/page", + snippet: "Google snippet", + pagemap: { + cse_image: [{ src: "https://google.example.com/image.png" }], + }, + }, + ], + searchInformation: { totalResults: "42" }, + }), + { status: 200, headers: { "content-type": "application/json" } } + ); + }; + + try { + const result = await handleSearch({ + query: "google custom search", + provider: "google-pse-search", + maxResults: 3, + searchType: "web", + country: "US", + language: "en", + credentials: { apiKey: "google-key", providerSpecificData: { cx: "engine-id" } }, + log: null, + }); + + const url = new URL(capturedUrl); + assert.equal(url.origin + url.pathname, "https://www.googleapis.com/customsearch/v1"); + assert.equal(url.searchParams.get("key"), "google-key"); + assert.equal(url.searchParams.get("cx"), "engine-id"); + assert.equal(url.searchParams.get("q"), "google custom search"); + assert.equal(url.searchParams.get("num"), "3"); + assert.equal(result.success, true); + assert.equal(result.data.results[0].metadata.image_url, "https://google.example.com/image.png"); + assert.equal(result.data.metrics.total_results_available, 42); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleSearch builds Linkup requests and normalizes searchResults payload", async () => { + const originalFetch = globalThis.fetch; + let captured; + + globalThis.fetch = async (url, init = {}) => { + captured = { + url: String(url), + headers: init.headers, + body: JSON.parse(String(init.body || "{}")), + }; + + return new Response( + JSON.stringify({ + results: [ + { + name: "Linkup result", + url: "https://linkup.example.com/page", + content: "Retrieved content", + type: "web", + }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } } + ); + }; + + try { + const result = await handleSearch({ + query: "linkup test", + provider: "linkup-search", + maxResults: 4, + searchType: "web", + domainFilter: ["example.com", "-blocked.com"], + credentials: { apiKey: "linkup-key" }, + log: null, + }); + + assert.equal(captured.url, "https://api.linkup.so/v1/search"); + assert.equal(captured.headers.Authorization, "Bearer linkup-key"); + assert.deepEqual(captured.body, { + q: "linkup test", + depth: "standard", + outputType: "searchResults", + maxResults: 4, + includeDomains: ["example.com"], + excludeDomains: ["blocked.com"], + }); + assert.equal(result.success, true); + assert.equal(result.data.results[0].title, "Linkup result"); + assert.equal(result.data.results[0].content.text, "Retrieved content"); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleSearch builds SearchAPI requests and normalizes organic results", async () => { + const originalFetch = globalThis.fetch; + let capturedUrl; + + globalThis.fetch = async (url) => { + capturedUrl = String(url); + + return new Response( + JSON.stringify({ + search_information: { total_results: 18 }, + organic_results: [ + { + title: "SearchAPI result", + link: "https://searchapi.example.com/page", + snippet: "SearchAPI snippet", + date: "2026-04-06", + source: "SearchAPI Source", + thumbnail: "https://searchapi.example.com/thumb.png", + }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } } + ); + }; + + try { + const result = await handleSearch({ + query: "searchapi query", + provider: "searchapi-search", + maxResults: 2, + searchType: "news", + country: "US", + language: "en", + credentials: { apiKey: "searchapi-key" }, + log: null, + }); + + const url = new URL(capturedUrl); + assert.equal(url.origin + url.pathname, "https://www.searchapi.io/api/v1/search"); + assert.equal(url.searchParams.get("engine"), "google_news"); + assert.equal(url.searchParams.get("api_key"), "searchapi-key"); + assert.equal(url.searchParams.get("gl"), "us"); + assert.equal(url.searchParams.get("hl"), "en"); + assert.equal(result.success, true); + assert.equal(result.data.results[0].metadata.author, "SearchAPI Source"); + assert.equal( + result.data.results[0].metadata.image_url, + "https://searchapi.example.com/thumb.png" + ); + assert.equal(result.data.metrics.total_results_available, 18); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleSearch builds SearXNG requests with custom baseUrl and no apiKey", async () => { + const originalFetch = globalThis.fetch; + let capturedUrl; + + globalThis.fetch = async (url) => { + capturedUrl = String(url); + + return new Response( + JSON.stringify({ + results: [ + { + title: "SearXNG result", + url: "https://searx.example.com/page", + content: "Metasearch snippet", + publishedDate: "2026-04-07", + engines: ["duckduckgo", "brave"], + }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } } + ); + }; + + try { + const result = await handleSearch({ + query: "privacy search", + provider: "searxng-search", + maxResults: 5, + searchType: "news", + language: "pt-BR", + credentials: { + providerSpecificData: { baseUrl: "http://localhost:9999" }, + }, + log: null, + }); + + const url = new URL(capturedUrl); + assert.equal(url.origin + url.pathname, "http://localhost:9999/search"); + assert.equal(url.searchParams.get("format"), "json"); + assert.equal(url.searchParams.get("categories"), "news"); + assert.equal(url.searchParams.get("language"), "pt-BR"); + assert.equal(result.success, true); + assert.equal(result.data.provider, "searxng-search"); + assert.equal(result.data.results[0].metadata.source_type, "duckduckgo, brave"); + } finally { + globalThis.fetch = originalFetch; + } +}); + test("handleSearch rejects queries with invalid control characters", async () => { const result = await handleSearch({ query: "bad\u0000query", diff --git a/tests/unit/search-registry.test.ts b/tests/unit/search-registry.test.ts index 8bcf74245b..dfac5537a9 100644 --- a/tests/unit/search-registry.test.ts +++ b/tests/unit/search-registry.test.ts @@ -6,21 +6,30 @@ import assert from "node:assert/strict"; // Tests for searchRegistry, searchCache, and response normalization // ═══════════════════════════════════════════════════════════════ -const { SEARCH_PROVIDERS, getSearchProvider, getAllSearchProviders, selectProvider } = - await import("../../open-sse/config/searchRegistry.ts"); +const { + SEARCH_PROVIDERS, + getSearchProvider, + getAllSearchProviders, + selectProvider, + supportsSearchType, +} = await import("../../open-sse/config/searchRegistry.ts"); const { computeCacheKey, getOrCoalesce, getCacheStats, SEARCH_CACHE_DEFAULT_TTL_MS } = await import("../../open-sse/services/searchCache.ts"); // ─── Registry Tests ────────────────────────────────────────── -test("SEARCH_PROVIDERS has all 5 providers", () => { +test("SEARCH_PROVIDERS has all 9 providers", () => { assert.ok(SEARCH_PROVIDERS["serper-search"], "serper should exist"); assert.ok(SEARCH_PROVIDERS["brave-search"], "brave should exist"); assert.ok(SEARCH_PROVIDERS["perplexity-search"], "perplexity-search should exist"); assert.ok(SEARCH_PROVIDERS["exa-search"], "exa should exist"); assert.ok(SEARCH_PROVIDERS["tavily-search"], "tavily should exist"); - assert.equal(Object.keys(SEARCH_PROVIDERS).length, 5); + assert.ok(SEARCH_PROVIDERS["google-pse-search"], "google-pse should exist"); + assert.ok(SEARCH_PROVIDERS["linkup-search"], "linkup should exist"); + assert.ok(SEARCH_PROVIDERS["searchapi-search"], "searchapi should exist"); + assert.ok(SEARCH_PROVIDERS["searxng-search"], "searxng should exist"); + assert.equal(Object.keys(SEARCH_PROVIDERS).length, 9); }); test("serper-search config is correct", () => { @@ -74,14 +83,52 @@ test("tavily config is correct", () => { assert.deepEqual(t.searchTypes, ["web", "news"]); }); +test("google-pse-search config is correct", () => { + const g = SEARCH_PROVIDERS["google-pse-search"]; + assert.equal(g.id, "google-pse-search"); + assert.equal(g.method, "GET"); + assert.equal(g.authHeader, "key"); + assert.equal(g.costPerQuery, 0.005); + assert.equal(g.maxMaxResults, 10); +}); + +test("linkup-search config is correct", () => { + const l = SEARCH_PROVIDERS["linkup-search"]; + assert.equal(l.id, "linkup-search"); + assert.equal(l.method, "POST"); + assert.equal(l.authHeader, "bearer"); + assert.deepEqual(l.searchTypes, ["web"]); +}); + +test("searchapi-search config is correct", () => { + const s = SEARCH_PROVIDERS["searchapi-search"]; + assert.equal(s.id, "searchapi-search"); + assert.equal(s.method, "GET"); + assert.equal(s.authHeader, "api_key"); + assert.deepEqual(s.searchTypes, ["web", "news"]); +}); + +test("searxng-search config is correct", () => { + const s = SEARCH_PROVIDERS["searxng-search"]; + assert.equal(s.id, "searxng-search"); + assert.equal(s.method, "GET"); + assert.equal(s.authType, "none"); + assert.equal(s.costPerQuery, 0); + assert.deepEqual(s.searchTypes, ["web", "news"]); +}); + test("getAllSearchProviders returns flat list", () => { const all = getAllSearchProviders(); - assert.equal(all.length, 5); + assert.equal(all.length, 9); assert.ok(all.some((p) => p.id === "serper-search")); assert.ok(all.some((p) => p.id === "brave-search")); assert.ok(all.some((p) => p.id === "perplexity-search")); assert.ok(all.some((p) => p.id === "exa-search")); assert.ok(all.some((p) => p.id === "tavily-search")); + assert.ok(all.some((p) => p.id === "google-pse-search")); + assert.ok(all.some((p) => p.id === "linkup-search")); + assert.ok(all.some((p) => p.id === "searchapi-search")); + assert.ok(all.some((p) => p.id === "searxng-search")); // Each entry should have id, name, searchTypes for (const p of all) { assert.ok(p.id); @@ -91,7 +138,7 @@ test("getAllSearchProviders returns flat list", () => { }); test("selectProvider with explicit provider returns that provider", () => { - const config = selectProvider("brave-search"); + const config = selectProvider("brave-search", "news"); assert.ok(config); assert.equal(config.id, "brave-search"); }); @@ -100,10 +147,23 @@ test("selectProvider with unknown provider returns null", () => { assert.equal(selectProvider("unknown"), null); }); -test("selectProvider without argument returns cheapest (serper)", () => { +test("selectProvider without argument returns cheapest provider", () => { const config = selectProvider(); assert.ok(config); - assert.equal(config.id, "serper-search"); // $0.001 < $0.005 + assert.equal(config.id, "searxng-search"); +}); + +test("selectProvider filters by search type support", () => { + const config = selectProvider(undefined, "news"); + assert.ok(config); + assert.equal(config.id, "searxng-search"); + assert.equal(selectProvider("linkup-search", "news"), null); +}); + +test("supportsSearchType reflects provider capabilities", () => { + assert.equal(supportsSearchType("linkup-search", "web"), true); + assert.equal(supportsSearchType("linkup-search", "news"), false); + assert.equal(supportsSearchType("searxng-search", "news"), true); }); // ─── Cache Key Tests ───────────────────────────────────────── @@ -256,6 +316,70 @@ test("v1SearchSchema accepts tavily provider", async () => { assert.equal(result.data.provider, "tavily-search"); }); +test("v1SearchSchema accepts new search providers", async () => { + const { v1SearchSchema } = await import("../../src/shared/validation/schemas.ts"); + + const providers = [ + "google-pse-search", + "linkup-search", + "searchapi-search", + "searxng-search", + ] as const; + + for (const provider of providers) { + const result = v1SearchSchema.safeParse({ query: "test", provider }); + assert.equal(result.success, true, `${provider} should be accepted`); + } +}); + +test("createProviderSchema allows SearXNG without apiKey", async () => { + const { createProviderSchema } = await import("../../src/shared/validation/schemas.ts"); + + const result = createProviderSchema.safeParse({ + provider: "searxng-search", + name: "Local SearXNG", + providerSpecificData: { baseUrl: "http://localhost:8888/search" }, + }); + + assert.equal(result.success, true); +}); + +test("createProviderSchema requires cx for Google PSE", async () => { + const { createProviderSchema } = await import("../../src/shared/validation/schemas.ts"); + + const missingCx = createProviderSchema.safeParse({ + provider: "google-pse-search", + apiKey: "google-key", + name: "Google PSE", + }); + assert.equal(missingCx.success, false); + + const valid = createProviderSchema.safeParse({ + provider: "google-pse-search", + apiKey: "google-key", + name: "Google PSE", + providerSpecificData: { cx: "engine-id" }, + }); + assert.equal(valid.success, true); +}); + +test("validateProviderApiKeySchema requires cx for Google PSE", async () => { + const { validateProviderApiKeySchema } = await import("../../src/shared/validation/schemas.ts"); + + const missingCx = validateProviderApiKeySchema.safeParse({ + provider: "google-pse-search", + apiKey: "google-key", + }); + assert.equal(missingCx.success, false); + + const valid = validateProviderApiKeySchema.safeParse({ + provider: "google-pse-search", + apiKey: "google-key", + cx: "engine-id", + }); + assert.equal(valid.success, true); +}); + test("v1SearchSchema applies defaults", async () => { const { v1SearchSchema } = await import("../../src/shared/validation/schemas.ts"); diff --git a/tests/unit/search-route.test.ts b/tests/unit/search-route.test.ts new file mode 100644 index 0000000000..0a947fe9cb --- /dev/null +++ b/tests/unit/search-route.test.ts @@ -0,0 +1,178 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-search-route-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const searchRoute = await import("../../src/app/api/v1/search/route.ts"); + +async function resetStorage() { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +async function seedConnection( + provider: string, + overrides: { + apiKey?: string | null; + authType?: string; + providerSpecificData?: Record; + } = {} +) { + return providersDb.createProviderConnection({ + provider, + authType: overrides.authType || "apikey", + name: `${provider}-${Math.random().toString(16).slice(2, 8)}`, + apiKey: overrides.apiKey ?? "test-key", + isActive: true, + testStatus: "active", + providerSpecificData: overrides.providerSpecificData || {}, + }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(() => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("v1 search GET lists all 9 search providers", async () => { + const response = await searchRoute.GET(); + const body = await response.json(); + const ids = body.data.map((item: { id: string }) => item.id); + + assert.equal(response.status, 200); + assert.equal(body.object, "list"); + assert.equal(body.data.length, 9); + assert.deepEqual(ids, [ + "serper-search", + "brave-search", + "perplexity-search", + "exa-search", + "tavily-search", + "google-pse-search", + "linkup-search", + "searchapi-search", + "searxng-search", + ]); +}); + +test("v1 search POST uses stored Linkup credentials and returns normalized results", async () => { + await seedConnection("linkup-search", { apiKey: "linkup-key" }); + + const originalFetch = globalThis.fetch; + let capturedUrl = ""; + let capturedInit: RequestInit | undefined; + + globalThis.fetch = async (url, init = {}) => { + capturedUrl = String(url); + capturedInit = init; + + return new Response( + JSON.stringify({ + results: [ + { + name: "Linkup result", + url: "https://example.com/article", + content: "Linkup snippet", + type: "web", + }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } } + ); + }; + + try { + const response = await searchRoute.POST( + new Request("http://localhost/api/v1/search", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + query: "omniroute linkup", + provider: "linkup-search", + max_results: 1, + search_type: "web", + }), + }) + ); + const body = await response.json(); + + assert.equal(response.status, 200); + assert.equal(capturedUrl, "https://api.linkup.so/v1/search"); + assert.equal( + (capturedInit?.headers as Record).Authorization, + "Bearer linkup-key" + ); + assert.equal(body.provider, "linkup-search"); + assert.equal(body.query, "omniroute linkup"); + assert.equal(body.results.length, 1); + assert.equal(body.results[0].title, "Linkup result"); + assert.equal(body.results[0].snippet, "Linkup snippet"); + assert.equal(body.results[0].citation.provider, "linkup-search"); + assert.equal(body.cached, false); + assert.equal(body.usage.queries_used, 1); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("v1 search POST accepts authless SearXNG with provider_options baseUrl", async () => { + const originalFetch = globalThis.fetch; + let capturedUrl = ""; + + globalThis.fetch = async (url) => { + capturedUrl = String(url); + return new Response( + JSON.stringify({ + results: [ + { + title: "SearXNG result", + url: "https://searx.example/result", + content: "Self-hosted response", + engines: ["duckduckgo"], + }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } } + ); + }; + + try { + const response = await searchRoute.POST( + new Request("http://localhost/api/v1/search", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + query: "self hosted meta search", + provider: "searxng-search", + search_type: "news", + provider_options: { + baseUrl: "http://127.0.0.1:9090/custom-search", + }, + }), + }) + ); + const body = await response.json(); + + assert.equal(response.status, 200); + assert.equal( + capturedUrl, + "http://127.0.0.1:9090/custom-search/search?q=self+hosted+meta+search&format=json&categories=news" + ); + assert.equal(body.provider, "searxng-search"); + assert.equal(body.results[0].title, "SearXNG result"); + assert.equal(body.results[0].citation.provider, "searxng-search"); + } finally { + globalThis.fetch = originalFetch; + } +}); diff --git a/tests/unit/spend-batch-writer.test.ts b/tests/unit/spend-batch-writer.test.ts new file mode 100644 index 0000000000..85b994bc3e --- /dev/null +++ b/tests/unit/spend-batch-writer.test.ts @@ -0,0 +1,150 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-spend-batch-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const costRules = await import("../../src/domain/costRules.ts"); +const domainState = await import("../../src/lib/db/domainState.ts"); +const { SpendBatchWriter, flushSpendBatchWriter, resetSpendBatchWriterForTests } = + await import("../../src/lib/spend/batchWriter.ts"); + +function quietLogger() { + return { + log() {}, + error() {}, + }; +} + +async function resetStorage() { + resetSpendBatchWriterForTests(); + costRules.resetCostData(); + core.resetDbInstance(); + + for (let attempt = 0; attempt < 10; attempt++) { + try { + if (fs.existsSync(TEST_DATA_DIR)) { + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + } + break; + } catch (error) { + if ((error?.code === "EBUSY" || error?.code === "EPERM") && attempt < 9) { + await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1))); + } else { + throw error; + } + } + } + + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +async function waitFor(fn: () => boolean, timeoutMs = 1_000) { + const startedAt = Date.now(); + while (Date.now() - startedAt < timeoutMs) { + if (fn()) return; + await new Promise((resolve) => setTimeout(resolve, 10)); + } + throw new Error("Timed out waiting for condition"); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + resetSpendBatchWriterForTests(); + costRules.resetCostData(); + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("spend batch writer auto-flushes at the configured threshold", async () => { + const persisted: Array<{ apiKeyId: string; cost: number; timestamp: number }> = []; + const writer = new SpendBatchWriter({ + maxBufferSize: 2, + persistEntries: async (entries) => { + persisted.push(...entries); + }, + logger: quietLogger(), + }); + + writer.increment("key-a", 1.25, 100); + writer.increment("key-b", 2.75, 200); + + await waitFor(() => persisted.length === 2); + + assert.deepEqual(persisted, [ + { apiKeyId: "key-a", cost: 1.25, timestamp: 100 }, + { apiKeyId: "key-b", cost: 2.75, timestamp: 200 }, + ]); + + await writer.stop(); +}); + +test("spend batch writer requeues failed flushes and flushes pending entries on stop", async () => { + const persisted: Array<{ apiKeyId: string; cost: number; timestamp: number }> = []; + let shouldFail = true; + + const writer = new SpendBatchWriter({ + persistEntries: async (entries) => { + if (shouldFail) { + shouldFail = false; + throw new Error("database temporarily unavailable"); + } + persisted.push(...entries); + }, + logger: quietLogger(), + }); + + writer.increment("key-a", 1.5, 111); + writer.increment("key-b", 2.5, 222); + + const failed = await writer.flush(); + assert.equal(failed.requeued, true); + assert.equal(writer.getPendingCostTotal("key-a", 0), 1.5); + assert.equal(writer.getPendingCostTotal("key-b", 0), 2.5); + + const stopped = await writer.stop(); + assert.equal(stopped.requeued, false); + assert.equal(stopped.flushedEntries, 2); + assert.equal(writer.getPendingCostTotal("key-a", 0), 0); + assert.equal(writer.getPendingCostTotal("key-b", 0), 0); + assert.deepEqual(persisted, [ + { apiKeyId: "key-a", cost: 1.5, timestamp: 111 }, + { apiKeyId: "key-b", cost: 2.5, timestamp: 222 }, + ]); +}); + +test("recordCost buffers writes while budget checks and summaries still see pending spend", async () => { + costRules.setBudget("key-live", { dailyLimitUsd: 5 }); + costRules.recordCost("key-live", 3.5); + costRules.recordCost("key-live", 1.0); + + assert.deepEqual(domainState.loadCostEntries("key-live", 0), []); + assert.equal(costRules.getDailyTotal("key-live"), 4.5); + assert.equal(costRules.getCostSummary("key-live").dailyTotal, 4.5); + assert.equal(costRules.checkBudget("key-live", 0).periodUsed, 4.5); + assert.equal(costRules.checkBudget("key-live", 1).allowed, false); + + const flushResult = await flushSpendBatchWriter(); + assert.equal(flushResult.flushedEntries, 2); + + const persistedEntries = domainState.loadCostEntries("key-live", 0); + assert.equal(persistedEntries.length, 2); + assert.equal(costRules.getDailyTotal("key-live"), 4.5); +}); + +test("deleteBudget discards pending spend before it reaches the database", async () => { + costRules.setBudget("key-drop", { dailyLimitUsd: 10 }); + costRules.recordCost("key-drop", 2); + costRules.deleteBudget("key-drop"); + + const flushResult = await flushSpendBatchWriter(); + assert.equal(flushResult.flushedEntries, 0); + assert.deepEqual(domainState.loadCostEntries("key-drop", 0), []); +}); diff --git a/tests/unit/sse-auth.test.ts b/tests/unit/sse-auth.test.ts index 3af43b3a2a..98b31f54c8 100644 --- a/tests/unit/sse-auth.test.ts +++ b/tests/unit/sse-auth.test.ts @@ -177,6 +177,26 @@ test("getProviderCredentialsWithQuotaPreflight skips exhausted preflight account assert.equal(selected.connectionId, healthy.id); }); +test("getProviderCredentials skips connections that exclude the requested model and selects the next eligible account", async () => { + const excluded = await seedConnection("openai", { + name: "excluded-first", + priority: 1, + providerSpecificData: { + excludedModels: ["gpt-4o*"], + }, + }); + const allowed = await seedConnection("openai", { + name: "allowed-second", + priority: 2, + apiKey: "sk-allowed", + }); + + const selected = await auth.getProviderCredentials("openai", null, null, "gpt-4o-mini"); + + assert.equal(selected.connectionId, allowed.id); + assert.notEqual(selected.connectionId, excluded.id); +}); + test("getProviderCredentialsWithQuotaPreflight returns allRateLimited when a forced connection is blocked by preflight", async () => { const blocked = await seedConnection("openai", { name: "quota-preflight-forced", diff --git a/tests/unit/tag-routing.test.ts b/tests/unit/tag-routing.test.ts new file mode 100644 index 0000000000..3d8995c3a9 --- /dev/null +++ b/tests/unit/tag-routing.test.ts @@ -0,0 +1,165 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-tag-routing-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const providersDb = await import("../../src/lib/db/providers.ts"); +const { handleComboChat } = await import("../../open-sse/services/combo.ts"); +const { getConnectionRoutingTags, matchesRoutingTags, resolveRequestRoutingTags } = + await import("../../src/domain/tagRouter.ts"); + +function createLog() { + return { + info() {}, + warn() {}, + debug() {}, + error() {}, + }; +} + +function okResponse(content: string) { + return new Response( + JSON.stringify({ + choices: [{ message: { content } }], + }), + { + status: 200, + headers: { "content-type": "application/json" }, + } + ); +} + +async function resetStorage() { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +async function seedConnection(provider: string, name: string, tags: string[]) { + return providersDb.createProviderConnection({ + provider, + authType: "apikey", + name, + apiKey: `sk-${name}`, + providerSpecificData: { tags }, + isActive: true, + }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + await resetStorage(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("tag router normalizes request metadata and matches connection tags", () => { + const resolved = resolveRequestRoutingTags({ + metadata: { + tags: [" Fast ", "cheap", "FAST"], + tag_match_mode: "all", + }, + }); + + assert.deepEqual(resolved.tags, ["fast", "cheap"]); + assert.equal(resolved.matchMode, "all"); + assert.deepEqual(getConnectionRoutingTags({ tags: ["Cheap", "EU-Region"] }), [ + "cheap", + "eu-region", + ]); + assert.equal(matchesRoutingTags(["cheap", "fast"], ["cheap"], "any"), true); + assert.equal(matchesRoutingTags(["cheap"], ["cheap", "fast"], "all"), false); +}); + +test("handleComboChat filters priority targets by metadata.tags using any-match by default", async () => { + const reliable = await seedConnection("openai", "reliable", ["reliable", "us-region"]); + const cheap = await seedConnection("fireworks", "cheap", ["cheap", "fast"]); + const attempts: Array<{ model: string; allowedConnectionIds: string[] | null }> = []; + + const response = await handleComboChat({ + body: { + model: "router", + messages: [{ role: "user", content: "route this cheaply" }], + metadata: { tags: ["cheap"] }, + }, + combo: { + name: "router", + strategy: "priority", + models: ["openai/gpt-4o-mini", "fireworks/gpt-4o-mini"], + }, + handleSingleModel: async (_body, modelStr, target) => { + attempts.push({ + model: modelStr, + allowedConnectionIds: Array.isArray(target?.allowedConnectionIds) + ? target.allowedConnectionIds + : null, + }); + return okResponse(modelStr); + }, + log: createLog(), + }); + + const payload = await response.json(); + assert.equal(response.status, 200); + assert.equal(attempts.length, 1); + assert.equal(attempts[0].model, "fireworks/gpt-4o-mini"); + assert.deepEqual(attempts[0].allowedConnectionIds, [cheap.id]); + assert.equal(payload.choices[0].message.content, "fireworks/gpt-4o-mini"); + assert.notEqual(reliable.id, cheap.id); +}); + +test("handleComboChat honors tag_match_mode=all and falls back to the full target set when nothing matches", async () => { + await seedConnection("openai", "openai-all", ["cheap", "fast"]); + await seedConnection("fireworks", "fireworks-partial", ["cheap"]); + + const attemptsAll: string[] = []; + const allModeResponse = await handleComboChat({ + body: { + model: "router-all", + messages: [{ role: "user", content: "need cheap and fast" }], + metadata: { tags: ["cheap", "fast"], tag_match_mode: "all" }, + }, + combo: { + name: "router-all", + strategy: "priority", + models: ["fireworks/gpt-4o-mini", "openai/gpt-4o-mini"], + }, + handleSingleModel: async (_body, modelStr) => { + attemptsAll.push(modelStr); + return okResponse(modelStr); + }, + log: createLog(), + }); + + assert.equal(allModeResponse.status, 200); + assert.deepEqual(attemptsAll, ["openai/gpt-4o-mini"]); + + const attemptsFallback: string[] = []; + const fallbackResponse = await handleComboChat({ + body: { + model: "router-fallback", + messages: [{ role: "user", content: "need eu-only" }], + metadata: { tags: ["eu-only"] }, + }, + combo: { + name: "router-fallback", + strategy: "priority", + models: ["openai/gpt-4o-mini", "fireworks/gpt-4o-mini"], + }, + handleSingleModel: async (_body, modelStr) => { + attemptsFallback.push(modelStr); + return okResponse(modelStr); + }, + log: createLog(), + }); + + assert.equal(fallbackResponse.status, 200); + assert.deepEqual(attemptsFallback, ["openai/gpt-4o-mini"]); +});