diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 49aefdbc8c..f25a058713 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -88,6 +88,11 @@ import { updateFromResponseBody, initializeRateLimits, } from "../services/rateLimitManager.ts"; +import { + acquire as acquireAccountSemaphore, + buildAccountSemaphoreKey, + markBlocked as markAccountSemaphoreBlocked, +} from "../services/accountSemaphore.ts"; import { generateSignature, getCachedResponse, @@ -419,6 +424,72 @@ function getHeaderValueCaseInsensitive( return null; } +function toFiniteNumberOrNull(value: unknown): number | null { + 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 : null; + } + return null; +} + +function isSemaphoreTimeoutError(error: unknown): error is Error & { code: string } { + return ( + !!error && + typeof error === "object" && + (error as { code?: unknown }).code === "SEMAPHORE_TIMEOUT" + ); +} + +function resolveAccountSemaphoreAccountKey( + connectionId: string | null | undefined, + credentials: Record | null | undefined +): string | null { + if (typeof connectionId === "string" && connectionId.trim().length > 0) { + return connectionId; + } + + const candidateKeys = [ + credentials?.connectionId, + credentials?.id, + credentials?.email, + credentials?.name, + credentials?.displayName, + ]; + + for (const candidate of candidateKeys) { + if (typeof candidate === "string" && candidate.trim().length > 0) { + return candidate.trim(); + } + } + + return null; +} + +function resolveAccountSemaphoreMaxConcurrency( + credentials: Record | null | undefined +): number | null { + return toFiniteNumberOrNull(credentials?.maxConcurrent); +} + +function resolveAccountSemaphoreKey({ + provider, + model, + connectionId, + credentials, +}: { + provider: string | null | undefined; + model: string; + connectionId: string | null | undefined; + credentials: Record | null | undefined; +}): string | null { + const accountKey = resolveAccountSemaphoreAccountKey(connectionId, credentials); + if (!accountKey || !provider) return null; + return buildAccountSemaphoreKey({ provider, model, accountKey }); +} + function buildClaudePromptCacheLogMeta( targetFormat: string, finalBody: Record | null | undefined, @@ -1720,6 +1791,15 @@ export async function handleChatCore({ const executeProviderRequest = async (modelToCall = effectiveModel, allowDedup = false) => { const execute = async () => { + const executionCredentials = getExecutionCredentials(); + const accountSemaphoreMaxConcurrency = + resolveAccountSemaphoreMaxConcurrency(executionCredentials); + const accountSemaphoreKey = resolveAccountSemaphoreKey({ + provider, + model: modelToCall, + connectionId, + credentials: executionCredentials, + }); let bodyToSend = translatedBody.model === modelToCall ? translatedBody @@ -1787,60 +1867,71 @@ export async function handleChatCore({ } } - const rawResult = await withRateLimit(provider, connectionId, modelToCall, async () => { - let attempts = 0; - const maxAttempts = provider === "qwen" ? 3 : 1; + const releaseAccountSemaphore = + accountSemaphoreKey && accountSemaphoreMaxConcurrency != null + ? await acquireAccountSemaphore(accountSemaphoreKey, { + maxConcurrency: accountSemaphoreMaxConcurrency, + }) + : () => {}; - while (attempts < maxAttempts) { - const res = await executor.execute({ - model: modelToCall, - body: bodyToSend, - stream: upstreamStream, - credentials: getExecutionCredentials(), - signal: streamController.signal, - log, - extendedContext, - upstreamExtraHeaders: buildUpstreamHeadersForExecute(modelToCall), - clientHeaders: buildExecutorClientHeaders(clientRawRequest?.headers, userAgent), - onCredentialsRefreshed, - }); + try { + const rawResult = await withRateLimit(provider, connectionId, modelToCall, async () => { + let attempts = 0; + const maxAttempts = provider === "qwen" ? 3 : 1; - // Qwen 429 strict quota backoff (wait 1.5s, 3s and retry) - if (provider === "qwen" && res.response.status === 429 && attempts < maxAttempts - 1) { - const bodyPeek = await res.response - .clone() - .text() - .catch(() => ""); - if (bodyPeek.toLowerCase().includes("exceeded your current quota")) { - const delay = 1500 * (attempts + 1); - log?.warn?.("QWEN_RETRY", `Quota 429 hit. Retrying in ${delay}ms...`); - await new Promise((r) => setTimeout(r, delay)); - attempts++; - continue; + while (attempts < maxAttempts) { + const res = await executor.execute({ + model: modelToCall, + body: bodyToSend, + stream: upstreamStream, + credentials: executionCredentials, + signal: streamController.signal, + log, + extendedContext, + upstreamExtraHeaders: buildUpstreamHeadersForExecute(modelToCall), + clientHeaders: buildExecutorClientHeaders(clientRawRequest?.headers, userAgent), + onCredentialsRefreshed, + }); + + // Qwen 429 strict quota backoff (wait 1.5s, 3s and retry) + if (provider === "qwen" && res.response.status === 429 && attempts < maxAttempts - 1) { + const bodyPeek = await res.response + .clone() + .text() + .catch(() => ""); + if (bodyPeek.toLowerCase().includes("exceeded your current quota")) { + const delay = 1500 * (attempts + 1); + log?.warn?.("QWEN_RETRY", `Quota 429 hit. Retrying in ${delay}ms...`); + await new Promise((r) => setTimeout(r, delay)); + attempts++; + continue; + } } + return res; } - return res; - } - }); + }); - if (stream) return rawResult; + if (stream) return rawResult; - // Non-stream responses need cloning for shared dedup consumers. - const status = rawResult.response.status; - const statusText = rawResult.response.statusText; - const headers = Array.from(rawResult.response.headers.entries()) as [string, string][]; - const payload = await rawResult.response.text(); + // Non-stream responses need cloning for shared dedup consumers. + const status = rawResult.response.status; + const statusText = rawResult.response.statusText; + const headers = Array.from(rawResult.response.headers.entries()) as [string, string][]; + const payload = await rawResult.response.text(); - return { - ...rawResult, - response: new Response(payload, { status, statusText, headers }), - _dedupSnapshot: { - status, - statusText, - headers, - payload, - }, - }; + return { + ...rawResult, + response: new Response(payload, { status, statusText, headers }), + _dedupSnapshot: { + status, + statusText, + headers, + payload, + }, + }; + } finally { + releaseAccountSemaphore(); + } }; if (allowDedup && dedupEnabled && dedupHash) { @@ -1904,6 +1995,28 @@ export async function handleChatCore({ ); } catch (error) { trackPendingRequest(model, provider, connectionId, false); + if (isSemaphoreTimeoutError(error)) { + appendRequestLog({ + model, + provider, + connectionId, + status: `FAILED ${error.code}`, + }).catch(() => {}); + if (isCombo) { + throw error; + } + const failureMessage = error.message || "Semaphore timeout"; + persistAttemptLogs({ + status: HTTP_STATUS.RATE_LIMITED, + error: failureMessage, + providerRequest: finalBody || translatedBody, + clientResponse: buildErrorBody(HTTP_STATUS.RATE_LIMITED, failureMessage), + claudeCacheMeta: claudePromptCacheLogMeta, + cacheSource: "upstream", + }); + persistFailureUsage(HTTP_STATUS.RATE_LIMITED, error.code); + return createErrorResult(HTTP_STATUS.RATE_LIMITED, failureMessage); + } const failureStatus = error.name === "AbortError" ? 499 @@ -2080,6 +2193,15 @@ export async function handleChatCore({ // each model has independent quota. A 429 on one model must NOT lock out // the entire connection — other models may still have quota available. const rateLimitCooldownMs = retryAfterMs || COOLDOWN_MS.rateLimit; + const accountSemaphoreKey = resolveAccountSemaphoreKey({ + provider, + model: currentModel, + connectionId, + credentials, + }); + if (accountSemaphoreKey) { + markAccountSemaphoreBlocked(accountSemaphoreKey, rateLimitCooldownMs); + } if ( lockModelIfPerModelQuota( provider, @@ -2109,17 +2231,27 @@ export async function handleChatCore({ } } else if (errorType === PROVIDER_ERROR_TYPES.QUOTA_EXHAUSTED) { // Providers with per-model quotas — lock the model only, not the connection + const quotaCooldownMs = retryAfterMs || COOLDOWN_MS.rateLimit; + const accountSemaphoreKey = resolveAccountSemaphoreKey({ + provider, + model: currentModel, + connectionId, + credentials, + }); + if (accountSemaphoreKey) { + markAccountSemaphoreBlocked(accountSemaphoreKey, quotaCooldownMs); + } if ( lockModelIfPerModelQuota( provider, connectionId, model, "quota_exhausted", - retryAfterMs || COOLDOWN_MS.rateLimit + quotaCooldownMs ) ) { console.warn( - `[provider] Node ${connectionId} model-only quota exhausted (${statusCode}) for ${model} - ${Math.ceil((retryAfterMs || COOLDOWN_MS.rateLimit) / 1000)}s (connection stays active)` + `[provider] Node ${connectionId} model-only quota exhausted (${statusCode}) for ${model} - ${Math.ceil(quotaCooldownMs / 1000)}s (connection stays active)` ); } else { await updateProviderConnection(connectionId, { diff --git a/open-sse/services/__tests__/accountSemaphore.test.ts b/open-sse/services/__tests__/accountSemaphore.test.ts new file mode 100644 index 0000000000..04049d0686 --- /dev/null +++ b/open-sse/services/__tests__/accountSemaphore.test.ts @@ -0,0 +1,164 @@ +import assert from "node:assert/strict"; +import { afterEach, describe, it } from "node:test"; + +import { + acquire, + buildAccountSemaphoreKey, + getStats, + markBlocked, + reset, + resetAll, +} from "../accountSemaphore"; + +afterEach(() => { + resetAll(); +}); + +describe("accountSemaphore", async () => { + it("queues requests beyond the account cap and drains on release", async () => { + const key = buildAccountSemaphoreKey({ + provider: "alibaba", + model: "qwen-max", + accountKey: "acct-1", + }); + + const releaseA = await acquire(key, { maxConcurrency: 2, timeoutMs: 200 }); + const releaseB = await acquire(key, { maxConcurrency: 2, timeoutMs: 200 }); + const queued = acquire(key, { maxConcurrency: 2, timeoutMs: 200 }); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + assert.deepEqual(getStats()[key], { + running: 2, + queued: 1, + maxConcurrency: 2, + blockedUntil: null, + }); + + releaseA(); + const releaseC = await queued; + + assert.deepEqual(getStats()[key], { + running: 2, + queued: 0, + maxConcurrency: 2, + blockedUntil: null, + }); + + releaseA(); + releaseB(); + releaseC(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + assert.equal(getStats()[key], undefined); + }); + + it("returns a no-op release when concurrency is bypassed", async () => { + const key = buildAccountSemaphoreKey({ + provider: "alibaba", + model: "qwen-max", + accountKey: "acct-bypass", + }); + + const release = await acquire(key, { maxConcurrency: 0, timeoutMs: 50 }); + + assert.deepEqual(getStats(), {}); + release(); + release(); + assert.deepEqual(getStats(), {}); + }); + + it("uses SEMAPHORE_TIMEOUT for timed out queued requests", async () => { + const key = buildAccountSemaphoreKey({ + provider: "alibaba", + model: "qwen-max", + accountKey: "acct-timeout", + }); + + const release = await acquire(key, { maxConcurrency: 1, timeoutMs: 100 }); + + await assert.rejects( + acquire(key, { maxConcurrency: 1, timeoutMs: 20 }), + (error: Error & { code?: string }) => error.code === "SEMAPHORE_TIMEOUT" + ); + + release(); + }); + + it("keeps release idempotent for finally blocks", async () => { + const key = buildAccountSemaphoreKey({ + provider: "alibaba", + model: "qwen-max", + accountKey: "acct-finally", + }); + + const release = await acquire(key, { maxConcurrency: 1, timeoutMs: 100 }); + + try { + throw new Error("boom"); + } catch { + release(); + release(); + } + + const secondRelease = await acquire(key, { maxConcurrency: 1, timeoutMs: 100 }); + secondRelease(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + assert.equal(getStats()[key], undefined); + }); + + it("supports temporary blocking and explicit reset hooks", async () => { + const key = buildAccountSemaphoreKey({ + provider: "alibaba", + model: "qwen-max", + accountKey: "acct-blocked", + }); + + markBlocked(key, 30); + const queued = acquire(key, { maxConcurrency: 1, timeoutMs: 100 }); + + await new Promise((resolve) => setTimeout(resolve, 10)); + assert.equal(getStats()[key]?.queued, 1); + + reset(key); + + await assert.rejects(queued, /Semaphore reset/); + assert.equal(getStats()[key], undefined); + }); + + it("preserves existing maxConcurrency when markBlocked is applied", async () => { + const key = buildAccountSemaphoreKey({ + provider: "alibaba", + model: "qwen-max", + accountKey: "acct-preserve-max", + }); + + const releaseA = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); + const releaseB = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); + + markBlocked(key, 30); + + assert.equal(getStats()[key]?.maxConcurrency, 3); + + releaseA(); + releaseB(); + + await new Promise((resolve) => setTimeout(resolve, 40)); + + const releaseC = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); + const releaseD = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); + const releaseE = await acquire(key, { maxConcurrency: 3, timeoutMs: 100 }); + + assert.deepEqual(getStats()[key], { + running: 3, + queued: 0, + maxConcurrency: 3, + blockedUntil: null, + }); + + releaseC(); + releaseD(); + releaseE(); + }); +}); diff --git a/open-sse/services/accountSemaphore.ts b/open-sse/services/accountSemaphore.ts new file mode 100644 index 0000000000..1235e5c990 --- /dev/null +++ b/open-sse/services/accountSemaphore.ts @@ -0,0 +1,290 @@ +/** + * Account Semaphore + * + * In-memory provider/account concurrency limiter keyed by provider, model, and account. + * Requests beyond the configured concurrency cap wait in a FIFO queue until a slot opens, + * the gate is unblocked, or the queue timeout expires. + */ + +export interface AccountSemaphoreKeyParts { + provider: string; + model: string; + accountKey: string; +} + +interface QueuedAcquire { + resolve: (release: () => void) => void; + reject: (error: Error) => void; + timer: ReturnType; +} + +interface AccountGate { + running: number; + maxConcurrency: number; + queue: QueuedAcquire[]; + blockedUntil: number | null; + cleanupTimer: ReturnType | null; +} + +export interface AcquireAccountSemaphoreOptions { + maxConcurrency?: number | null; + timeoutMs?: number; +} + +export interface AccountSemaphoreStatsEntry { + running: number; + queued: number; + maxConcurrency: number; + blockedUntil: string | null; +} + +const DEFAULT_TIMEOUT_MS = 30_000; + +const gates = new Map(); + +/** + * Build the canonical account semaphore key. + */ +export function buildAccountSemaphoreKey({ + provider, + model, + accountKey, +}: AccountSemaphoreKeyParts): string { + return `${String(provider)}:${String(model)}:${String(accountKey)}`; +} + +function isBypassed(maxConcurrency?: number | null): boolean { + return maxConcurrency == null || maxConcurrency <= 0; +} + +function createNoopReleaseFn(): () => void { + let released = false; + + return () => { + if (released) return; + released = true; + }; +} + +function ensureGate(semaphoreKey: string, maxConcurrency: number): AccountGate { + const existing = gates.get(semaphoreKey); + if (existing) { + existing.maxConcurrency = maxConcurrency; + return existing; + } + + const created: AccountGate = { + running: 0, + maxConcurrency, + queue: [], + blockedUntil: null, + cleanupTimer: null, + }; + gates.set(semaphoreKey, created); + return created; +} + +function isBlocked(gate: AccountGate): boolean { + if (!gate.blockedUntil) return false; + if (Date.now() >= gate.blockedUntil) { + gate.blockedUntil = null; + return false; + } + return true; +} + +function clearCleanupTimer(gate: AccountGate): void { + if (!gate.cleanupTimer) return; + clearTimeout(gate.cleanupTimer); + gate.cleanupTimer = null; +} + +function cleanupGateIfIdle(semaphoreKey: string): void { + const gate = gates.get(semaphoreKey); + if (!gate) return; + if (gate.running > 0 || gate.queue.length > 0 || isBlocked(gate)) return; + clearCleanupTimer(gate); + gates.delete(semaphoreKey); +} + +function scheduleCleanup(semaphoreKey: string): void { + const gate = gates.get(semaphoreKey); + if (!gate) return; + clearCleanupTimer(gate); + + gate.cleanupTimer = setTimeout(() => { + gate.cleanupTimer = null; + cleanupGateIfIdle(semaphoreKey); + }, 0); + + gate.cleanupTimer.unref?.(); +} + +function drainQueue(semaphoreKey: string): void { + const gate = gates.get(semaphoreKey); + if (!gate) return; + + while (gate.queue.length > 0 && gate.running < gate.maxConcurrency && !isBlocked(gate)) { + const next = gate.queue.shift(); + if (!next) break; + clearTimeout(next.timer); + gate.running++; + next.resolve(createReleaseFn(semaphoreKey)); + } + + if (gate.running === 0 && gate.queue.length === 0) { + scheduleCleanup(semaphoreKey); + } +} + +function createReleaseFn(semaphoreKey: string): () => void { + let released = false; + + return () => { + if (released) return; + released = true; + + const gate = gates.get(semaphoreKey); + if (!gate) return; + if (gate.running > 0) { + gate.running--; + } + + if (gate.queue.length > 0) { + drainQueue(semaphoreKey); + return; + } + + scheduleCleanup(semaphoreKey); + }; +} + +function createSemaphoreTimeoutError( + semaphoreKey: string, + timeoutMs: number +): Error & { code: string } { + const error = new Error(`Semaphore timeout after ${timeoutMs}ms for ${semaphoreKey}`) as Error & { + code: string; + }; + error.code = "SEMAPHORE_TIMEOUT"; + return error; +} + +/** + * Acquire a slot for a provider/model/account tuple. + * Returns an idempotent release function that is safe to call in finally blocks. + */ +export function acquire( + semaphoreKey: string, + { maxConcurrency = null, timeoutMs = DEFAULT_TIMEOUT_MS }: AcquireAccountSemaphoreOptions = {} +): Promise<() => void> { + if (isBypassed(maxConcurrency)) { + return Promise.resolve(createNoopReleaseFn()); + } + + const gate = ensureGate(semaphoreKey, maxConcurrency); + clearCleanupTimer(gate); + + if (gate.running < gate.maxConcurrency && !isBlocked(gate)) { + gate.running++; + return Promise.resolve(createReleaseFn(semaphoreKey)); + } + + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + const nextGate = gates.get(semaphoreKey); + if (!nextGate) { + reject(createSemaphoreTimeoutError(semaphoreKey, timeoutMs)); + return; + } + + const queueIndex = nextGate.queue.findIndex((item) => item.timer === timer); + if (queueIndex !== -1) { + nextGate.queue.splice(queueIndex, 1); + } + + if (nextGate.running === 0 && nextGate.queue.length === 0) { + scheduleCleanup(semaphoreKey); + } + + reject(createSemaphoreTimeoutError(semaphoreKey, timeoutMs)); + }, timeoutMs); + + timer.unref?.(); + gate.queue.push({ resolve, reject, timer }); + }); +} + +/** + * Temporarily block new acquisitions for a key while allowing in-flight requests to finish. + */ +export function markBlocked(semaphoreKey: string, cooldownMs: number): void { + const safeCooldownMs = Number.isFinite(cooldownMs) && cooldownMs > 0 ? cooldownMs : 0; + if (safeCooldownMs <= 0) { + const gate = gates.get(semaphoreKey); + if (!gate) return; + gate.blockedUntil = null; + drainQueue(semaphoreKey); + return; + } + + const gate = gates.get(semaphoreKey) ?? ensureGate(semaphoreKey, 1); + clearCleanupTimer(gate); + gate.blockedUntil = Date.now() + safeCooldownMs; + + const timer = setTimeout(() => { + const nextGate = gates.get(semaphoreKey); + if (!nextGate) return; + if (nextGate.blockedUntil && Date.now() >= nextGate.blockedUntil) { + nextGate.blockedUntil = null; + drainQueue(semaphoreKey); + if (nextGate.running === 0 && nextGate.queue.length === 0) { + scheduleCleanup(semaphoreKey); + } + } + }, safeCooldownMs + 50); + + timer.unref?.(); +} + +/** + * Return the current in-memory semaphore snapshot. + */ +export function getStats(): Record { + const stats: Record = {}; + + for (const [key, gate] of gates) { + stats[key] = { + running: gate.running, + queued: gate.queue.length, + maxConcurrency: gate.maxConcurrency, + blockedUntil: gate.blockedUntil ? new Date(gate.blockedUntil).toISOString() : null, + }; + } + + return stats; +} + +/** + * Reset a single key and reject queued waiters. + */ +export function reset(semaphoreKey: string): void { + const gate = gates.get(semaphoreKey); + if (!gate) return; + + clearCleanupTimer(gate); + for (const entry of gate.queue) { + clearTimeout(entry.timer); + entry.reject(new Error("Semaphore reset")); + } + gates.delete(semaphoreKey); +} + +/** + * Reset all keys and reject queued waiters. + */ +export function resetAll(): void { + for (const key of gates.keys()) { + reset(key); + } +} diff --git a/src/app/(dashboard)/dashboard/combos/page.tsx b/src/app/(dashboard)/dashboard/combos/page.tsx index e5932b019a..64a4b91853 100644 --- a/src/app/(dashboard)/dashboard/combos/page.tsx +++ b/src/app/(dashboard)/dashboard/combos/page.tsx @@ -142,8 +142,10 @@ const ADVANCED_FIELD_HELP_FALLBACK = { retryDelay: "Initial delay between retries. Higher values reduce burst pressure.", timeout: "Maximum request time before aborting. Set higher for long generations.", healthcheck: "Skips unhealthy models/providers from routing decisions when enabled.", - concurrencyPerModel: "Max simultaneous requests sent to each model in round-robin.", - queueTimeout: "How long a request can wait in queue before timeout in round-robin.", + concurrencyPerModel: + "Round-robin combo/model limit: max simultaneous requests sent to each model target. This is separate from any provider account-only cap.", + queueTimeout: + "How long a request can wait for a round-robin model slot before timing out. This queue is separate from any account-only concurrency cap.", }; const STRATEGY_RECOMMENDATIONS_FALLBACK = { diff --git a/src/app/(dashboard)/dashboard/providers/[id]/page.tsx b/src/app/(dashboard)/dashboard/providers/[id]/page.tsx index 2b9d0abfc4..014c55dd42 100644 --- a/src/app/(dashboard)/dashboard/providers/[id]/page.tsx +++ b/src/app/(dashboard)/dashboard/providers/[id]/page.tsx @@ -542,6 +542,7 @@ interface EditConnectionModalConnection { name?: string; email?: string; priority?: number; + maxConcurrent?: number | null; authType?: string; provider?: string; providerSpecificData?: Record; @@ -5807,6 +5808,7 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec const [formData, setFormData] = useState({ name: "", priority: 1, + maxConcurrent: "", apiKey: "", healthCheckInterval: 60, baseUrl: "", @@ -5871,6 +5873,10 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec setFormData({ name: connection.name || "", priority: connection.priority || 1, + maxConcurrent: + connection.maxConcurrent === null || connection.maxConcurrent === undefined + ? "" + : String(connection.maxConcurrent), apiKey: "", healthCheckInterval: connection.healthCheckInterval ?? 60, baseUrl: existingBaseUrl || defaultBaseUrl, @@ -5963,9 +5969,21 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec setSaving(true); setSaveError(null); try { + const trimmedMaxConcurrent = formData.maxConcurrent.trim(); + let parsedMaxConcurrent: number | null = null; + if (trimmedMaxConcurrent) { + const numericMaxConcurrent = Number(trimmedMaxConcurrent); + if (!Number.isInteger(numericMaxConcurrent) || numericMaxConcurrent < 0) { + setSaveError("Max concurrent must be a whole number greater than or equal to 0."); + return; + } + parsedMaxConcurrent = numericMaxConcurrent; + } + const updates: any = { name: formData.name, priority: formData.priority, + maxConcurrent: parsedMaxConcurrent, healthCheckInterval: formData.healthCheckInterval, }; @@ -6211,6 +6229,30 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec setFormData({ ...formData, priority: Number.parseInt(e.target.value) || 1 }) } /> + { + const nextValue = e.target.value; + setFormData({ ...formData, maxConcurrent: nextValue }); + if (saveError && nextValue.trim()) { + const numericValue = Number(nextValue); + if (Number.isInteger(numericValue) && numericValue >= 0) { + setSaveError(null); + } + } + }} + placeholder="0" + hint={t("accountConcurrencyCapHint")} + /> + {saveError && ( +
+ {saveError} +
+ )} {!isOAuth && ( <>
@@ -6256,11 +6298,6 @@ function EditConnectionModal({ isOpen, connection, onSave, onClose }: EditConnec {validationResult === "success" ? t("valid") : t("invalid")} )} - {saveError && ( -
- {saveError} -
- )}
))} diff --git a/src/app/api/providers/[id]/route.ts b/src/app/api/providers/[id]/route.ts index 88ee650319..ea163abef9 100644 --- a/src/app/api/providers/[id]/route.ts +++ b/src/app/api/providers/[id]/route.ts @@ -115,6 +115,7 @@ export async function PUT(request: Request, { params }: { params: Promise<{ id: rateLimitedUntil, lastTested, healthCheckInterval, + maxConcurrent, providerSpecificData: incomingPsd, } = body; @@ -139,6 +140,7 @@ export async function PUT(request: Request, { params }: { params: Promise<{ id: if (rateLimitedUntil !== undefined) updateData.rateLimitedUntil = rateLimitedUntil; if (lastTested !== undefined) updateData.lastTested = lastTested; if (healthCheckInterval !== undefined) updateData.healthCheckInterval = healthCheckInterval; + if (maxConcurrent !== undefined) updateData.maxConcurrent = maxConcurrent; // Merge providerSpecificData (partial update — preserve existing keys not sent by caller) if (incomingPsd !== undefined && incomingPsd !== null && typeof incomingPsd === "object") { @@ -160,7 +162,8 @@ export async function PUT(request: Request, { params }: { params: Promise<{ id: } updateData.providerSpecificData = - normalizeProviderSpecificData(existing.provider, mergedPsd) || {}; + normalizeProviderSpecificData(existing.provider as string | null | undefined, mergedPsd) || + {}; } const updated = await updateProviderConnection(id, updateData); diff --git a/src/i18n/messages/en.json b/src/i18n/messages/en.json index 9e1726d107..04f1f65d79 100644 --- a/src/i18n/messages/en.json +++ b/src/i18n/messages/en.json @@ -1896,6 +1896,8 @@ "save": "Save", "editConnection": "Edit Connection", "accountName": "Account name", + "accountConcurrencyCapLabel": "Provider account max concurrent requests", + "accountConcurrencyCapHint": "Maximum simultaneous requests for this provider account. Leave empty or set 0 for no limit. Helps prevent limits enforced by the provider itself before requests fail.", "email": "Email", "healthCheckMinutes": "Health Check (min)", "healthCheckHint": "Proactive token refresh interval. 0 = disabled.", @@ -2303,6 +2305,8 @@ "defaultSafetyNet": "Default Safety Net", "rpm": "RPM", "minGap": "Min Gap", + "comboConcurrencyLabel": "Combo round-robin max concurrent requests per model", + "comboConcurrencyHint": "Maximum simultaneous requests per model during combo round-robin. This setting only applies to combo strategies.", "maxConcurrent": "Max Concurrent", "activeLimiters": "Active Limiters", "noActiveLimiters": "No active rate limiters yet.", diff --git a/src/i18n/messages/ko.json b/src/i18n/messages/ko.json index b32c5bcbee..9d0d35db55 100644 --- a/src/i18n/messages/ko.json +++ b/src/i18n/messages/ko.json @@ -1845,6 +1845,8 @@ "save": "저장", "editConnection": "연결 편집", "accountName": "계정 이름", + "accountConcurrencyCapLabel": "Provider 계정 최대 동시 요청 수", + "accountConcurrencyCapHint": "이 provider 계정의 최대 동시 요청 수입니다. 비워 두거나 0으로 설정하면 제한이 없습니다. provider 자체가 부과하는 제한에 먼저 걸리지 않도록 돕습니다.", "email": "이메일", "healthCheckMinutes": "상태 점검(분)", "healthCheckHint": "사전 토큰 새로 고침 간격. 0 = 비활성화됨.", @@ -2215,6 +2217,8 @@ "defaultSafetyNet": "기본 안전망", "rpm": "RPM", "minGap": "민갭", + "comboConcurrencyLabel": "콤보 라운드로빈 모델별 최대 동시 요청 수", + "comboConcurrencyHint": "Combo 라운드로빈에서 모델별 최대 동시 요청 수입니다. 이 설정은 combo 전략에만 적용됩니다.", "maxConcurrent": "최대 동시", "activeLimiters": "활성 리미터", "noActiveLimiters": "아직 활성 속도 제한기가 없습니다.", diff --git a/src/lib/db/migrations/028_provider_connection_max_concurrent.sql b/src/lib/db/migrations/028_provider_connection_max_concurrent.sql new file mode 100644 index 0000000000..c9d455c20f --- /dev/null +++ b/src/lib/db/migrations/028_provider_connection_max_concurrent.sql @@ -0,0 +1,10 @@ +-- 028_provider_connection_max_concurrent.sql +-- Adds account-native concurrency cap to provider_connections. +-- This defines the maximum concurrent requests allowed for a specific account (connection). +-- Coexists with existing combo-level concurrencyPerModel, which is separate. + +-- Add max_concurrent column (NULL = unlimited, uses provider defaults or combo rules) +ALTER TABLE provider_connections ADD COLUMN max_concurrent INTEGER; + +-- Index for provider-level filtering +CREATE INDEX IF NOT EXISTS idx_pc_max_concurrent ON provider_connections(provider, max_concurrent); \ No newline at end of file diff --git a/src/lib/db/providers.ts b/src/lib/db/providers.ts index 51dbcf03d6..b5e2c4c03d 100644 --- a/src/lib/db/providers.ts +++ b/src/lib/db/providers.ts @@ -21,6 +21,26 @@ interface DbLike { prepare: (sql: string) => StatementLike; } +function withNullableMaxConcurrent( + record: JsonRecord, + source: JsonRecord | null | undefined +): JsonRecord { + if (!source || !Object.hasOwn(source, "maxConcurrent")) { + return record; + } + + const sourceMaxConcurrent = source.maxConcurrent; + const normalizedMaxConcurrent = + typeof sourceMaxConcurrent === "number" || sourceMaxConcurrent === null + ? sourceMaxConcurrent + : record.maxConcurrent; + + return { + ...record, + maxConcurrent: normalizedMaxConcurrent, + }; +} + function toRecord(value: unknown): JsonRecord { return value && typeof value === "object" ? (value as JsonRecord) : {}; } @@ -56,13 +76,19 @@ export async function getProviderConnections(filter: JsonRecord = {}) { sql += " ORDER BY priority ASC, updated_at DESC"; const rows = db.prepare(sql).all(params); - return rows.map((r) => decryptConnectionFields(cleanNulls(rowToCamel(r)))); + return rows.map((r) => { + const camelRow = rowToCamel(r); + return decryptConnectionFields(withNullableMaxConcurrent(cleanNulls(camelRow), camelRow)); + }); } export async function getProviderConnectionById(id: string) { const db = getDbInstance() as unknown as DbLike; const row = db.prepare("SELECT * FROM provider_connections WHERE id = ?").get(id); - return row ? decryptConnectionFields(cleanNulls(rowToCamel(row))) : null; + if (!row) return null; + + const camelRow = rowToCamel(row); + return decryptConnectionFields(withNullableMaxConcurrent(cleanNulls(camelRow), camelRow)); } export async function createProviderConnection(data: JsonRecord) { @@ -133,7 +159,7 @@ export async function createProviderConnection(data: JsonRecord) { ); _updateConnectionRow(db, existingId, merged); backupDbFile("pre-write"); - return cleanNulls(merged); + return withNullableMaxConcurrent(cleanNulls(merged), merged); } // Generate name: prefer explicit name, then email, then a stable short-ID label. @@ -195,6 +221,7 @@ export async function createProviderConnection(data: JsonRecord) { "consecutiveUseCount", "rateLimitProtection", "group", + "maxConcurrent", ]; for (const field of optionalFields) { if (data[field] !== undefined && data[field] !== null) { @@ -213,7 +240,7 @@ export async function createProviderConnection(data: JsonRecord) { backupDbFile("pre-write"); invalidateDbCache("connections"); // Bust connections read cache - return cleanNulls(connection); + return withNullableMaxConcurrent(cleanNulls(connection), connection); } function _insertConnectionRow(db: DbLike, conn: JsonRecord) { @@ -227,7 +254,8 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) { rate_limited_until, health_check_interval, last_health_check_at, last_tested, api_key, id_token, provider_specific_data, expires_in, display_name, global_priority, default_model, - token_type, consecutive_use_count, rate_limit_protection, last_used_at, "group", created_at, updated_at + token_type, consecutive_use_count, rate_limit_protection, last_used_at, "group", max_concurrent, + created_at, updated_at ) VALUES ( @id, @provider, @authType, @name, @email, @priority, @isActive, @accessToken, @refreshToken, @expiresAt, @tokenExpiresAt, @@ -236,7 +264,8 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) { @rateLimitedUntil, @healthCheckInterval, @lastHealthCheckAt, @lastTested, @apiKey, @idToken, @providerSpecificData, @expiresIn, @displayName, @globalPriority, @defaultModel, - @tokenType, @consecutiveUseCount, @rateLimitProtection, @lastUsedAt, @group, @createdAt, @updatedAt + @tokenType, @consecutiveUseCount, @rateLimitProtection, @lastUsedAt, @group, @maxConcurrent, + @createdAt, @updatedAt ) ` ).run({ @@ -279,6 +308,7 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) { conn.rateLimitProtection === true || conn.rateLimitProtection === 1 ? 1 : 0, lastUsedAt: conn.lastUsedAt || null, group: conn.group || null, + maxConcurrent: conn.maxConcurrent ?? null, createdAt: conn.createdAt, updatedAt: conn.updatedAt, }); @@ -304,6 +334,7 @@ function _updateConnectionRow(db: DbLike, id: string, data: JsonRecord) { rate_limit_protection = @rateLimitProtection, last_used_at = @lastUsedAt, "group" = @group, + max_concurrent = @maxConcurrent, updated_at = @updatedAt WHERE id = @id ` @@ -347,6 +378,7 @@ function _updateConnectionRow(db: DbLike, id: string, data: JsonRecord) { data.rateLimitProtection === true || data.rateLimitProtection === 1 ? 1 : 0, lastUsedAt: data.lastUsedAt || null, group: data.group || null, + maxConcurrent: data.maxConcurrent ?? null, updatedAt: now, }); } @@ -378,7 +410,7 @@ export async function updateProviderConnection(id: string, data: JsonRecord) { _reorderConnections(db, providerId); } - return cleanNulls(merged); + return withNullableMaxConcurrent(cleanNulls(merged), merged); } export async function deleteProviderConnection(id: string) { diff --git a/src/shared/validation/schemas.ts b/src/shared/validation/schemas.ts index ae120f3da7..0307f35531 100644 --- a/src/shared/validation/schemas.ts +++ b/src/shared/validation/schemas.ts @@ -1406,6 +1406,7 @@ export const updateProviderConnectionSchema = z lastTested: z.union([z.string(), z.null()]).optional(), healthCheckInterval: z.coerce.number().int().min(0).optional(), group: z.union([z.string().max(100), z.null()]).optional(), + maxConcurrent: z.union([z.null(), z.coerce.number().int().min(0)]).optional(), // Partial patch of per-connection provider-specific settings (e.g. quota toggles) providerSpecificData: z .record(z.string(), z.unknown()) diff --git a/tests/unit/account-concurrency-cap.test.ts b/tests/unit/account-concurrency-cap.test.ts new file mode 100644 index 0000000000..425f27b565 --- /dev/null +++ b/tests/unit/account-concurrency-cap.test.ts @@ -0,0 +1,116 @@ +import { after, beforeEach, describe, it } 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-account-cap-")); +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 { updateProviderConnectionSchema } = await import("../../src/shared/validation/schemas.ts"); + +type Connection = Awaited>; +type StoredConnection = Awaited>; + +function assertConnection(connection: Connection): asserts connection is NonNullable { + assert.ok(connection); +} + +function assertStoredConnection( + connection: StoredConnection +): asserts connection is NonNullable { + assert.ok(connection); +} + +function getConnectionId(connection: NonNullable): string { + assert.equal(typeof connection.id, "string"); + return connection.id as string; +} + +async function resetStorage() { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +async function createConnection(maxConcurrent: number | null): Promise { + return providersDb.createProviderConnection({ + provider: "openai", + authType: "apikey", + name: `openai-${String(maxConcurrent)}-${Math.random().toString(16).slice(2, 8)}`, + apiKey: "sk-test", + maxConcurrent, + }); +} + +beforeEach(async () => { + await resetStorage(); +}); + +after(() => { + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +describe("maxConcurrent DB round-trip", () => { + it("stores and retrieves maxConcurrent=3", async () => { + const created = await createConnection(3); + assertConnection(created); + const connectionId = getConnectionId(created); + + const stored = await providersDb.getProviderConnectionById(connectionId); + assertStoredConnection(stored); + assert.equal(stored.maxConcurrent, 3); + }); + + it("stores and retrieves maxConcurrent=null", async () => { + const created = await createConnection(null); + assertConnection(created); + const connectionId = getConnectionId(created); + + const stored = await providersDb.getProviderConnectionById(connectionId); + assertStoredConnection(stored); + assert.equal(stored.maxConcurrent, null); + }); + + it("stores and retrieves maxConcurrent=0", async () => { + const created = await createConnection(3); + assertConnection(created); + const connectionId = getConnectionId(created); + + const updated = await providersDb.updateProviderConnection(connectionId, { maxConcurrent: 0 }); + assertConnection(updated); + assert.equal(updated.maxConcurrent, 0); + + const stored = await providersDb.getProviderConnectionById(connectionId); + assertStoredConnection(stored); + assert.equal(stored.maxConcurrent, 0); + }); +}); + +describe("maxConcurrent validation", () => { + it("rejects negative values", () => { + const result = updateProviderConnectionSchema.safeParse({ maxConcurrent: -1 }); + + assert.equal(result.success, false); + }); + + it("accepts positive integers", () => { + const result = updateProviderConnectionSchema.safeParse({ maxConcurrent: 3 }); + + assert.equal(result.success, true); + assert.equal(result.data.maxConcurrent, 3); + }); + + it("accepts null/undefined as unlimited", () => { + const nullResult = updateProviderConnectionSchema.safeParse({ maxConcurrent: null }); + const undefinedResult = updateProviderConnectionSchema.safeParse({ name: "unchanged" }); + + assert.equal(nullResult.success, true); + assert.equal(nullResult.data.maxConcurrent, null); + assert.equal(undefinedResult.success, true); + assert.equal(undefinedResult.data.maxConcurrent, undefined); + }); +});