mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-06 15:22:12 +03:00
feat(api): wire shared admission across LLM routes
Acquire admission once after API-key policy, preserve lazy raw-request snapshots, and bind lease settlement to JSON, SSE, abort, deadline, and failure lifecycles. Expose a low-cardinality health summary and preserve non-SSE Ollama errors unchanged.
This commit is contained in:
@@ -56,6 +56,7 @@ export async function GET() {
|
||||
sessionManagerModule,
|
||||
credentialHealthModule,
|
||||
localHealthModule,
|
||||
adaptiveAdmissionModule,
|
||||
settingsResult,
|
||||
connectionsResult,
|
||||
] = await Promise.allSettled([
|
||||
@@ -67,6 +68,7 @@ export async function GET() {
|
||||
import("@omniroute/open-sse/services/sessionManager.ts"),
|
||||
import("@/lib/credentialHealth/cache"),
|
||||
import("@/lib/localHealthCheck"),
|
||||
import("@omniroute/open-sse/services/admission/runtime.ts"),
|
||||
getCachedSettings(),
|
||||
getProviderConnections(),
|
||||
]);
|
||||
@@ -145,6 +147,14 @@ export async function GET() {
|
||||
: {};
|
||||
const settings = settingsResult.status === "fulfilled" ? settingsResult.value : {};
|
||||
const connections = connectionsResult.status === "fulfilled" ? connectionsResult.value : [];
|
||||
const adaptiveAdmission =
|
||||
adaptiveAdmissionModule.status === "fulfilled"
|
||||
? readHealthValue(
|
||||
"adaptive admission",
|
||||
() => adaptiveAdmissionModule.value.getAdaptiveAdmissionRuntime().snapshot(),
|
||||
null
|
||||
)
|
||||
: null;
|
||||
|
||||
const payload = buildHealthPayload({
|
||||
appVersion: APP_CONFIG.version,
|
||||
@@ -169,6 +179,7 @@ export async function GET() {
|
||||
activeSessions,
|
||||
activeSessionsByKey,
|
||||
credentialHealth,
|
||||
adaptiveAdmission,
|
||||
});
|
||||
|
||||
healthPayloadCache = { payload, expiresAt: Date.now() + HEALTH_PAYLOAD_TTL_MS };
|
||||
@@ -186,6 +197,7 @@ export async function GET() {
|
||||
lockouts: [],
|
||||
quotaMonitor: { ...fallbackQuotaMonitorSummary, monitors: [] },
|
||||
sessions: { activeCount: 0, stickyBoundCount: 0, byApiKey: {}, top: [] },
|
||||
adaptiveAdmission: null,
|
||||
dedup: { inflightRequests: 0 },
|
||||
});
|
||||
}
|
||||
|
||||
@@ -82,6 +82,7 @@ export async function POST(request: Request) {
|
||||
method: request.method,
|
||||
headers: request.headers,
|
||||
body: JSON.stringify(normalized),
|
||||
signal: request.signal,
|
||||
});
|
||||
// #3571 — translate the chat-pipeline response back to the legacy
|
||||
// text-completion shape so OpenAI Completion clients (e.g. TabbyML) work.
|
||||
@@ -90,7 +91,7 @@ export async function POST(request: Request) {
|
||||
// echo the compression header on the way out.
|
||||
return withCompressionHeaderEcho(
|
||||
await asTextCompletionResponse(
|
||||
await handleChat(newRequest, buildClientRawRequest(request, body)),
|
||||
await handleChat(newRequest, () => buildClientRawRequest(request, body)),
|
||||
typeof body.model === "string" ? body.model : undefined
|
||||
),
|
||||
compressionRequestHeader
|
||||
@@ -106,7 +107,10 @@ export async function POST(request: Request) {
|
||||
// Re-read body.model so the response echoes the caller's requested identifier.
|
||||
let requestedModel: string | undefined;
|
||||
try {
|
||||
const bodyForModel = await request.clone().json().catch(() => null);
|
||||
const bodyForModel = await request
|
||||
.clone()
|
||||
.json()
|
||||
.catch(() => null);
|
||||
if (bodyForModel && typeof bodyForModel.model === "string") {
|
||||
requestedModel = bodyForModel.model;
|
||||
}
|
||||
|
||||
@@ -98,7 +98,8 @@ export async function POST(request, { params }) {
|
||||
method: request.method,
|
||||
headers: request.headers,
|
||||
body: JSON.stringify(body),
|
||||
signal: request.signal,
|
||||
});
|
||||
|
||||
return await handleChat(newRequest, buildClientRawRequest(request, rawBody));
|
||||
return await handleChat(newRequest, () => buildClientRawRequest(request, rawBody));
|
||||
}
|
||||
|
||||
@@ -89,9 +89,7 @@ export async function POST(request, { params }) {
|
||||
action = modelAction.includes(":streamGenerateContent")
|
||||
? ":streamGenerateContent"
|
||||
: ":generateContent";
|
||||
model = modelAction
|
||||
.replace(":streamGenerateContent", "")
|
||||
.replace(":generateContent", "");
|
||||
model = modelAction.replace(":streamGenerateContent", "").replace(":generateContent", "");
|
||||
}
|
||||
|
||||
const validation = validateBody(v1betaGeminiGenerateSchema, rawBody);
|
||||
@@ -113,9 +111,10 @@ export async function POST(request, { params }) {
|
||||
method: "POST",
|
||||
headers: request.headers,
|
||||
body: JSON.stringify(convertedBody),
|
||||
signal: request.signal,
|
||||
});
|
||||
|
||||
const response = await handleChat(newRequest, buildClientRawRequest(request, rawBody));
|
||||
const response = await handleChat(newRequest, () => buildClientRawRequest(request, rawBody));
|
||||
|
||||
if (stream) {
|
||||
// Transform OpenAI SSE => Gemini SSE on the fly. The @google/genai SDK
|
||||
|
||||
@@ -1,5 +1,63 @@
|
||||
import type { AdaptiveAdmissionPublicSnapshot } from "@omniroute/open-sse/services/admission/runtime.ts";
|
||||
|
||||
type JsonRecord = Record<string, unknown>;
|
||||
|
||||
/** Low-card adaptive-admission health summary — no tenant/request/body/queue details. */
|
||||
export type AdaptiveAdmissionHealthSummary = {
|
||||
mode: AdaptiveAdmissionPublicSnapshot["mode"];
|
||||
currentLimit: number;
|
||||
minLimit: number;
|
||||
maxLimit: number;
|
||||
activeCost: number;
|
||||
activeCount: number;
|
||||
queuedCost: number;
|
||||
queuedCount: number;
|
||||
admittedCount: number;
|
||||
rejectedCount: number;
|
||||
wouldAdmitCount: number;
|
||||
wouldQueueCount: number;
|
||||
wouldRejectCount: number;
|
||||
utilization: number;
|
||||
pressure: AdaptiveAdmissionPublicSnapshot["pressure"];
|
||||
resourceSeverity: AdaptiveAdmissionPublicSnapshot["resourceSeverity"];
|
||||
resourceReason: AdaptiveAdmissionPublicSnapshot["resourceReason"];
|
||||
resourceObservedAtMs: number;
|
||||
pressureGuardRejectCount: number;
|
||||
shutdown: boolean;
|
||||
};
|
||||
|
||||
/**
|
||||
* Explicit allowlisted projection of the public adaptive-admission snapshot.
|
||||
* Never spreads the snapshot — extra keys (tenant, body, queue items, paths) are dropped.
|
||||
*/
|
||||
export function projectAdaptiveAdmissionSummary(
|
||||
snapshot: AdaptiveAdmissionPublicSnapshot | null | undefined
|
||||
): AdaptiveAdmissionHealthSummary | null {
|
||||
if (!snapshot || typeof snapshot !== "object") return null;
|
||||
return {
|
||||
mode: snapshot.mode,
|
||||
currentLimit: snapshot.currentLimit,
|
||||
minLimit: snapshot.minLimit,
|
||||
maxLimit: snapshot.maxLimit,
|
||||
activeCost: snapshot.activeCost,
|
||||
activeCount: snapshot.activeCount,
|
||||
queuedCost: snapshot.queuedCost,
|
||||
queuedCount: snapshot.queuedCount,
|
||||
admittedCount: snapshot.admittedCount,
|
||||
rejectedCount: snapshot.rejectedCount,
|
||||
wouldAdmitCount: snapshot.wouldAdmitCount,
|
||||
wouldQueueCount: snapshot.wouldQueueCount,
|
||||
wouldRejectCount: snapshot.wouldRejectCount,
|
||||
utilization: snapshot.utilization,
|
||||
pressure: snapshot.pressure,
|
||||
resourceSeverity: snapshot.resourceSeverity,
|
||||
resourceReason: snapshot.resourceReason,
|
||||
resourceObservedAtMs: snapshot.resourceObservedAtMs,
|
||||
pressureGuardRejectCount: snapshot.pressureGuardRejectCount,
|
||||
shutdown: snapshot.shutdown,
|
||||
};
|
||||
}
|
||||
|
||||
interface CircuitBreakerStatus {
|
||||
name: string;
|
||||
state: string;
|
||||
@@ -88,6 +146,8 @@ interface BuildHealthPayloadOptions {
|
||||
unknown: number;
|
||||
stale: number;
|
||||
};
|
||||
/** Optional injected public adaptive-admission snapshot; projected, never raw-spread. */
|
||||
adaptiveAdmission?: AdaptiveAdmissionPublicSnapshot | null;
|
||||
}
|
||||
|
||||
function limitMonitors(monitors: QuotaMonitorSnapshot[], maxItems = 8): QuotaMonitorSnapshot[] {
|
||||
@@ -227,6 +287,7 @@ export function buildHealthPayload({
|
||||
activeSessions,
|
||||
activeSessionsByKey = {},
|
||||
credentialHealth,
|
||||
adaptiveAdmission = null,
|
||||
}: BuildHealthPayloadOptions) {
|
||||
const timestamp = new Date().toISOString();
|
||||
const system = {
|
||||
@@ -321,6 +382,7 @@ export function buildHealthPayload({
|
||||
},
|
||||
sessions: buildSessionsSummary({ activeSessions, activeSessionsByKey }),
|
||||
credentialHealth, // may be undefined if credentialHealth module not loaded
|
||||
adaptiveAdmission: projectAdaptiveAdmissionSummary(adaptiveAdmission),
|
||||
dedup: {
|
||||
inflightRequests,
|
||||
},
|
||||
|
||||
@@ -1,5 +1,8 @@
|
||||
import { randomUUID } from "crypto";
|
||||
import { resolveChatRequestBody } from "./requestBody";
|
||||
import * as chatAdmission from "./chatAdmission.ts";
|
||||
import { buildClientRawRequest, resolveDispatchClientRawRequest } from "./chat/clientRawRequest.ts";
|
||||
export { buildClientRawRequest, resolveDispatchClientRawRequest };
|
||||
import { normalizeReasoningRequest } from "@/shared/reasoning/effortStandardization";
|
||||
import { resolveRoutingModel, RoutingModelOps } from "./resolveRoutingModel";
|
||||
import {
|
||||
@@ -64,6 +67,7 @@ import { guardrailRegistry, resolveDisabledGuardrails } from "@/lib/guardrails";
|
||||
import {
|
||||
resolveModelOrError,
|
||||
checkPipelineGates,
|
||||
checkResourcePressureBeforeProviderWork,
|
||||
executeChatWithBreaker,
|
||||
handleNoCredentials,
|
||||
safeResolveProxy,
|
||||
@@ -232,16 +236,12 @@ const comboPromoteDeps = { updateCombo, info: log.info, warn: log.warn };
|
||||
|
||||
export { shouldTripProviderBreakerForResult } from "./chatPredicates";
|
||||
|
||||
/**
|
||||
* Handle chat completion request
|
||||
* Supports: OpenAI, Claude, Gemini, OpenAI Responses API formats
|
||||
* Format detection and translation handled by translator
|
||||
*/
|
||||
export async function handleChat(
|
||||
async function handleChatImplementation(
|
||||
request: any,
|
||||
clientRawRequest: any = null,
|
||||
preParsedBody: any = null,
|
||||
correlationId?: string
|
||||
correlationId: string | undefined,
|
||||
admissionContext: chatAdmission.ChatAdmissionContext
|
||||
) {
|
||||
const peerRejection = rejectPeerRequest(request?.headers, log.warn, errorResponse);
|
||||
if (peerRejection) return peerRejection;
|
||||
@@ -357,11 +357,7 @@ export async function handleChat(
|
||||
}
|
||||
}
|
||||
|
||||
// buildClientRawRequest already deep-clones the body, so pass `body` directly — the
|
||||
// prior local clone was a redundant second full-body copy on the hot path (#5152).
|
||||
if (!clientRawRequest) {
|
||||
clientRawRequest = buildClientRawRequest(request, body);
|
||||
}
|
||||
const deferredClientRawBody = chatAdmission.captureDeferredClientRawBody(body);
|
||||
|
||||
// T01 — Accept-header streaming opt-in (#302 / #5305). A bare `Accept:
|
||||
// text/event-stream` with `stream` omitted opts a curl/httpx-style client into
|
||||
@@ -488,6 +484,12 @@ export async function handleChat(
|
||||
const bypassProviderQuotaPolicy = hasProviderQuotaBypassScope(apiKeyInfo?.scopes);
|
||||
telemetry.endPhase();
|
||||
|
||||
const admissionRejection = await admissionContext.acquire(apiKeyInfo?.id, request, body);
|
||||
if (admissionRejection) return admissionRejection;
|
||||
clientRawRequest = chatAdmission.resolveClientRawAfterAdmission(clientRawRequest, () =>
|
||||
deferredClientRawBody.withClientBody((clientBody) => buildClientRawRequest(request, clientBody))
|
||||
);
|
||||
|
||||
// Guardrail pre-call pipeline — prompt injection, PII masking, and future custom rules.
|
||||
telemetry.startPhase("validate");
|
||||
const preCallGuardrails = await guardrailRegistry.runPreCallHooks(body, {
|
||||
@@ -972,18 +974,9 @@ export async function handleChat(
|
||||
return withCorrelationId(withSessionHeader(response, sessionId), reqId);
|
||||
}
|
||||
|
||||
// The clientRawRequest envelope lives in ./chat/clientRawRequest.ts. Imported for local use
|
||||
// below and re-exported for the historical public surface.
|
||||
import { buildClientRawRequest, resolveDispatchClientRawRequest } from "./chat/clientRawRequest.ts";
|
||||
export { buildClientRawRequest, resolveDispatchClientRawRequest };
|
||||
export const handleChat = chatAdmission.withChatAdmission(handleChatImplementation);
|
||||
|
||||
/**
|
||||
* Handle single model chat request
|
||||
*
|
||||
* Refactored: model resolution, logging, pipeline gates, and chat execution
|
||||
* extracted to focused helpers. This function orchestrates the credential
|
||||
* retry loop.
|
||||
*/
|
||||
/** Handle one resolved model through gates, credentials, and retry/fallback. */
|
||||
async function handleSingleModelChat(
|
||||
body: any,
|
||||
modelStr: string,
|
||||
@@ -1147,7 +1140,9 @@ async function handleSingleModelChat(
|
||||
? "fixed combo step connection"
|
||||
: undefined;
|
||||
|
||||
// 2. Pipeline gates (availability + provider circuit breaker)
|
||||
// 2. Local pressure precedes availability/breaker gates and account selection.
|
||||
const pressureGuard = checkResourcePressureBeforeProviderWork();
|
||||
if (pressureGuard) return pressureGuard.response;
|
||||
const providerProfile = await getRuntimeProviderProfile(provider);
|
||||
const gate = await checkPipelineGates(provider, model, {
|
||||
ignoreCircuitBreaker: forceLiveComboTest || hasForcedConnection,
|
||||
@@ -1426,7 +1421,7 @@ async function handleSingleModelChat(
|
||||
clientRawRequest,
|
||||
runtimeOptions.modelAbortSignal
|
||||
);
|
||||
const { result, tlsFingerprintUsed } = await executeChatWithBreaker({
|
||||
const execution = await executeChatWithBreaker({
|
||||
bypassCircuitBreaker: forceLiveComboTest || hasForcedConnection,
|
||||
breaker,
|
||||
body: requestBody,
|
||||
@@ -1455,6 +1450,10 @@ async function handleSingleModelChat(
|
||||
routingComboId: runtimeOptions?.routingComboId ?? null,
|
||||
});
|
||||
if (telemetry) telemetry.endPhase();
|
||||
if ("localResourcePressureResult" in execution) {
|
||||
return execution.localResourcePressureResult.response;
|
||||
}
|
||||
const { result, tlsFingerprintUsed } = execution;
|
||||
|
||||
const proxyLatency = Date.now() - proxyStartTime;
|
||||
const providerAlias = PROVIDER_ID_TO_ALIAS[provider] || provider;
|
||||
|
||||
239
src/sse/handlers/chatAdmission.ts
Normal file
239
src/sse/handlers/chatAdmission.ts
Normal file
@@ -0,0 +1,239 @@
|
||||
/**
|
||||
* Shared handleChat adaptive-admission lifecycle wrapper.
|
||||
*
|
||||
* Owns a per-call context that acquires exactly once after API-key policy and
|
||||
* attaches/releases the admitted lease around the handler response or throw.
|
||||
* No AsyncLocalStorage, no route registry — one higher-order wrapper only.
|
||||
*/
|
||||
|
||||
import {
|
||||
getAdaptiveAdmissionRuntime,
|
||||
type AdaptiveAdmissionAdmitted,
|
||||
type AdaptiveAdmissionFailureOutcome,
|
||||
type AdaptiveAdmissionRuntime,
|
||||
} from "@omniroute/open-sse/services/admission/runtime.ts";
|
||||
|
||||
/** Single fairness bucket for unauthenticated / keyless traffic. Opaque; never a raw key. */
|
||||
export const ANONYMOUS_ADMISSION_TENANT_KEY = "anonymous";
|
||||
|
||||
export type ChatAdmissionContext = {
|
||||
/**
|
||||
* Acquire once against the process runtime.
|
||||
* Returns a sanitized rejection Response, or null when admitted / already acquired.
|
||||
*/
|
||||
acquire(
|
||||
apiKeyId: string | null | undefined,
|
||||
request: { signal?: AbortSignal | null },
|
||||
body: unknown
|
||||
): Promise<Response | null>;
|
||||
};
|
||||
|
||||
type AdmittedState = {
|
||||
runtime: AdaptiveAdmissionRuntime;
|
||||
admitted: AdaptiveAdmissionAdmitted;
|
||||
};
|
||||
|
||||
export function resolveAdmissionTenantKey(apiKeyId: string | null | undefined): string {
|
||||
return typeof apiKeyId === "string" && apiKeyId.length > 0
|
||||
? apiKeyId
|
||||
: ANONYMOUS_ADMISSION_TENANT_KEY;
|
||||
}
|
||||
|
||||
const CANCEL_NAMES = new Set(["AbortError"]);
|
||||
const CANCEL_CODES = new Set(["ABORT_ERR", "ERR_CANCELED"]);
|
||||
const TIMEOUT_NAMES = new Set(["TimeoutError"]);
|
||||
const TIMEOUT_CODES = new Set(["ETIMEDOUT", "ESOCKETTIMEDOUT", "TIMEOUT", "ERR_TIMEOUT"]);
|
||||
|
||||
function asStringField(err: object, key: string): string {
|
||||
const value = (err as Record<string, unknown>)[key];
|
||||
return typeof value === "string" ? value : "";
|
||||
}
|
||||
|
||||
function asStatus(err: object): number | null {
|
||||
const status = (err as Record<string, unknown>).status;
|
||||
if (typeof status === "number") return status;
|
||||
const statusCode = (err as Record<string, unknown>).statusCode;
|
||||
return typeof statusCode === "number" ? statusCode : null;
|
||||
}
|
||||
|
||||
/** Classify thrown handler failure; never exposes raw errors to clients. */
|
||||
export function classifyHandlerFailure(
|
||||
err: unknown,
|
||||
signal?: AbortSignal | null
|
||||
): AdaptiveAdmissionFailureOutcome {
|
||||
if (signal?.aborted) return "cancelled";
|
||||
if (!err || typeof err !== "object") return "upstream_error";
|
||||
|
||||
const name = asStringField(err, "name");
|
||||
const code = asStringField(err, "code");
|
||||
if (CANCEL_NAMES.has(name) || CANCEL_CODES.has(code)) return "cancelled";
|
||||
if (TIMEOUT_NAMES.has(name) || TIMEOUT_CODES.has(code)) return "timeout";
|
||||
|
||||
const status = asStatus(err);
|
||||
if (status === 408 || status === 504) return "timeout";
|
||||
if (status !== null && status >= 400 && status < 500) return "local_reject";
|
||||
return "upstream_error";
|
||||
}
|
||||
|
||||
const CLIENT_RAW_MUTABLE_FIELDS = ["model", "reasoning", "reasoning_effort", "thinking"] as const;
|
||||
type ClientRawFieldState = {
|
||||
key: (typeof CLIENT_RAW_MUTABLE_FIELDS)[number];
|
||||
present: boolean;
|
||||
value: unknown;
|
||||
};
|
||||
|
||||
function captureClientRawFields(body: Record<string, unknown>): ClientRawFieldState[] {
|
||||
return CLIENT_RAW_MUTABLE_FIELDS.map((key) => {
|
||||
const present = Object.hasOwn(body, key);
|
||||
return { key, present, value: present ? body[key] : undefined };
|
||||
});
|
||||
}
|
||||
|
||||
function clientRawFieldsEqual(a: ClientRawFieldState[], b: ClientRawFieldState[]): boolean {
|
||||
return a.every(
|
||||
(field, index) =>
|
||||
field.key === b[index]?.key &&
|
||||
field.present === b[index]?.present &&
|
||||
Object.is(field.value, b[index]?.value)
|
||||
);
|
||||
}
|
||||
|
||||
function applyClientRawFields(body: Record<string, unknown>, fields: ClientRawFieldState[]): void {
|
||||
for (const field of fields) {
|
||||
if (field.present) body[field.key] = field.value;
|
||||
else delete body[field.key];
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Capture only the fixed fields mutated before admission. The full bounded observability
|
||||
* snapshot is built after admission without enumerating or cloning the body beforehand.
|
||||
*/
|
||||
export function captureDeferredClientRawBody(body: unknown): {
|
||||
withClientBody<T>(build: (clientBody: unknown) => T): T;
|
||||
} {
|
||||
const target =
|
||||
body !== null && typeof body === "object" ? (body as Record<string, unknown>) : null;
|
||||
const originalFields = target ? captureClientRawFields(target) : null;
|
||||
|
||||
return {
|
||||
withClientBody(build) {
|
||||
if (!target || !originalFields) return build(body);
|
||||
const workingFields = captureClientRawFields(target);
|
||||
if (clientRawFieldsEqual(originalFields, workingFields)) return build(target);
|
||||
|
||||
applyClientRawFields(target, originalFields);
|
||||
try {
|
||||
return build(target);
|
||||
} finally {
|
||||
applyClientRawFields(target, workingFields);
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
/** Resolve lazy/eager client-raw after admission; invoke factories at most once. */
|
||||
export function resolveClientRawAfterAdmission(
|
||||
clientRawRequest: unknown,
|
||||
build: () => unknown
|
||||
): unknown {
|
||||
if (typeof clientRawRequest === "function") {
|
||||
return (clientRawRequest as () => unknown)();
|
||||
}
|
||||
if (clientRawRequest) return clientRawRequest;
|
||||
return build();
|
||||
}
|
||||
|
||||
export function createChatAdmissionContext(
|
||||
getRuntime: () => AdaptiveAdmissionRuntime = getAdaptiveAdmissionRuntime
|
||||
): ChatAdmissionContext & { getAdmittedState(): AdmittedState | null } {
|
||||
let state: AdmittedState | null = null;
|
||||
let acquireStarted = false;
|
||||
|
||||
return {
|
||||
getAdmittedState: () => state,
|
||||
async acquire(apiKeyId, request, body) {
|
||||
// Exactly once per logical request — never re-enter the runtime.
|
||||
if (state || acquireStarted) return null;
|
||||
acquireStarted = true;
|
||||
|
||||
const runtime = getRuntime();
|
||||
const streaming =
|
||||
body !== null && typeof body === "object" && (body as { stream?: unknown }).stream === true;
|
||||
|
||||
const result = await runtime.acquire({
|
||||
tenantKey: resolveAdmissionTenantKey(apiKeyId),
|
||||
body,
|
||||
signal: request?.signal ?? undefined,
|
||||
streaming,
|
||||
});
|
||||
|
||||
if (result.status === "rejected") {
|
||||
return result.response;
|
||||
}
|
||||
|
||||
state = { runtime, admitted: result };
|
||||
return null;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
type HandleChatImplementation = (
|
||||
request: any,
|
||||
clientRawRequest: any,
|
||||
preParsedBody: any,
|
||||
correlationId: string | undefined,
|
||||
admissionContext: ChatAdmissionContext
|
||||
) => Promise<Response>;
|
||||
|
||||
export type WithChatAdmissionOptions = {
|
||||
/** Test seam: override process-global runtime resolution. */
|
||||
getRuntime?: () => AdaptiveAdmissionRuntime;
|
||||
};
|
||||
|
||||
/**
|
||||
* Thin public wrapper: create per-call context, run implementation, attach/release lease.
|
||||
*/
|
||||
export function withChatAdmission(
|
||||
implementation: HandleChatImplementation,
|
||||
options: WithChatAdmissionOptions = {}
|
||||
) {
|
||||
return async function handleChat(
|
||||
request: any,
|
||||
clientRawRequest: any = null,
|
||||
preParsedBody: any = null,
|
||||
correlationId?: string
|
||||
): Promise<Response> {
|
||||
const admissionContext = createChatAdmissionContext(
|
||||
options.getRuntime ?? getAdaptiveAdmissionRuntime
|
||||
);
|
||||
try {
|
||||
const response = await implementation(
|
||||
request,
|
||||
clientRawRequest,
|
||||
preParsedBody,
|
||||
correlationId,
|
||||
admissionContext
|
||||
);
|
||||
const admittedState = admissionContext.getAdmittedState();
|
||||
if (!admittedState) return response;
|
||||
|
||||
const { runtime, admitted } = admittedState;
|
||||
return runtime.attachResponseLifecycle(response, admitted.lease, {
|
||||
admittedAtMs: admitted.admittedAtMs,
|
||||
signal: request?.signal ?? undefined,
|
||||
});
|
||||
} catch (err) {
|
||||
const admittedState = admissionContext.getAdmittedState();
|
||||
if (admittedState) {
|
||||
const { runtime, admitted } = admittedState;
|
||||
runtime.releaseHandlerFailure(
|
||||
admitted.lease,
|
||||
classifyHandlerFailure(err, request?.signal),
|
||||
{ admittedAtMs: admitted.admittedAtMs }
|
||||
);
|
||||
}
|
||||
throw err;
|
||||
}
|
||||
};
|
||||
}
|
||||
@@ -13,6 +13,10 @@ import {
|
||||
PROVIDER_ID_TO_ALIAS,
|
||||
} from "@omniroute/open-sse/config/providerModels.ts";
|
||||
import { handleChatCore } from "@omniroute/open-sse/handlers/chatCore.ts";
|
||||
import {
|
||||
checkResourcePressureGuard,
|
||||
type ResourcePressureGuardResult,
|
||||
} from "@omniroute/open-sse/utils/resourcePressure.ts";
|
||||
import {
|
||||
errorResponse,
|
||||
modelCooldownResponse,
|
||||
@@ -64,6 +68,10 @@ type ExecuteChatWithBreakerOptions = {
|
||||
[key: string]: any;
|
||||
};
|
||||
|
||||
type ExecuteChatWithBreakerResult =
|
||||
| { result: any; tlsFingerprintUsed: boolean }
|
||||
| { localResourcePressureResult: ResourcePressureGuardResult; tlsFingerprintUsed: false };
|
||||
|
||||
function getHeaderValue(headers: Record<string, unknown> | null | undefined, name: string) {
|
||||
if (!headers || typeof headers !== "object") return "";
|
||||
const lowerName = name.toLowerCase();
|
||||
@@ -368,6 +376,14 @@ export async function checkPipelineGates(
|
||||
return null;
|
||||
}
|
||||
|
||||
export function checkResourcePressureBeforeProviderWork(): ResourcePressureGuardResult | null {
|
||||
try {
|
||||
return checkResourcePressureGuard();
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
export async function executeChatWithBreaker({
|
||||
bypassCircuitBreaker,
|
||||
breaker,
|
||||
@@ -396,7 +412,7 @@ export async function executeChatWithBreaker({
|
||||
correlationId = null,
|
||||
modelPinned = false,
|
||||
routingComboId = null,
|
||||
}: ExecuteChatWithBreakerOptions): Promise<{ result: any; tlsFingerprintUsed: boolean }> {
|
||||
}: ExecuteChatWithBreakerOptions): Promise<ExecuteChatWithBreakerResult> {
|
||||
let tlsFingerprintUsed = false;
|
||||
const normalizedTrafficType: TrafficType =
|
||||
typeof trafficType === "string" && trafficType.trim().toLowerCase() === "shadow"
|
||||
@@ -410,6 +426,11 @@ export async function executeChatWithBreaker({
|
||||
const capture = <T>(fn: () => T): T =>
|
||||
appliedProxySink ? runWithAppliedProxyCapture(appliedProxySink, fn) : fn();
|
||||
|
||||
const pressureGuard = checkResourcePressureBeforeProviderWork();
|
||||
if (pressureGuard) {
|
||||
return { localResourcePressureResult: pressureGuard, tlsFingerprintUsed: false };
|
||||
}
|
||||
|
||||
try {
|
||||
const chatFn = () =>
|
||||
capture(() =>
|
||||
@@ -434,6 +455,7 @@ export async function executeChatWithBreaker({
|
||||
correlationId,
|
||||
modelPinned,
|
||||
routingComboId,
|
||||
skipResourcePressureGuard: true,
|
||||
onCredentialsRefreshed: async (newCreds: any) => {
|
||||
await updateProviderCredentials(credentials.connectionId, {
|
||||
accessToken: newCreds.accessToken,
|
||||
|
||||
Reference in New Issue
Block a user