From 88a10dfcc43214ab77550afafc49e35e1aca1b19 Mon Sep 17 00:00:00 2001 From: Nick Roth Date: Thu, 12 Feb 2026 20:11:33 -0600 Subject: [PATCH 1/2] feat(auth): add model-level rate limit locking for multi-bucket providers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Providers like Antigravity maintain separate quota buckets per model family (e.g. Claude vs Gemini). A 429 on claude-opus previously locked the entire account, preventing gemini-pro requests even though its quota was full. This adds in-memory per-model locking so that only the specific model is skipped during account selection while other models remain accessible. Changes: - Add model-aware lock tracking in auth.js (Map) - Pass model context from chat handler to auth service - Multi-bucket behavior gated to known providers (MULTI_BUCKET_PROVIDERS set) - No database schema changes — locks are in-memory and clear on restart Closes #110 --- src/sse/handlers/chat.js | 4 +- src/sse/services/auth.js | 94 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 90 insertions(+), 8 deletions(-) diff --git a/src/sse/handlers/chat.js b/src/sse/handlers/chat.js index 18b537da054..cddbaac6a9a 100644 --- a/src/sse/handlers/chat.js +++ b/src/sse/handlers/chat.js @@ -120,7 +120,7 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re let lastStatus = null; while (true) { - const credentials = await getProviderCredentials(provider, excludeConnectionId); + const credentials = await getProviderCredentials(provider, excludeConnectionId, model); // All accounts unavailable if (!credentials || credentials.allRateLimited) { @@ -170,7 +170,7 @@ async function handleSingleModelChat(body, modelStr, clientRawRequest = null, re if (result.success) return result.response; // Mark account unavailable (auto-calculates cooldown with exponential backoff) - const { shouldFallback } = await markAccountUnavailable(credentials.connectionId, result.status, result.error, provider); + const { shouldFallback } = await markAccountUnavailable(credentials.connectionId, result.status, result.error, provider, model); if (shouldFallback) { log.warn("AUTH", `Account ${accountId}... unavailable (${result.status}), trying fallback`); diff --git a/src/sse/services/auth.js b/src/sse/services/auth.js index 31fe92f788a..83de7981c93 100644 --- a/src/sse/services/auth.js +++ b/src/sse/services/auth.js @@ -5,13 +5,70 @@ import * as log from "../utils/logger.js"; // Mutex to prevent race conditions during account selection let selectionMutex = Promise.resolve(); +// --------------------------------------------------------------------------- +// Model-level rate-limit locking (in-memory) +// +// Providers like Antigravity maintain *separate* quota buckets per model +// family (e.g. Claude vs Gemini). When a 429 arrives for one model, the +// standard account-level DB lock would block ALL models on that account, +// wasting quota that is still available for other model families. +// +// This module tracks model-specific locks in memory so that only the +// affected model is skipped during account selection while the rest of the +// account's quota remains accessible. +// +// Locks are intentionally in-memory: they clear on restart and require no +// database schema migration. +// --------------------------------------------------------------------------- + +/** Providers known to have independent per-model quota buckets */ +const MULTI_BUCKET_PROVIDERS = new Set(["antigravity"]); + +/** Map<"connectionId:model", expiryTimestamp> */ +const modelLocks = new Map(); + +/** Default lock duration for model-level rate limits (5 minutes) */ +const DEFAULT_MODEL_LOCK_MS = 5 * 60 * 1000; + +/** + * Check whether a specific model is temporarily locked on a connection. + * Expired locks are cleaned up lazily. + */ +function isModelLocked(connectionId, model) { + if (!connectionId || !model) return false; + const key = `${connectionId}:${model}`; + const expiry = modelLocks.get(key); + if (!expiry) return false; + if (expiry > Date.now()) return true; + modelLocks.delete(key); // clean up expired + return false; +} + +/** + * Lock a model on a specific connection for `durationMs` milliseconds. + */ +function lockModel(connectionId, model, durationMs) { + if (!connectionId || !model) return; + const key = `${connectionId}:${model}`; + modelLocks.set(key, Date.now() + durationMs); + log.warn("AUTH", `Model lock: ${model} on ${connectionId.slice(0, 8)} for ${Math.round(durationMs / 1000)}s`); +} + +/** + * Check whether a provider uses per-model quota buckets. + */ +function isMultiBucketProvider(provider) { + return MULTI_BUCKET_PROVIDERS.has(provider); +} + /** * Get provider credentials from localDb * Filters out unavailable accounts and returns the selected account based on strategy * @param {string} provider - Provider name * @param {string|null} excludeConnectionId - Connection ID to exclude (for retry with next account) + * @param {string|null} model - Model name for per-model rate limit filtering */ -export async function getProviderCredentials(provider, excludeConnectionId = null) { +export async function getProviderCredentials(provider, excludeConnectionId = null, model = null) { // Acquire mutex to prevent race conditions const currentMutex = selectionMutex; let resolveMutex; @@ -21,7 +78,7 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul await currentMutex; const connections = await getProviderConnections({ provider, isActive: true }); - log.debug("AUTH", `${provider} | total connections: ${connections.length}, excludeId: ${excludeConnectionId || "none"}`); + log.debug("AUTH", `${provider} | total connections: ${connections.length}, excludeId: ${excludeConnectionId || "none"}, model: ${model || "any"}`); if (connections.length === 0) { // Check all connections (including inactive) to see if rate limited @@ -43,9 +100,11 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul } // Filter out unavailable accounts and excluded connection + const multiBucket = isMultiBucketProvider(provider); const availableConnections = connections.filter(c => { if (excludeConnectionId && c.id === excludeConnectionId) return false; if (isAccountUnavailable(c.rateLimitedUntil)) return false; + if (multiBucket && model && isModelLocked(c.id, model)) return false; return true; }); @@ -53,8 +112,9 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul connections.forEach(c => { const excluded = excludeConnectionId && c.id === excludeConnectionId; const rateLimited = isAccountUnavailable(c.rateLimitedUntil); - if (excluded || rateLimited) { - log.debug("AUTH", ` → ${c.id?.slice(0, 8)} | ${excluded ? "excluded" : ""} ${rateLimited ? `rateLimited until ${c.rateLimitedUntil}` : ""}`); + const modelLocked = multiBucket && model && isModelLocked(c.id, model); + if (excluded || rateLimited || modelLocked) { + log.debug("AUTH", ` → ${c.id?.slice(0, 8)} | ${excluded ? "excluded" : ""} ${rateLimited ? `rateLimited until ${c.rateLimitedUntil}` : ""} ${modelLocked ? `modelLocked(${model})` : ""}`); } }); @@ -73,6 +133,15 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul lastErrorCode: earliestConn?.errorCode || null }; } + if (multiBucket && model) { + log.warn("AUTH", `${provider} | all accounts model-locked for ${model}`); + return { + allRateLimited: true, + retryAfter: new Date(Date.now() + 60000).toISOString(), + retryAfterHuman: "reset after 1m", + lastError: `All accounts rate limited for model ${model}` + }; + } log.warn("AUTH", `${provider} | all ${connections.length} accounts unavailable`); return null; } @@ -144,10 +213,17 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul } /** - * Mark account as unavailable — reads backoffLevel from DB, calculates cooldown with exponential backoff, saves new level + * Mark account as unavailable — reads backoffLevel from DB, calculates cooldown with exponential backoff, saves new level. + * For multi-bucket providers (e.g. Antigravity), 429 errors lock only the specific model in memory + * rather than the entire account in the database. + * @param {string} connectionId + * @param {number} status - HTTP status code from upstream + * @param {string} errorText + * @param {string|null} provider + * @param {string|null} model - The specific model that triggered the error * @returns {{ shouldFallback: boolean, cooldownMs: number }} */ -export async function markAccountUnavailable(connectionId, status, errorText, provider = null) { +export async function markAccountUnavailable(connectionId, status, errorText, provider = null, model = null) { // Read current connection to get backoffLevel const connections = await getProviderConnections({ provider }); const conn = connections.find(c => c.id === connectionId); @@ -156,6 +232,12 @@ export async function markAccountUnavailable(connectionId, status, errorText, pr const { shouldFallback, cooldownMs, newBackoffLevel } = checkFallbackError(status, errorText, backoffLevel); if (!shouldFallback) return { shouldFallback: false, cooldownMs: 0 }; + if (isMultiBucketProvider(provider) && status === 429 && model) { + const lockDuration = cooldownMs > 0 ? cooldownMs : DEFAULT_MODEL_LOCK_MS; + lockModel(connectionId, model, lockDuration); + return { shouldFallback: true, cooldownMs: 0 }; + } + const rateLimitedUntil = getUnavailableUntil(cooldownMs); const reason = typeof errorText === "string" ? errorText.slice(0, 100) : "Provider error"; From c247be572dfaadfd49e4838c5ba73fe3edca06cf Mon Sep 17 00:00:00 2001 From: Nick Roth Date: Thu, 12 Mar 2026 12:58:21 -0500 Subject: [PATCH 2/2] feat: strip markdown code blocks from JSON responses When models (like kimi) return JSON wrapped in markdown code blocks (\`\`\`json...\`\`\`), strip the markers before sending to client. Changes: - Add markdown stripping to streaming text_delta responses - Add markdown stripping to non-streaming Claude format responses - Translate response_format to system prompts for JSON schema/object modes - Add jsonExtractor utility for consistent markdown removal This fixes AI_APICallError: Invalid JSON response when using AI SDK generateObject with kimi models through 9router. --- .../handlers/chatCore/nonStreamingHandler.js | 15 ++++---- .../translator/request/openai-to-claude.js | 9 +++-- .../translator/response/claude-to-openai.js | 5 ++- open-sse/utils/jsonExtractor.js | 34 +++++++++++++++++++ open-sse/utils/stream.js | 23 ++++++++----- 5 files changed, 63 insertions(+), 23 deletions(-) create mode 100644 open-sse/utils/jsonExtractor.js diff --git a/open-sse/handlers/chatCore/nonStreamingHandler.js b/open-sse/handlers/chatCore/nonStreamingHandler.js index 62efc03dfb0..2aa0a6367aa 100644 --- a/open-sse/handlers/chatCore/nonStreamingHandler.js +++ b/open-sse/handlers/chatCore/nonStreamingHandler.js @@ -1,9 +1,8 @@ import { FORMATS } from "../../translator/formats.js"; import { needsTranslation } from "../../translator/index.js"; -import { ollamaBodyToOpenAI } from "../../translator/response/ollama-to-openai.js"; import { addBufferToUsage, filterUsageForFormat } from "../../utils/usageTracking.js"; import { createErrorResult } from "../../utils/error.js"; -import { HTTP_STATUS } from "../../config/runtimeConfig.js"; +import { HTTP_STATUS } from "../../config/constants.js"; import { parseSSEToOpenAIResponse } from "./sseToJsonHandler.js"; import { buildRequestDetail, extractRequestConfig, extractUsageFromResponse, saveUsageStats } from "./requestDetail.js"; import { appendRequestLog, saveRequestDetail } from "@/lib/usageDb.js"; @@ -77,8 +76,11 @@ export function translateNonStreamingResponse(responseBody, targetFormat, source const toolCalls = []; for (const block of responseBody.content) { - if (block.type === "text") textContent += block.text; - else if (block.type === "thinking") thinkingContent += block.thinking || ""; + if (block.type === "text") { + let text = block.text; + text = text.replace(/^\s*```\s*json\s*\n?/i, "").replace(/\n?\s*```\s*$/i, ""); + textContent += text; + } else if (block.type === "tool_use") { toolCalls.push({ id: block.id, type: "function", function: { name: block.name, arguments: JSON.stringify(block.input || {}) } }); } @@ -112,11 +114,6 @@ export function translateNonStreamingResponse(responseBody, targetFormat, source return result; } - // Ollama - if (targetFormat === FORMATS.OLLAMA) { - return ollamaBodyToOpenAI(responseBody); - } - return responseBody; } diff --git a/open-sse/translator/request/openai-to-claude.js b/open-sse/translator/request/openai-to-claude.js index bb2c3153365..3a252bcfaec 100644 --- a/open-sse/translator/request/openai-to-claude.js +++ b/open-sse/translator/request/openai-to-claude.js @@ -1,6 +1,6 @@ import { register } from "../index.js"; import { FORMATS } from "../formats.js"; -import { CLAUDE_SYSTEM_PROMPT } from "../../config/appConstants.js"; +import { CLAUDE_SYSTEM_PROMPT } from "../../config/constants.js"; import { adjustMaxTokens } from "../helpers/maxTokensHelper.js"; // Empty prefix matches real Claude Code behavior (no tool name prefix). @@ -106,10 +106,9 @@ export function openaiToClaudeRequest(model, body, stream) { if (responseFormat.type === "json_schema" && responseFormat.json_schema?.schema) { const schemaJson = JSON.stringify(responseFormat.json_schema.schema, null, 2); systemParts.push(`You must respond with valid JSON that strictly follows this JSON schema: -\`\`\`json -${schemaJson} -\`\`\` -Respond ONLY with the JSON object, no other text.`); +Schema: ${schemaJson} + +IMPORTANT: Return ONLY the raw JSON object without markdown formatting, code blocks, or backticks.`); } else if (responseFormat.type === "json_object") { systemParts.push("You must respond with valid JSON. Respond ONLY with a JSON object, no other text."); } diff --git a/open-sse/translator/response/claude-to-openai.js b/open-sse/translator/response/claude-to-openai.js index 2dfe0dbbf73..794e7b221ea 100644 --- a/open-sse/translator/response/claude-to-openai.js +++ b/open-sse/translator/response/claude-to-openai.js @@ -69,7 +69,10 @@ export function claudeToOpenAIResponse(chunk, state) { if (chunk.index === state.serverToolBlockIndex) break; const delta = chunk.delta; if (delta?.type === "text_delta" && delta.text) { - results.push(createChunk(state, { content: delta.text })); + let text = delta.text; + // Strip markdown code block markers + text = text.replace(/^\s*```\s*json\s*\n?/i, "").replace(/\n?\s*```\s*$/i, ""); + results.push(createChunk(state, { content: text })); } else if (delta?.type === "thinking_delta" && delta.thinking) { results.push(createChunk(state, { reasoning_content: delta.thinking })); } else if (delta?.type === "input_json_delta" && delta.partial_json) { diff --git a/open-sse/utils/jsonExtractor.js b/open-sse/utils/jsonExtractor.js new file mode 100644 index 00000000000..3cf7ffb6096 --- /dev/null +++ b/open-sse/utils/jsonExtractor.js @@ -0,0 +1,34 @@ +/** + * Extract JSON from content that may be wrapped in markdown code blocks + * Only processes content when responseFormat is specified + * @param {string} content - Raw response content + * @param {object} responseFormat - The response_format from the request body + * @returns {string} - Extracted JSON or original content + */ +export function extractJSON(content, responseFormat) { + // Only process if response_format is specified + if (!responseFormat || !content || typeof content !== 'string') { + return content; + } + + // Check if it's a structured output request + const isStructuredOutput = responseFormat.type === 'json_schema' || + responseFormat.type === 'json_object'; + + if (!isStructuredOutput) { + return content; + } + + // Try to find JSON between markdown code blocks + const codeBlockMatch = content.match(/```(?:json)?\s*\n?([\s\S]*?)\n?```/i); + if (codeBlockMatch && codeBlockMatch[1]) { + const extracted = codeBlockMatch[1].trim(); + // Validate it looks like JSON + if (extracted.startsWith('{') || extracted.startsWith('[')) { + return extracted; + } + } + + // Return original if no markdown found or doesn't look like JSON + return content; +} diff --git a/open-sse/utils/stream.js b/open-sse/utils/stream.js index 4f5e4a01cb4..03693bdecf9 100644 --- a/open-sse/utils/stream.js +++ b/open-sse/utils/stream.js @@ -1,3 +1,4 @@ +import { extractJSON } from './jsonExtractor.js'; import { translateResponse, initState } from "../translator/index.js"; import { FORMATS } from "../translator/formats.js"; import { trackPendingRequest, appendRequestLog } from "@/lib/usageDb.js"; @@ -159,7 +160,7 @@ export function createSSEStream(options = {}) { // Translate mode if (!trimmed) continue; - const parsed = parseSSELine(trimmed, targetFormat); + const parsed = parseSSELine(trimmed); if (!parsed) continue; if (parsed && parsed.done) { @@ -278,13 +279,16 @@ export function createSSEStream(options = {}) { // Some clients (e.g. OpenClaw) expect the OpenAI-style sentinel: // data: [DONE]\n\n // Without it they can hang until timeout and trigger failover. - const doneOutput = "data: [DONE]\n\n"; - reqLogger?.appendConvertedChunk?.(doneOutput); - controller.enqueue(sharedEncoder.encode(doneOutput)); + // Only emit data: [DONE] for streaming requests + if (body?.stream === true) { + const doneOutput = "data: [DONE]\n\n"; + reqLogger?.appendConvertedChunk?.(doneOutput); + controller.enqueue(sharedEncoder.encode(doneOutput)); + } if (onStreamComplete) { onStreamComplete({ - content: accumulatedContent, + content: extractJSON(accumulatedContent, body?.response_format), thinking: accumulatedThinking }, usage, ttftAt); } @@ -330,9 +334,12 @@ export function createSSEStream(options = {}) { } } - const doneOutput = "data: [DONE]\n\n"; - reqLogger?.appendConvertedChunk?.(doneOutput); + // Only emit data: [DONE] for streaming requests + if (body?.stream === true) { + const doneOutput = "data: [DONE]\n\n"; + reqLogger?.appendConvertedChunk?.(doneOutput); controller.enqueue(sharedEncoder.encode(doneOutput)); + } if (!hasValidUsage(state?.usage) && totalContentLength > 0) { state.usage = estimateUsage(body, totalContentLength, sourceFormat); @@ -346,7 +353,7 @@ export function createSSEStream(options = {}) { if (onStreamComplete) { onStreamComplete({ - content: accumulatedContent, + content: extractJSON(accumulatedContent, body?.response_format), thinking: accumulatedThinking }, state?.usage, ttftAt); }