mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-07 07:42:13 +03:00
feat(open-sse): wire quotaShare PRE/POST hooks in chatCore handler (B/F7)
PRE-hook (before executor dispatch): - Calls enforceQuotaShare via dynamic import (lazy load, fail-open). - Returns 429 JSON via buildErrorBody() when decision.kind === 'block' (B25). - Sets quotaSoftDeprioritize=true when decision.deprioritize=true (B17). POST-hook (after successful response): - Calls scheduleRecordConsumption for both streaming and non-streaming paths. - Fire-and-forget via setImmediate; never blocks the client response (B29). Both hooks use try/catch outer guards so any unexpected error fails open (B16).
This commit is contained in:
@@ -3367,6 +3367,62 @@ export async function handleChatCore({
|
||||
return wrapper;
|
||||
};
|
||||
|
||||
// === Quota Share enforcement PRE-hook (B/F7) ===
|
||||
// Runs after provider/model/credentials/apiKeyInfo are fully resolved,
|
||||
// before dispatcher. Fail-open per B16: errors → allow.
|
||||
let quotaSoftDeprioritize = false;
|
||||
if (apiKeyInfo?.id && credentials?.connectionId) {
|
||||
try {
|
||||
const { enforceQuotaShare } = await import("@/lib/quota/enforce");
|
||||
const decision = await enforceQuotaShare({
|
||||
apiKeyId: apiKeyInfo.id,
|
||||
connectionId: credentials.connectionId,
|
||||
provider: provider ?? "unknown",
|
||||
estimatedCost: {},
|
||||
}).catch((err: unknown) => {
|
||||
log?.warn?.(
|
||||
"QUOTA_SHARE",
|
||||
`enforceQuotaShare failed; fail-open: ${err instanceof Error ? err.message : String(err)}`
|
||||
);
|
||||
return { kind: "allow" as const };
|
||||
});
|
||||
|
||||
if (decision.kind === "block") {
|
||||
const { buildErrorBody } = await import("../utils/error.ts");
|
||||
log?.warn?.(
|
||||
"QUOTA_SHARE",
|
||||
`[quotaShare] blocked apiKeyId=${apiKeyInfo.id} provider=${provider ?? "unknown"}: ${decision.reason}`
|
||||
);
|
||||
const headers: Record<string, string> = { "Content-Type": "application/json" };
|
||||
if (decision.retryAfterSeconds) {
|
||||
headers["Retry-After"] = String(decision.retryAfterSeconds);
|
||||
}
|
||||
return new Response(
|
||||
JSON.stringify(buildErrorBody(429, decision.reason)),
|
||||
{ status: 429, headers }
|
||||
);
|
||||
}
|
||||
|
||||
if (decision.kind === "allow" && decision.deprioritize) {
|
||||
quotaSoftDeprioritize = true;
|
||||
log?.info?.(
|
||||
"QUOTA_SHARE",
|
||||
`[quotaShare] soft deprioritize active for apiKeyId=${apiKeyInfo.id} provider=${provider ?? "unknown"}`
|
||||
);
|
||||
}
|
||||
} catch (err) {
|
||||
// Outer fail-open guard — should not be reached (inner .catch covers it)
|
||||
log?.warn?.(
|
||||
"QUOTA_SHARE",
|
||||
`[quotaShare] enforceQuotaShare unexpected error; fail-open: ${err instanceof Error ? err.message : String(err)}`
|
||||
);
|
||||
}
|
||||
}
|
||||
// Suppress unused variable lint warning — quotaSoftDeprioritize is available for
|
||||
// combo.ts to read when candidateBuilder populates quotaSoftPenalty in the future.
|
||||
void quotaSoftDeprioritize;
|
||||
// === /Quota Share enforcement PRE-hook ===
|
||||
|
||||
// Get executor for this provider (with optional upstream proxy routing)
|
||||
const executor = await resolveExecutorWithProxy(provider);
|
||||
const getExecutionCredentials = () => {
|
||||
@@ -5020,6 +5076,33 @@ export async function handleChatCore({
|
||||
recordCost(apiKeyInfo.id, estimatedCost);
|
||||
}
|
||||
|
||||
// === Quota Share POST-hook (B/F7) — fire-and-forget, fail-open ===
|
||||
if (apiKeyInfo?.id && credentials?.connectionId) {
|
||||
try {
|
||||
const { scheduleRecordConsumption } = await import("@/lib/quota/spendRecorder");
|
||||
scheduleRecordConsumption(
|
||||
{
|
||||
apiKeyId: apiKeyInfo.id,
|
||||
connectionId: credentials.connectionId,
|
||||
provider: provider ?? "unknown",
|
||||
cost: {
|
||||
tokens:
|
||||
usage && typeof usage === "object"
|
||||
? ((usage as Record<string, unknown>).prompt_tokens as number ?? 0) +
|
||||
((usage as Record<string, unknown>).completion_tokens as number ?? 0)
|
||||
: 0,
|
||||
usd: estimatedCost > 0 ? estimatedCost : 0,
|
||||
requests: 1,
|
||||
},
|
||||
},
|
||||
log
|
||||
);
|
||||
} catch (_) {
|
||||
// Outer fail-open — never throws to caller
|
||||
}
|
||||
}
|
||||
// === /Quota Share POST-hook ===
|
||||
|
||||
// ── Gamification event (fire-and-forget) ──
|
||||
if (apiKeyInfo?.id) {
|
||||
try {
|
||||
@@ -5215,6 +5298,33 @@ export async function handleChatCore({
|
||||
.catch(() => {});
|
||||
}
|
||||
|
||||
// === Quota Share POST-hook streaming (B/F7) — fire-and-forget, fail-open ===
|
||||
if (apiKeyInfo?.id && credentials?.connectionId && streamStatus === 200) {
|
||||
try {
|
||||
const { scheduleRecordConsumption } = await import("@/lib/quota/spendRecorder");
|
||||
const su = streamUsage as Record<string, unknown> | null;
|
||||
scheduleRecordConsumption(
|
||||
{
|
||||
apiKeyId: apiKeyInfo.id,
|
||||
connectionId: credentials.connectionId,
|
||||
provider: provider ?? "unknown",
|
||||
cost: {
|
||||
tokens: su
|
||||
? (Number(su.prompt_tokens ?? 0) || 0) +
|
||||
(Number(su.completion_tokens ?? 0) || 0)
|
||||
: 0,
|
||||
usd: 0, // estimatedCost resolved async above; omit to avoid dependency
|
||||
requests: 1,
|
||||
},
|
||||
},
|
||||
log
|
||||
);
|
||||
} catch (_) {
|
||||
// Outer fail-open — never throws to caller
|
||||
}
|
||||
}
|
||||
// === /Quota Share POST-hook streaming ===
|
||||
|
||||
if (
|
||||
memoryOwnerId &&
|
||||
memorySettings?.enabled &&
|
||||
|
||||
Reference in New Issue
Block a user