From 8ca40e7971af08443e9a63aba0fc544cd5b60c59 Mon Sep 17 00:00:00 2001 From: Xiangzhe Date: Mon, 3 Aug 2026 07:58:45 +0800 Subject: [PATCH] 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. --- open-sse/utils/ollamaTransform.ts | 5 + src/app/api/monitoring/health/route.ts | 12 + src/app/api/v1/completions/route.ts | 8 +- .../[provider]/chat/completions/route.ts | 3 +- src/app/api/v1beta/models/[...path]/route.ts | 7 +- src/lib/monitoring/observability.ts | 62 +++ src/sse/handlers/chat.ts | 49 +- src/sse/handlers/chatAdmission.ts | 239 +++++++++ src/sse/handlers/chatHelpers.ts | 24 +- .../monitoring-health-cache.test.ts | 9 +- .../adaptive-admission-route-matrix.test.ts | 454 ++++++++++++++++ .../chat-adaptive-admission-binding.test.ts | 443 ++++++++++++++++ tests/unit/chat-admission-wrapper.test.ts | 483 ++++++++++++++++++ ...ute-chat-resource-pressure-breaker.test.ts | 246 +++++++++ tests/unit/observability-payloads.test.ts | 110 ++++ tests/unit/ollama-transform.test.ts | 44 ++ 16 files changed, 2163 insertions(+), 35 deletions(-) create mode 100644 src/sse/handlers/chatAdmission.ts create mode 100644 tests/unit/adaptive-admission-route-matrix.test.ts create mode 100644 tests/unit/chat-adaptive-admission-binding.test.ts create mode 100644 tests/unit/chat-admission-wrapper.test.ts create mode 100644 tests/unit/execute-chat-resource-pressure-breaker.test.ts diff --git a/open-sse/utils/ollamaTransform.ts b/open-sse/utils/ollamaTransform.ts index 62fc36ec29..12844c87e5 100644 --- a/open-sse/utils/ollamaTransform.ts +++ b/open-sse/utils/ollamaTransform.ts @@ -11,6 +11,11 @@ type PendingToolCall = { // Transform OpenAI SSE stream to Ollama JSON lines format export function transformToOllama(response, model) { + // Only successful SSE responses belong to the NDJSON transformer. Preserve errors, + // bodyless responses, and successful JSON responses without losing status/body/headers. + const contentType = String(response.headers?.get?.("content-type") || "").toLowerCase(); + if (!response.ok || !response.body || !contentType.includes("text/event-stream")) return response; + let buffer = ""; let pendingToolCalls: Record = {}; const completedToolCalls: PendingToolCall[] = []; diff --git a/src/app/api/monitoring/health/route.ts b/src/app/api/monitoring/health/route.ts index 18d3f10de8..8d7f5f9b2c 100644 --- a/src/app/api/monitoring/health/route.ts +++ b/src/app/api/monitoring/health/route.ts @@ -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 }, }); } diff --git a/src/app/api/v1/completions/route.ts b/src/app/api/v1/completions/route.ts index d6a2254b9c..f11f7cb904 100644 --- a/src/app/api/v1/completions/route.ts +++ b/src/app/api/v1/completions/route.ts @@ -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; } diff --git a/src/app/api/v1/providers/[provider]/chat/completions/route.ts b/src/app/api/v1/providers/[provider]/chat/completions/route.ts index 39f9c5977f..f064163b9e 100644 --- a/src/app/api/v1/providers/[provider]/chat/completions/route.ts +++ b/src/app/api/v1/providers/[provider]/chat/completions/route.ts @@ -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)); } diff --git a/src/app/api/v1beta/models/[...path]/route.ts b/src/app/api/v1beta/models/[...path]/route.ts index 79d56ae614..465ecdf78a 100644 --- a/src/app/api/v1beta/models/[...path]/route.ts +++ b/src/app/api/v1beta/models/[...path]/route.ts @@ -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 diff --git a/src/lib/monitoring/observability.ts b/src/lib/monitoring/observability.ts index a7036f7480..9974c48e55 100644 --- a/src/lib/monitoring/observability.ts +++ b/src/lib/monitoring/observability.ts @@ -1,5 +1,63 @@ +import type { AdaptiveAdmissionPublicSnapshot } from "@omniroute/open-sse/services/admission/runtime.ts"; + type JsonRecord = Record; +/** 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, }, diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 510314f455..8a03241b6c 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -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; diff --git a/src/sse/handlers/chatAdmission.ts b/src/sse/handlers/chatAdmission.ts new file mode 100644 index 0000000000..78387e2742 --- /dev/null +++ b/src/sse/handlers/chatAdmission.ts @@ -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; +}; + +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)[key]; + return typeof value === "string" ? value : ""; +} + +function asStatus(err: object): number | null { + const status = (err as Record).status; + if (typeof status === "number") return status; + const statusCode = (err as Record).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): 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, 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(build: (clientBody: unknown) => T): T; +} { + const target = + body !== null && typeof body === "object" ? (body as Record) : 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; + +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 { + 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; + } + }; +} diff --git a/src/sse/handlers/chatHelpers.ts b/src/sse/handlers/chatHelpers.ts index 7faf961975..272735940d 100644 --- a/src/sse/handlers/chatHelpers.ts +++ b/src/sse/handlers/chatHelpers.ts @@ -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 | 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 { let tlsFingerprintUsed = false; const normalizedTrafficType: TrafficType = typeof trafficType === "string" && trafficType.trim().toLowerCase() === "shadow" @@ -410,6 +426,11 @@ export async function executeChatWithBreaker({ const capture = (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, diff --git a/tests/integration/monitoring-health-cache.test.ts b/tests/integration/monitoring-health-cache.test.ts index 61a918609b..3c7a66b397 100644 --- a/tests/integration/monitoring-health-cache.test.ts +++ b/tests/integration/monitoring-health-cache.test.ts @@ -24,8 +24,13 @@ const { GET, DELETE } = await import("../../src/app/api/monitoring/health/route. async function healthTimestamp(): Promise { const res = await GET(); - const body = (await res.json()) as { timestamp?: string }; + const body = (await res.json()) as { + timestamp?: string; + adaptiveAdmission?: unknown; + }; assert.ok(body.timestamp, "health payload should carry a timestamp"); + // Adaptive admission is always projected (summary object or null) — never omitted. + assert.ok("adaptiveAdmission" in body, "health payload must include adaptiveAdmission"); return body.timestamp as string; } @@ -54,7 +59,7 @@ test("DELETE (circuit-breaker reset) invalidates the cache immediately", async ( new Request("http://localhost/api/monitoring/health", { method: "DELETE", headers: { cookie: `auth_token=${authToken}` }, - }), + }) ); assert.ok(delRes.status < 400, `DELETE should succeed, got ${delRes.status}`); await new Promise((r) => setTimeout(r, 5)); // ensure the clock advances past ms precision diff --git a/tests/unit/adaptive-admission-route-matrix.test.ts b/tests/unit/adaptive-admission-route-matrix.test.ts new file mode 100644 index 0000000000..cd1677bf4f --- /dev/null +++ b/tests/unit/adaptive-admission-route-matrix.test.ts @@ -0,0 +1,454 @@ +/** + * Behavioral matrix: adaptive-admission enforce rejection across the 10 real + * shared LLM POST route modules. Uses a test-owned POST table (not production + * registries/globs). Asserts standardized 503 contract, zero provider fetch, + * provider-health isolation, and runtime reject accounting. + * + * DB isolation: only Node/assert + harness are static imports; createChatPipelineHarness + * must run before any dynamic runtime/resource/DB/route import so DATA_DIR is set first. + */ +import test from "node:test"; +import assert from "node:assert/strict"; +import { createHash } from "node:crypto"; +import { createChatPipelineHarness } from "../integration/_chatPipelineHarness.ts"; + +const harness = await createChatPipelineHarness("adaptive-admission-route-matrix"); +assert.ok( + harness.TEST_DATA_DIR.includes("adaptive-admission-route-matrix") || + harness.TEST_DATA_DIR.includes("omniroute-"), + "task-private harness DATA_DIR must be set before DB imports" +); +console.log(`[adaptive-admission-route-matrix] DATA_DIR=${harness.TEST_DATA_DIR}`); + +const { BaseExecutor, resetStorage, seedConnection, cleanup } = harness; + +const { + getAdaptiveAdmissionRuntime, + reloadAdaptiveAdmissionRuntime, + resetAdaptiveAdmissionRuntimeForTests, +} = await import("../../open-sse/services/admission/runtime.ts"); +const { reloadResourcePressureRuntime } = await import("../../open-sse/utils/resourcePressure.ts"); +const { getProviderConnectionById } = await import("../../src/lib/db/providers.ts"); +const { getCircuitBreaker, resetAllCircuitBreakers, STATE } = + await import("../../src/shared/utils/circuitBreaker.ts"); +const core = await import("../../src/lib/db/core.ts"); +const relayProxies = await import("../../src/lib/db/relayProxies.ts"); + +const chatCompletionsRoute = await import("../../src/app/api/v1/chat/completions/route.ts"); +const messagesRoute = await import("../../src/app/api/v1/messages/route.ts"); +const responsesRoute = await import("../../src/app/api/v1/responses/route.ts"); +const responsesCatchAllRoute = await import("../../src/app/api/v1/responses/[...path]/route.ts"); +const completionsRoute = await import("../../src/app/api/v1/completions/route.ts"); +const ollamaRoute = await import("../../src/app/api/v1/api/chat/route.ts"); +const antigravityRoute = await import("../../src/app/api/v1/antigravity/route.ts"); +const providerPinnedRoute = + await import("../../src/app/api/v1/providers/[provider]/chat/completions/route.ts"); +const relayRoute = await import("../../src/app/api/v1/relay/chat/completions/route.ts"); +const geminiRoute = await import("../../src/app/api/v1beta/models/[...path]/route.ts"); + +const originalFetch = globalThis.fetch; +const MiB = 1024 ** 2; +const MODEL = "openai/gpt-4o-mini"; +const ADMISSION_MESSAGE = "Request too large for current capacity"; + +type RouteCase = { + name: string; + invoke: (request: Request) => Promise; + buildRequest: () => Request; +}; + +function padContent(label: string, targetBytes = 2048): string { + const base = `${label}-admission-matrix-`; + return base + "y".repeat(Math.max(0, targetBytes - base.length)); +} + +function jsonRequest( + url: string, + body: unknown, + headers: Record = {}, + signal?: AbortSignal +): Request { + return new Request(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + Accept: "application/json", + ...headers, + }, + body: JSON.stringify(body), + signal, + }); +} + +function chatBody(label: string) { + return { + model: MODEL, + stream: false, + messages: [{ role: "user", content: padContent(label) }], + }; +} + +function messagesBody(label: string) { + return { + model: MODEL, + max_tokens: 64, + stream: false, + messages: [{ role: "user", content: padContent(label) }], + }; +} + +function responsesBody(label: string) { + return { + model: MODEL, + stream: false, + input: [{ role: "user", content: padContent(label) }], + }; +} + +function completionsBody(label: string) { + return { + model: MODEL, + stream: false, + prompt: padContent(label, 3072), + }; +} + +function antigravityBody(label: string) { + return { + model: MODEL, + project: "admission-matrix-project", + request: { + contents: [{ role: "user", parts: [{ text: padContent(label) }] }], + }, + }; +} + +function geminiBody(label: string) { + return { + contents: [{ role: "user", parts: [{ text: padContent(label) }] }], + }; +} + +function insertRelayToken(rawToken: string) { + const db = core.getDbInstance(); + const id = "rl_admission_matrix"; + const now = Math.floor(Date.now() / 1000); + const tokenHash = createHash("sha256").update(rawToken).digest("hex"); + db.prepare( + ` + INSERT INTO relay_tokens (id, name, token_hash, token_prefix, description, combo_id, allowed_models, + max_tokens_per_request, max_requests_per_minute, max_requests_per_day, max_cost_per_day, + enabled, created_at, updated_at, expires_at, metadata) + VALUES (?, ?, ?, ?, '', NULL, '["*"]', 128000, 1000, 100000, 0, 1, ?, ?, NULL, '{}') + ` + ).run(id, "admission-matrix-relay", tokenHash, "rl_matrix", now, now); + const token = relayProxies.getRelayToken(id); + if (!token) throw new Error("failed to insert matrix relay token"); + return { token, rawToken }; +} + +function reloadNormalResourcePressure() { + reloadResourcePressureRuntime({ + heapThresholdMb: 10_000, + immediateHeapUsedMb: () => 1, + sample: async () => ({ + observedAtMs: Date.now(), + v8: { heapUsedBytes: MiB, heapLimitBytes: 10_000 * MiB }, + process: { + rssBytes: MiB, + externalBytes: 0, + arrayBuffersBytes: 0, + availableBytes: null, + constrainedBytes: null, + }, + cgroup: { currentBytes: null, maxBytes: null, highBytes: null, events: null }, + psi: null, + }), + }); +} + +function reloadEnforceOversized() { + reloadAdaptiveAdmissionRuntime({ + config: { + mode: "enforce", + minLimit: 1, + initialLimit: 1, + maxLimit: 1, + maxQueueCount: 1, + maxQueueCost: 1, + defaultMaxWaitMs: 50, + windowMs: 50, + cost: { + maxRequestCost: 100, + baseCost: 1, + bodyBytesPerUnit: 1, + tokensPerUnit: 1, + messagesPerUnit: 1, + toolsPerUnit: 1, + fanoutPerUnit: 1, + streamingClassCost: 1, + nonStreamingClassCost: 1, + }, + }, + checkResourcePressure: () => null, + }); +} + +function connectionFailureState(connection: Record | null) { + assert.ok(connection); + return { + isActive: connection.isActive, + testStatus: connection.testStatus, + rateLimitedUntil: connection.rateLimitedUntil ?? null, + backoffLevel: connection.backoffLevel ?? null, + lastError: connection.lastError ?? null, + lastErrorAt: connection.lastErrorAt ?? null, + lastErrorType: connection.lastErrorType ?? null, + lastErrorSource: connection.lastErrorSource ?? null, + errorCode: connection.errorCode ?? null, + }; +} + +function breakerSnapshot(breaker: ReturnType) { + const status = breaker.getStatus(); + return { + state: status.state, + failureCount: status.failureCount, + successCount: breaker.successCount, + }; +} + +async function assertAdmissionOversized(response: Response, fetchCalls: number) { + assert.equal(response.status, 503); + assert.equal(fetchCalls, 0); + const contentType = String(response.headers.get("content-type") || ""); + assert.match(contentType, /application\/json/i); + const payload = (await response.json()) as { + error?: { code?: string; type?: string; message?: string }; + }; + assert.equal(payload.error?.type, "server_error"); + assert.equal(payload.error?.code, "admission_oversized"); + assert.equal(payload.error?.message, ADMISSION_MESSAGE); +} + +// Test-owned table of the exact 10 canonical shared LLM POST handlers. +const ROUTE_CASES: RouteCase[] = [ + { + name: "chat.completions", + invoke: (request) => chatCompletionsRoute.POST(request), + buildRequest: () => + jsonRequest("http://localhost/v1/chat/completions", chatBody("chat-completions")), + }, + { + name: "messages", + invoke: (request) => messagesRoute.POST(request, {}), + buildRequest: () => jsonRequest("http://localhost/v1/messages", messagesBody("messages")), + }, + { + name: "responses", + invoke: (request) => responsesRoute.POST(request, {}), + buildRequest: () => + jsonRequest("http://localhost/v1/responses", responsesBody("responses"), { + Accept: "application/json", + }), + }, + { + name: "responses.catch-all", + invoke: (request) => responsesCatchAllRoute.POST(request), + buildRequest: () => + jsonRequest( + "http://localhost/v1/responses/input_items", + responsesBody("responses-catch-all"), + { Accept: "application/json" } + ), + }, + { + name: "completions.legacy", + invoke: (request) => completionsRoute.POST(request), + buildRequest: () => + jsonRequest("http://localhost/v1/completions", completionsBody("legacy-completions")), + }, + { + name: "ollama.api.chat", + invoke: (request) => ollamaRoute.POST(request), + buildRequest: () => jsonRequest("http://localhost/api/chat", chatBody("ollama")), + }, + { + name: "antigravity", + invoke: (request) => antigravityRoute.POST(request), + buildRequest: () => + jsonRequest("http://localhost/v1/antigravity", antigravityBody("antigravity")), + }, + { + name: "providers.pinned", + invoke: (request) => + providerPinnedRoute.POST(request, { params: Promise.resolve({ provider: "openai" }) }), + buildRequest: () => + jsonRequest("http://localhost/v1/providers/openai/chat/completions", { + model: "gpt-4o-mini", + stream: false, + messages: [{ role: "user", content: padContent("provider-pinned") }], + }), + }, + { + name: "relay.chat.completions", + invoke: (request) => relayRoute.POST(request), + buildRequest: () => { + throw new Error("relay buildRequest is set per-test after token insert"); + }, + }, + { + name: "gemini.v1beta.generateContent", + invoke: (request) => + geminiRoute.POST(request, { + params: Promise.resolve({ path: ["openai", "gpt-4o-mini:generateContent"] }), + }), + buildRequest: () => + jsonRequest( + "http://localhost/v1beta/models/openai/gpt-4o-mini:generateContent", + geminiBody("gemini") + ), + }, +]; + +test.beforeEach(async () => { + BaseExecutor.RETRY_CONFIG.delayMs = 0; + await resetStorage(); + resetAllCircuitBreakers(); + resetAdaptiveAdmissionRuntimeForTests(); + reloadNormalResourcePressure(); + reloadEnforceOversized(); + globalThis.fetch = originalFetch; + delete process.env.OMNIROUTE_RELAY_BACKEND; + delete process.env.RELAY_ROUTING_BACKEND; +}); + +test.afterEach(async () => { + globalThis.fetch = originalFetch; + resetAdaptiveAdmissionRuntimeForTests(); + delete process.env.OMNIROUTE_RELAY_BACKEND; + delete process.env.RELAY_ROUTING_BACKEND; + await resetStorage(); +}); + +test.after(async () => { + globalThis.fetch = originalFetch; + resetAdaptiveAdmissionRuntimeForTests(); + await cleanup(); +}); + +test( + "adaptive admission enforce rejects all 10 shared LLM POST routes with standardized contract", + { timeout: 30_000 }, + async () => { + const connection = await seedConnection("openai", { + name: "admission-matrix-openai", + apiKey: "sk-openai-admission-matrix", + }); + const connectionId = String(connection.id); + const beforeConnection = connectionFailureState( + (await getProviderConnectionById(connectionId)) as Record | null + ); + const breaker = getCircuitBreaker("openai"); + const beforeBreaker = breakerSnapshot(breaker); + assert.equal(beforeBreaker.state, STATE.CLOSED); + + const rawRelayToken = `relay_matrix_${createHash("sha256").update("admission").digest("hex").slice(0, 24)}`; + insertRelayToken(rawRelayToken); + process.env.OMNIROUTE_RELAY_BACKEND = "ts"; + + const cases: RouteCase[] = ROUTE_CASES.map((routeCase) => { + if (routeCase.name !== "relay.chat.completions") return routeCase; + return { + ...routeCase, + buildRequest: () => + jsonRequest("http://localhost/api/v1/relay/chat/completions", chatBody("relay"), { + Authorization: `Bearer ${rawRelayToken}`, + }), + }; + }); + + assert.equal(cases.length, 10); + + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return new Response("provider must not run under admission reject", { status: 500 }); + }; + + const beforeRuntime = getAdaptiveAdmissionRuntime().snapshot(); + assert.equal(beforeRuntime.activeCount, 0); + assert.equal(beforeRuntime.queuedCount, 0); + + for (const routeCase of cases) { + const rejectedBefore = getAdaptiveAdmissionRuntime().snapshot().rejectedCount; + const response = await routeCase.invoke(routeCase.buildRequest()); + await assertAdmissionOversized(response, fetchCalls); + + const afterCase = getAdaptiveAdmissionRuntime().snapshot(); + assert.equal( + afterCase.rejectedCount, + rejectedBefore + 1, + `${routeCase.name}: rejectedCount must increment once` + ); + assert.equal(afterCase.activeCount, 0, `${routeCase.name}: activeCount must return to 0`); + assert.equal(afterCase.queuedCount, 0, `${routeCase.name}: queuedCount must return to 0`); + assert.equal(fetchCalls, 0, `${routeCase.name}: fetch must stay 0`); + } + + const afterRuntime = getAdaptiveAdmissionRuntime().snapshot(); + assert.equal(afterRuntime.rejectedCount, beforeRuntime.rejectedCount + cases.length); + assert.equal(afterRuntime.activeCount, 0); + assert.equal(afterRuntime.queuedCount, 0); + assert.equal(fetchCalls, 0); + + assert.deepEqual( + connectionFailureState( + (await getProviderConnectionById(connectionId)) as Record | null + ), + beforeConnection + ); + assert.deepEqual(breakerSnapshot(breaker), beforeBreaker); + } +); + +test( + "provider-pinned route propagates request AbortSignal into admission rejection", + { timeout: 5_000 }, + async () => { + await seedConnection("openai", { + name: "admission-abort-openai", + apiKey: "sk-openai-admission-abort", + }); + reloadEnforceOversized(); + + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return new Response("provider must not run", { status: 500 }); + }; + + const ac = new AbortController(); + ac.abort(); + const response = await providerPinnedRoute.POST( + jsonRequest( + "http://localhost/v1/providers/openai/chat/completions", + { + model: "gpt-4o-mini", + stream: false, + messages: [{ role: "user", content: padContent("abort-provider") }], + }, + {}, + ac.signal + ), + { params: Promise.resolve({ provider: "openai" }) } + ); + + assert.equal(response.status, 499); + assert.equal(fetchCalls, 0); + const payload = (await response.json()) as { error?: { code?: string; type?: string } }; + assert.equal(payload.error?.code, "admission_aborted"); + assert.equal(payload.error?.type, "client_disconnected"); + assert.equal(getAdaptiveAdmissionRuntime().snapshot().activeCount, 0); + } +); diff --git a/tests/unit/chat-adaptive-admission-binding.test.ts b/tests/unit/chat-adaptive-admission-binding.test.ts new file mode 100644 index 0000000000..d2c8e65a15 --- /dev/null +++ b/tests/unit/chat-adaptive-admission-binding.test.ts @@ -0,0 +1,443 @@ +/** + * Shared handleChat ↔ adaptive admission binding tests. + * Proves policy-seam acquire, lazy client-raw, early pre-acquire returns, + * enforce rejection before provider work, and default shadow admission. + */ +import test from "node:test"; +import assert from "node:assert/strict"; +import { createChatPipelineHarness } from "../integration/_chatPipelineHarness.ts"; + +const harness = await createChatPipelineHarness("chat-adaptive-admission-binding"); +const { BaseExecutor, buildRequest, handleChat, resetStorage, seedConnection } = harness; +const { + getAdaptiveAdmissionRuntime, + reloadAdaptiveAdmissionRuntime, + resetAdaptiveAdmissionRuntimeForTests, +} = await import("../../open-sse/services/admission/runtime.ts"); +const { buildClientRawRequest } = await import("../../src/sse/handlers/chat/clientRawRequest.ts"); +const { getProviderConnectionById } = await import("../../src/lib/db/providers.ts"); +const { reloadResourcePressureRuntime } = await import("../../open-sse/utils/resourcePressure.ts"); +const { getCircuitBreaker, resetAllCircuitBreakers, STATE } = + await import("../../src/shared/utils/circuitBreaker.ts"); + +const originalFetch = globalThis.fetch; +const MiB = 1024 ** 2; + +function reloadNormalResourcePressure() { + reloadResourcePressureRuntime({ + heapThresholdMb: 10_000, + immediateHeapUsedMb: () => 1, + sample: async () => ({ + observedAtMs: Date.now(), + v8: { heapUsedBytes: MiB, heapLimitBytes: 10_000 * MiB }, + process: { + rssBytes: MiB, + externalBytes: 0, + arrayBuffersBytes: 0, + availableBytes: null, + constrainedBytes: null, + }, + cgroup: { currentBytes: null, maxBytes: null, highBytes: null, events: null }, + psi: null, + }), + }); +} + +function reloadCriticalResourcePressure() { + reloadResourcePressureRuntime({ + heapThresholdMb: 100, + immediateHeapUsedMb: () => 500, + sample: async () => { + throw new Error("critical request path must not await the async sampler"); + }, + }); +} + +function connectionFailureState(connection: Record | null) { + assert.ok(connection); + return { + isActive: connection.isActive, + testStatus: connection.testStatus, + rateLimitedUntil: connection.rateLimitedUntil ?? null, + backoffLevel: connection.backoffLevel ?? null, + lastError: connection.lastError ?? null, + lastErrorAt: connection.lastErrorAt ?? null, + lastErrorType: connection.lastErrorType ?? null, + lastErrorSource: connection.lastErrorSource ?? null, + errorCode: connection.errorCode ?? null, + }; +} + +test.beforeEach(async () => { + BaseExecutor.RETRY_CONFIG.delayMs = 0; + await resetStorage(); + resetAdaptiveAdmissionRuntimeForTests(); + reloadNormalResourcePressure(); + // Default process runtime is shadow; leave it unless a test reloads enforce. + reloadAdaptiveAdmissionRuntime({ + config: { + mode: "shadow", + minLimit: 8, + initialLimit: 64, + maxLimit: 1000, + maxQueueCount: 128, + maxQueueCost: 2000, + defaultMaxWaitMs: 5_000, + windowMs: 1_000, + }, + checkResourcePressure: () => null, + }); + globalThis.fetch = originalFetch; +}); + +test.afterEach(async () => { + globalThis.fetch = originalFetch; + resetAdaptiveAdmissionRuntimeForTests(); + await resetStorage(); +}); + +test.after(async () => { + globalThis.fetch = originalFetch; + resetAdaptiveAdmissionRuntimeForTests(); + await harness.cleanup(); +}); + +test("invalid body early-return creates no admission lease activity", async () => { + const before = getAdaptiveAdmissionRuntime().snapshot(); + const response = await handleChat( + new Request("http://localhost/v1/chat/completions", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: "{not-json", + }) + ); + assert.equal(response.status, 400); + const after = getAdaptiveAdmissionRuntime().snapshot(); + assert.equal(after.admittedCount, before.admittedCount); + assert.equal(after.activeCount, 0); + assert.equal(after.rejectedCount, before.rejectedCount); +}); + +test("schema-invalid request never acquires an admission lease", async () => { + const before = getAdaptiveAdmissionRuntime().snapshot(); + const response = await handleChat( + buildRequest({ + body: { + model: "openai/gpt-4o-mini", + messages: "not-an-array", + }, + }) + ); + assert.equal(response.status, 400); + const after = getAdaptiveAdmissionRuntime().snapshot(); + assert.equal(after.admittedCount, before.admittedCount); + assert.equal(after.activeCount, 0); +}); + +test("default shadow admits and releases active lease on JSON result", async () => { + await seedConnection("openai", { apiKey: "sk-openai-shadow-admit" }); + + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return new Response( + JSON.stringify({ + id: "chatcmpl-test", + object: "chat.completion", + choices: [ + { index: 0, message: { role: "assistant", content: "ok" }, finish_reason: "stop" }, + ], + }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + }; + + const before = getAdaptiveAdmissionRuntime().snapshot(); + const response = await handleChat( + buildRequest({ + body: { + model: "openai/gpt-4o-mini", + stream: false, + messages: [{ role: "user", content: "hi" }], + }, + }) + ); + assert.equal(response.status, 200); + assert.equal(fetchCalls, 1); + const after = getAdaptiveAdmissionRuntime().snapshot(); + assert.equal(after.activeCount, 0); + assert.equal(after.admittedCount, before.admittedCount + 1); +}); + +test("shared SSE response holds the lease until consumer cancellation", async () => { + await seedConnection("openai", { apiKey: "sk-openai-stream-admit" }); + + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + const stream = new ReadableStream({ + start(controller) { + controller.enqueue( + new TextEncoder().encode( + `data: ${JSON.stringify({ + id: "chatcmpl-stream", + object: "chat.completion.chunk", + choices: [{ index: 0, delta: { content: "hi" }, finish_reason: null }], + })}\n\n` + ) + ); + }, + }); + return new Response(stream, { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }); + }; + + const before = getAdaptiveAdmissionRuntime().snapshot(); + const response = await handleChat( + buildRequest({ + body: { + model: "openai/gpt-4o-mini", + stream: true, + messages: [{ role: "user", content: "stream" }], + }, + }) + ); + + assert.equal(response.status, 200); + assert.equal(fetchCalls, 1); + assert.equal(getAdaptiveAdmissionRuntime().snapshot().activeCount, 1); + assert.equal(getAdaptiveAdmissionRuntime().snapshot().admittedCount, before.admittedCount + 1); + + await response.body!.cancel(); + assert.equal(getAdaptiveAdmissionRuntime().snapshot().activeCount, 0); +}); + +function reloadEnforceOversized() { + // cost >> limit forces immediate ADMISSION_OVERSIZED (not clamped-to-limit admit). + reloadAdaptiveAdmissionRuntime({ + config: { + mode: "enforce", + minLimit: 1, + initialLimit: 1, + maxLimit: 1, + maxQueueCount: 1, + maxQueueCost: 1, + defaultMaxWaitMs: 50, + windowMs: 50, + cost: { + maxRequestCost: 100, + baseCost: 1, + bodyBytesPerUnit: 1, + tokensPerUnit: 1, + messagesPerUnit: 1, + toolsPerUnit: 1, + fanoutPerUnit: 1, + streamingClassCost: 1, + nonStreamingClassCost: 1, + }, + }, + checkResourcePressure: () => null, + }); +} + +function oversizedBody(prefix: string) { + return { + model: "openai/gpt-4o-mini", + stream: false, + messages: Array.from({ length: 20 }, (_, i) => ({ + role: "user", + content: `${prefix}-${i}-${"x".repeat(64)}`, + })), + }; +} + +test("enforce oversized/queue rejection returns standardized 503 before provider fetch", async () => { + await seedConnection("openai", { apiKey: "sk-openai-enforce-reject" }); + reloadEnforceOversized(); + + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return new Response("should-not-run", { status: 200 }); + }; + + const response = await handleChat(buildRequest({ body: oversizedBody("message") })); + + assert.equal(response.status, 503); + const payload = await response.json(); + assert.match(String(payload.error?.code || ""), /^admission_/); + assert.equal(fetchCalls, 0); + assert.equal(getAdaptiveAdmissionRuntime().snapshot().activeCount, 0); +}); + +test("lazy client-raw factory is not invoked on admission rejection", async () => { + reloadEnforceOversized(); + + let factoryCalls = 0; + const body = oversizedBody("lazy"); + const request = buildRequest({ body }); + + const response = await handleChat(request, () => { + factoryCalls += 1; + return buildClientRawRequest(request, body); + }); + + assert.equal(response.status, 503); + assert.equal(factoryCalls, 0); +}); + +test("lazy client-raw factory is invoked exactly once after admission", async () => { + await seedConnection("openai", { apiKey: "sk-openai-lazy-raw" }); + reloadAdaptiveAdmissionRuntime({ + config: { + mode: "shadow", + minLimit: 8, + initialLimit: 64, + maxLimit: 1000, + maxQueueCount: 128, + maxQueueCost: 2000, + defaultMaxWaitMs: 5_000, + windowMs: 1_000, + }, + checkResourcePressure: () => null, + }); + + let factoryCalls = 0; + globalThis.fetch = async () => + new Response( + JSON.stringify({ + id: "chatcmpl-lazy", + object: "chat.completion", + choices: [ + { index: 0, message: { role: "assistant", content: "ok" }, finish_reason: "stop" }, + ], + }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + + const body = { + model: "openai/gpt-4o-mini", + stream: false, + messages: [{ role: "user", content: "lazy once" }], + }; + const request = buildRequest({ body }); + await handleChat(request, () => { + factoryCalls += 1; + return buildClientRawRequest(request, body); + }); + + assert.equal(factoryCalls, 1); + assert.equal(getAdaptiveAdmissionRuntime().snapshot().activeCount, 0); +}); + +test( + "execution-time resource pressure bypasses provider/account accounting", + { timeout: 2_000 }, + async () => { + const connection = await seedConnection("openai", { + name: "pressure-isolation", + apiKey: "sk-openai-pressure-isolation", + }); + const connectionId = String(connection.id); + const beforeConnection = connectionFailureState( + (await getProviderConnectionById(connectionId)) as Record | null + ); + const breaker = getCircuitBreaker("openai"); + const beforeBreaker = breaker.getStatus(); + const beforeSuccessCount = breaker.successCount; + + reloadCriticalResourcePressure(); + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return new Response("provider must not run", { status: 500 }); + }; + + const response = await handleChat( + buildRequest({ + body: { + model: "openai/gpt-4o-mini", + messages: [{ role: "user", content: "shed locally" }], + }, + }) + ); + + assert.equal(response.status, 503); + assert.equal(response.headers.get("Retry-After"), "5"); + const payload = await response.json(); + assert.equal(payload.error.code, "resource_pressure"); + assert.equal(fetchCalls, 0); + assert.deepEqual( + connectionFailureState( + (await getProviderConnectionById(connectionId)) as Record | null + ), + beforeConnection + ); + const afterBreaker = breaker.getStatus(); + assert.equal(afterBreaker.state, beforeBreaker.state); + assert.equal(afterBreaker.failureCount, beforeBreaker.failureCount); + assert.equal(breaker.successCount, beforeSuccessCount); + } +); + +test("resource pressure takes precedence over an open provider breaker", async () => { + const breaker = getCircuitBreaker("openai"); + for (let i = 0; i < 20 && breaker.getStatus().state !== STATE.OPEN; i += 1) { + breaker._onFailure(); + } + const before = breaker.getStatus(); + const beforeSuccessCount = breaker.successCount; + assert.equal(before.state, STATE.OPEN); + + reloadCriticalResourcePressure(); + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return new Response("provider must not run", { status: 500 }); + }; + + const response = await handleChat( + buildRequest({ + body: { + model: "openai/gpt-4o-mini", + messages: [{ role: "user", content: "pressure before breaker" }], + }, + }) + ); + assert.equal(response.status, 503); + assert.equal((await response.json()).error.code, "resource_pressure"); + assert.equal(fetchCalls, 0); + + const after = breaker.getStatus(); + assert.equal(after.state, STATE.OPEN); + assert.equal(after.failureCount, before.failureCount); + assert.equal(breaker.successCount, beforeSuccessCount); +}); + +test("local admission rejection does not mutate a supplied provider breaker", async () => { + resetAllCircuitBreakers(); + const breaker = getCircuitBreaker("openai"); + const before = breaker.getStatus(); + const beforeSuccessCount = breaker.successCount; + assert.equal(before.state, STATE.CLOSED); + assert.equal(before.failureCount, 0); + + reloadEnforceOversized(); + + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return new Response("nope", { status: 200 }); + }; + + const response = await handleChat(buildRequest({ body: oversizedBody("breaker") })); + assert.equal(response.status, 503); + assert.equal(fetchCalls, 0); + + const after = breaker.getStatus(); + assert.equal(after.state, STATE.CLOSED); + assert.equal(after.failureCount, before.failureCount); + assert.equal(breaker.successCount, beforeSuccessCount); +}); diff --git a/tests/unit/chat-admission-wrapper.test.ts b/tests/unit/chat-admission-wrapper.test.ts new file mode 100644 index 0000000000..f3dd910e4f --- /dev/null +++ b/tests/unit/chat-admission-wrapper.test.ts @@ -0,0 +1,483 @@ +/** + * Focused unit tests for the shared handleChat adaptive-admission lifecycle wrapper. + * No provider/network work — pure wrapper + context seams. + */ +import { describe, it, beforeEach, afterEach } from "node:test"; +import assert from "node:assert/strict"; +import { + ANONYMOUS_ADMISSION_TENANT_KEY, + captureDeferredClientRawBody, + classifyHandlerFailure, + createChatAdmissionContext, + resolveAdmissionTenantKey, + withChatAdmission, + type ChatAdmissionContext, +} from "../../src/sse/handlers/chatAdmission.ts"; +import { + createAdaptiveAdmissionRuntime, + type AdaptiveAdmissionRuntime, +} from "../../open-sse/services/admission/runtime.ts"; +import type { AdaptiveAdmissionConfig } from "../../open-sse/services/admission/types.ts"; + +class FakeClock { + nowMs = 0; + private nextId = 1; + private timers = new Map void }>(); + + now = () => this.nowMs; + + setTimer = (fn: () => void, delayMs: number): number => { + const id = this.nextId++; + this.timers.set(id, { due: this.nowMs + Math.max(0, delayMs), fn }); + return id; + }; + + clearTimer = (id: number): void => { + this.timers.delete(id); + }; + + advance(ms: number): void { + const target = this.nowMs + ms; + while (true) { + let nextId: number | undefined; + let nextDue = Number.POSITIVE_INFINITY; + for (const [id, t] of this.timers) { + if (t.due <= target && t.due < nextDue) { + nextDue = t.due; + nextId = id; + } + } + if (nextId === undefined) { + this.nowMs = target; + return; + } + const timer = this.timers.get(nextId)!; + this.timers.delete(nextId); + this.nowMs = timer.due; + timer.fn(); + } + } +} + +function enforceConfig(overrides: Partial = {}): AdaptiveAdmissionConfig { + return { + mode: "enforce", + minLimit: 1, + maxLimit: 4, + initialLimit: 1, + maxQueueCount: 1, + maxQueueCost: 4, + defaultMaxWaitMs: 50, + windowMs: 50, + ...overrides, + }; +} + +function makeRuntime( + clock: FakeClock, + config: AdaptiveAdmissionConfig = enforceConfig() +): AdaptiveAdmissionRuntime { + return createAdaptiveAdmissionRuntime({ + config, + clock: { + now: clock.now, + setTimer: clock.setTimer, + clearTimer: clock.clearTimer, + }, + checkResourcePressure: () => null, + getResourcePressureObservation: () => ({ + signals: null, + state: { + severity: "normal", + reason: "none", + elevatedStreak: 0, + recoveryStreak: 0, + lastTransitionAtMs: 0, + observedAtMs: 0, + }, + }), + nowMs: clock.now, + }); +} + +describe("resolveAdmissionTenantKey", () => { + it("uses only opaque api key id; never falls through to empty/raw", () => { + assert.equal(resolveAdmissionTenantKey("key-abc"), "key-abc"); + assert.equal(resolveAdmissionTenantKey(""), ANONYMOUS_ADMISSION_TENANT_KEY); + assert.equal(resolveAdmissionTenantKey(null), ANONYMOUS_ADMISSION_TENANT_KEY); + assert.equal(resolveAdmissionTenantKey(undefined), ANONYMOUS_ADMISSION_TENANT_KEY); + }); +}); + +describe("captureDeferredClientRawBody", () => { + it("captures only fixed mutable fields and restores client-visible values after admission", () => { + let enumerations = 0; + const target: Record = { + model: "no-think/openai/model", + reasoning: { effort: "high" }, + untouched: "value", + }; + const body = new Proxy(target, { + ownKeys() { + enumerations += 1; + return Reflect.ownKeys(target); + }, + }); + + const deferred = captureDeferredClientRawBody(body); + assert.equal(enumerations, 0, "pre-admission capture must not enumerate the body"); + + body.model = "openai/model"; + body.reasoning_effort = "none"; + delete body.reasoning; + + const captured = deferred.withClientBody((clientBody) => ({ + model: (clientBody as Record).model, + reasoning: (clientBody as Record).reasoning, + hasEffort: Object.hasOwn(clientBody as object, "reasoning_effort"), + })); + + assert.deepEqual(captured, { + model: "no-think/openai/model", + reasoning: { effort: "high" }, + hasEffort: false, + }); + assert.equal(body.model, "openai/model", "working body must be restored after snapshot build"); + assert.equal(body.reasoning_effort, "none"); + assert.equal(Object.hasOwn(body, "reasoning"), false); + }); +}); + +describe("classifyHandlerFailure", () => { + it("classifies abort / timeout / 4xx / else correctly", () => { + const aborted = new AbortController(); + aborted.abort(); + assert.equal(classifyHandlerFailure(new Error("x"), aborted.signal), "cancelled"); + + const abortErr = new Error("aborted"); + abortErr.name = "AbortError"; + assert.equal(classifyHandlerFailure(abortErr), "cancelled"); + + const timeoutErr = new Error("timed out"); + timeoutErr.name = "TimeoutError"; + assert.equal(classifyHandlerFailure(timeoutErr), "timeout"); + + assert.equal(classifyHandlerFailure(Object.assign(new Error("t"), { status: 504 })), "timeout"); + assert.equal( + classifyHandlerFailure(Object.assign(new Error("bad"), { status: 400 })), + "local_reject" + ); + assert.equal(classifyHandlerFailure(new Error("upstream boom")), "upstream_error"); + }); +}); + +describe("createChatAdmissionContext", () => { + let clock: FakeClock; + let runtime: AdaptiveAdmissionRuntime; + + beforeEach(() => { + clock = new FakeClock(); + runtime = makeRuntime(clock); + }); + + afterEach(() => { + runtime.dispose(); + }); + + it("does not acquire when never called", async () => { + const ctx = createChatAdmissionContext(() => runtime); + assert.equal(ctx.getAdmittedState(), null); + assert.equal(runtime.snapshot().activeCount, 0); + assert.equal(runtime.snapshot().admittedCount, 0); + }); + + it("acquires once and rejects a second acquire without re-entering runtime", async () => { + // Capacity must clear default feature cost; this case only locks once-semantics. + runtime.dispose(); + runtime = makeRuntime( + clock, + enforceConfig({ + initialLimit: 64, + minLimit: 8, + maxLimit: 100, + maxQueueCount: 8, + maxQueueCost: 200, + }) + ); + const ctx = createChatAdmissionContext(() => runtime); + const first = await ctx.acquire( + "tenant-a", + { signal: undefined }, + { + messages: [{ role: "user", content: "hi" }], + stream: false, + } + ); + assert.equal(first, null); + assert.ok(ctx.getAdmittedState()); + assert.equal(runtime.snapshot().activeCount, 1); + + const second = await ctx.acquire("tenant-b", {}, { messages: [] }); + assert.equal(second, null); + assert.equal(runtime.snapshot().activeCount, 1); + assert.equal(runtime.snapshot().admittedCount, 1); + + ctx.getAdmittedState()!.admitted.lease.release("success"); + }); + + it("returns standardized 503 rejection without holding a lease", async () => { + const tiny = makeRuntime( + clock, + enforceConfig({ + initialLimit: 1, + minLimit: 1, + maxLimit: 1, + maxQueueCount: 1, + maxQueueCost: 1, + cost: { maxRequestCost: 1, baseCost: 1 }, + }) + ); + const holdCtx = createChatAdmissionContext(() => tiny); + assert.equal(await holdCtx.acquire("hold", {}, { messages: [] }), null); + + const rejectCtx = createChatAdmissionContext(() => tiny); + const rejectPromise = rejectCtx.acquire( + "waiter", + {}, + { messages: [{ role: "user", content: "x" }] } + ); + clock.advance(50); + const rejection = await rejectPromise; + assert.ok(rejection); + assert.equal(rejection!.status, 503); + const body = await rejection!.json(); + assert.match(String(body.error?.code || ""), /^admission_/); + assert.equal(rejectCtx.getAdmittedState(), null); + + holdCtx.getAdmittedState()!.admitted.lease.release("success"); + tiny.dispose(); + }); +}); + +describe("withChatAdmission lifecycle", () => { + let clock: FakeClock; + let runtime: AdaptiveAdmissionRuntime; + + beforeEach(() => { + clock = new FakeClock(); + runtime = makeRuntime(clock, enforceConfig({ mode: "shadow", initialLimit: 8, maxLimit: 20 })); + }); + + afterEach(() => { + runtime.dispose(); + }); + + function wrap( + impl: ( + request: unknown, + clientRaw: unknown, + body: unknown, + correlationId: string | undefined, + ctx: ChatAdmissionContext + ) => Promise + ) { + return withChatAdmission(impl as never, { getRuntime: () => runtime }); + } + + it("early return before acquire creates no lease / runtime activity", async () => { + const handle = wrap(async () => new Response(JSON.stringify({ ok: true }), { status: 400 })); + const res = await handle({ signal: undefined }, null, null); + assert.equal(res.status, 400); + assert.equal(runtime.snapshot().activeCount, 0); + assert.equal(runtime.snapshot().admittedCount, 0); + }); + + it("JSON success releases active lease before return", async () => { + const handle = wrap(async (_req, _raw, _body, _id, ctx) => { + const rejection = await ctx.acquire("k1", {}, { messages: [], stream: false }); + assert.equal(rejection, null); + assert.equal(runtime.snapshot().activeCount, 1); + return new Response(JSON.stringify({ choices: [] }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }); + const res = await handle({}, null, null); + assert.equal(res.status, 200); + assert.equal(runtime.snapshot().activeCount, 0); + assert.equal(runtime.snapshot().admittedCount, 1); + }); + + it("SSE keeps lease through open stream and releases once on cancel", async () => { + const handle = wrap(async (_req, _raw, _body, _id, ctx) => { + const rejection = await ctx.acquire("k-sse", {}, { messages: [], stream: true }); + assert.equal(rejection, null); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data: hi\n\n")); + }, + }); + return new Response(stream, { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }); + }); + const res = await handle({}, null, null); + assert.equal(runtime.snapshot().activeCount, 1); + await res.body!.cancel(); + assert.equal(runtime.snapshot().activeCount, 0); + }); + + it("handler throw after acquisition releases once as upstream_error and rethrows", async () => { + const outcomes: string[] = []; + const release = runtime.releaseHandlerFailure; + runtime.releaseHandlerFailure = (lease, outcome, options) => { + outcomes.push(outcome); + release.call(runtime, lease, outcome, options); + }; + const handle = wrap(async (_req, _raw, _body, _id, ctx) => { + assert.equal(await ctx.acquire("k-err", {}, { messages: [] }), null); + throw new Error("provider exploded"); + }); + await assert.rejects(() => handle({}, null, null), /provider exploded/); + assert.deepEqual(outcomes, ["upstream_error"]); + assert.equal(runtime.snapshot().activeCount, 0); + }); + + it("attach failure releases exactly once before rethrow", async () => { + const outcomes: string[] = []; + const release = runtime.releaseHandlerFailure; + runtime.releaseHandlerFailure = (lease, outcome, options) => { + outcomes.push(outcome); + release.call(runtime, lease, outcome, options); + }; + runtime.attachResponseLifecycle = () => { + throw new Error("attach failed"); + }; + + const handle = wrap(async (_req, _raw, _body, _id, ctx) => { + assert.equal(await ctx.acquire("k-attach", {}, { messages: [] }), null); + return new Response("ok"); + }); + + await assert.rejects(() => handle({}, null, null), /attach failed/); + assert.deepEqual(outcomes, ["upstream_error"]); + assert.equal(runtime.snapshot().activeCount, 0); + }); + + it("timeout-classified throw releases as timeout", async () => { + const handle = wrap(async (_req, _raw, _body, _id, ctx) => { + assert.equal(await ctx.acquire("k-to", {}, { messages: [] }), null); + throw Object.assign(new Error("gateway timeout"), { status: 504 }); + }); + await assert.rejects(() => handle({}, null, null), /gateway timeout/); + assert.equal(runtime.snapshot().activeCount, 0); + }); + + it("queue deadline rejection never invokes inner work after rejection", async () => { + const tiny = makeRuntime( + clock, + enforceConfig({ + initialLimit: 1, + minLimit: 1, + maxLimit: 1, + maxQueueCount: 1, + maxQueueCost: 1, + defaultMaxWaitMs: 40, + cost: { maxRequestCost: 1, baseCost: 1 }, + }) + ); + + const holdHandle = withChatAdmission( + async (_req, _raw, _body, _id, ctx) => { + assert.equal(await ctx.acquire("hold", {}, { messages: [] }), null); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode("data: hold\n\n")); + }, + }); + return new Response(stream, { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }); + }, + { getRuntime: () => tiny } + ); + const holdRes = await holdHandle({}, null, null); + assert.equal(tiny.snapshot().activeCount, 1); + + let innerCalls = 0; + const rejectHandle = withChatAdmission( + async (_req, _raw, _body, _id, ctx) => { + const rejection = await ctx.acquire("waiter", {}, { messages: [] }); + if (rejection) return rejection; + innerCalls += 1; + return new Response("inner", { status: 200 }); + }, + { getRuntime: () => tiny } + ); + + const pending = rejectHandle({}, null, null); + clock.advance(40); + const rejected = await pending; + assert.equal(rejected.status, 503); + assert.equal(innerCalls, 0); + + await holdRes.body!.cancel(); + tiny.dispose(); + }); + + it("queued request abort settles without inner provider work", async () => { + const tiny = makeRuntime( + clock, + enforceConfig({ + initialLimit: 1, + minLimit: 1, + maxLimit: 1, + maxQueueCount: 4, + maxQueueCost: 16, + defaultMaxWaitMs: 5_000, + cost: { maxRequestCost: 1, baseCost: 1 }, + }) + ); + + const holdHandle = withChatAdmission( + async (_req, _raw, _body, _id, ctx) => { + assert.equal(await ctx.acquire("hold", {}, { messages: [] }), null); + return new Response( + new ReadableStream({ + start(c) { + c.enqueue(new TextEncoder().encode("data: h\n\n")); + }, + }), + { status: 200, headers: { "Content-Type": "text/event-stream" } } + ); + }, + { getRuntime: () => tiny } + ); + const holdRes = await holdHandle({}, null, null); + + let innerCalls = 0; + const ac = new AbortController(); + const waitHandle = withChatAdmission( + async (req, _raw, _body, _id, ctx) => { + const rejection = await ctx.acquire("waiter", req, { messages: [] }); + if (rejection) return rejection; + innerCalls += 1; + return new Response("inner", { status: 200 }); + }, + { getRuntime: () => tiny } + ); + + const pending = waitHandle({ signal: ac.signal }, null, null); + // Allow queue promise to arm, then abort without wall-clock sleep. + await Promise.resolve(); + ac.abort(); + const rejected = await pending; + assert.ok(rejected.status === 499 || rejected.status === 503); + assert.equal(innerCalls, 0); + + await holdRes.body!.cancel(); + tiny.dispose(); + }); +}); diff --git a/tests/unit/execute-chat-resource-pressure-breaker.test.ts b/tests/unit/execute-chat-resource-pressure-breaker.test.ts new file mode 100644 index 0000000000..4bf1f54fe7 --- /dev/null +++ b/tests/unit/execute-chat-resource-pressure-breaker.test.ts @@ -0,0 +1,246 @@ +/** + * Resource-pressure isolation: executeChatWithBreaker must shed BEFORE the + * provider breaker path and must not call handleChatCore on pressure 503. + * Direct handleChatCore retains default guard protection. + */ +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-pressure-breaker-")); +process.env.DATA_DIR = TEST_DATA_DIR; + +const core = await import("../../src/lib/db/core.ts"); +const { executeChatWithBreaker } = await import("../../src/sse/handlers/chatHelpers.ts"); +const { getCircuitBreaker, resetAllCircuitBreakers, STATE } = + await import("../../src/shared/utils/circuitBreaker.ts"); +const { reloadResourcePressureRuntime, checkResourcePressureGuard } = + await import("../../open-sse/utils/resourcePressure.ts"); +const { handleChatCore } = await import("../../open-sse/handlers/chatCore.ts"); + +const MiB = 1024 ** 2; + +async function resetStorage() { + resetAllCircuitBreakers(); + core.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); + // Restore a non-shedding resource pressure runtime between tests. + reloadResourcePressureRuntime({ + heapThresholdMb: 10_000, + immediateHeapUsedMb: () => 1, + sample: async () => ({ + observedAtMs: Date.now(), + v8: { heapUsedBytes: MiB, heapLimitBytes: 10_000 * MiB }, + process: { + rssBytes: MiB, + externalBytes: 0, + arrayBuffersBytes: 0, + availableBytes: null, + constrainedBytes: null, + }, + cgroup: { currentBytes: null, maxBytes: null, highBytes: null, events: null }, + psi: null, + }), + }); +} + +test.beforeEach(async () => { + await resetStorage(); +}); + +test.after(async () => { + await resetStorage(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); +}); + +test("executeChatWithBreaker returns typed pressure 503 before normal, bypass, and shadow breaker paths", async () => { + reloadResourcePressureRuntime({ + heapThresholdMb: 100, + immediateHeapUsedMb: () => 999, + sample: async () => { + throw new Error("sampler must not run on request path"); + }, + }); + + // Sanity: process singleton sheds. + const direct = checkResourcePressureGuard(); + assert.ok(direct); + assert.equal(direct!.status, 503); + + const breaker = getCircuitBreaker("openai-pressure-iso"); + const before = breaker.getStatus(); + const beforeSuccessCount = breaker.successCount; + assert.equal(before.state, STATE.CLOSED); + assert.equal(before.failureCount, 0); + + // If handleChatCore were entered it would attempt real provider work / DB. + // Use credentials that would fail loudly if chatCore ran deep. + const credentials = { + connectionId: "conn_pressure_iso", + apiKey: "sk-pressure-iso", + providerSpecificData: {}, + }; + + let canExecuteCalls = 0; + const originalCanExecute = breaker.canExecute.bind(breaker); + breaker.canExecute = () => { + canExecuteCalls += 1; + return originalCanExecute(); + }; + + const baseExecution = { + bypassCircuitBreaker: false, + breaker, + body: { model: "openai/gpt-4o-mini", messages: [{ role: "user", content: "x" }] }, + provider: "openai", + model: "gpt-4o-mini", + refreshedCredentials: credentials, + proxyInfo: null, + log: console, + clientRawRequest: { endpoint: "/v1/chat/completions", headers: {}, body: {} }, + credentials, + apiKeyInfo: null, + userAgent: "", + comboName: null, + comboStrategy: null, + isCombo: false, + extendedContext: false, + comboStepId: null, + comboExecutionKey: null, + }; + const run = ( + overrides: { + bypassCircuitBreaker?: boolean; + trafficType?: "production" | "shadow"; + } = {} + ) => executeChatWithBreaker({ ...baseExecution, ...overrides }); + + const executions = await Promise.all([ + run(), + run({ bypassCircuitBreaker: true }), + run({ trafficType: "shadow" }), + ]); + const pressureResponses: Response[] = []; + for (const execution of executions) { + assert.equal(execution.tlsFingerprintUsed, false); + if (!("localResourcePressureResult" in execution)) { + assert.fail("provider execution result escaped the local pressure guard"); + } + assert.equal(execution.localResourcePressureResult.response.status, 503); + pressureResponses.push(execution.localResourcePressureResult.response); + } + const payload = await pressureResponses[0].json(); + assert.equal(payload.error.code, "resource_pressure"); + assert.match(payload.error.message, /resource pressure/i); + + const after = breaker.getStatus(); + assert.equal(canExecuteCalls, 0); + assert.equal(after.state, STATE.CLOSED); + assert.equal(after.failureCount, before.failureCount); + assert.equal(breaker.successCount, beforeSuccessCount); +}); + +test("direct handleChatCore default still applies resource pressure guard", async () => { + reloadResourcePressureRuntime({ + heapThresholdMb: 100, + immediateHeapUsedMb: () => 500, + sample: async () => ({ + observedAtMs: Date.now(), + v8: { heapUsedBytes: 500 * MiB, heapLimitBytes: 1000 * MiB }, + process: { + rssBytes: 200 * MiB, + externalBytes: 0, + arrayBuffersBytes: 0, + availableBytes: null, + constrainedBytes: null, + }, + cgroup: { currentBytes: null, maxBytes: null, highBytes: null, events: null }, + psi: null, + }), + }); + + const result = await ( + handleChatCore as unknown as (opts: Record) => Promise<{ + success?: boolean; + status?: number; + error?: string; + response?: Response; + }> + )({ + body: { model: "openai/gpt-4o-mini", messages: [{ role: "user", content: "hi" }] }, + modelInfo: { provider: "openai", model: "gpt-4o-mini" }, + credentials: { connectionId: "c1", apiKey: "sk-x", providerSpecificData: {} }, + log: console, + clientRawRequest: { endpoint: "/v1/chat/completions", headers: new Map() }, + }); + + assert.equal(result.success, false); + assert.equal(result.status, 503); + assert.ok(result.response); + const payload = await result.response.json(); + assert.equal(payload.error.code, "resource_pressure"); +}); + +test("handleChatCore skipResourcePressureGuard bypasses the inside-core fuse", async () => { + reloadResourcePressureRuntime({ + heapThresholdMb: 100, + immediateHeapUsedMb: () => 500, + sample: async () => ({ + observedAtMs: Date.now(), + v8: { heapUsedBytes: 500 * MiB, heapLimitBytes: 1000 * MiB }, + process: { + rssBytes: 200 * MiB, + externalBytes: 0, + arrayBuffersBytes: 0, + availableBytes: null, + constrainedBytes: null, + }, + cgroup: { currentBytes: null, maxBytes: null, highBytes: null, events: null }, + psi: null, + }), + }); + + const originalFetch = globalThis.fetch; + let fetchCalls = 0; + globalThis.fetch = async () => { + fetchCalls += 1; + return new Response(JSON.stringify({ error: { message: "upstream" } }), { status: 502 }); + }; + + try { + // With skip=true the pressure fuse is not applied; chatCore proceeds and hits fetch. + const result = await ( + handleChatCore as unknown as (opts: Record) => Promise<{ + success?: boolean; + status?: number; + error?: string; + response?: Response; + }> + )({ + body: { model: "openai/gpt-4o-mini", messages: [{ role: "user", content: "hi" }] }, + modelInfo: { provider: "openai", model: "gpt-4o-mini" }, + credentials: { connectionId: "c1", apiKey: "sk-x", providerSpecificData: {} }, + log: console, + clientRawRequest: { endpoint: "/v1/chat/completions", headers: new Map() }, + skipResourcePressureGuard: true, + }); + + assert.ok(fetchCalls > 0, "skip must let chatCore reach provider work"); + // Must NOT be the resource_pressure 503 from the fuse. + if (result?.response) { + try { + const payload = await result.response.clone().json(); + assert.notEqual(payload?.error?.code, "resource_pressure"); + } catch { + // non-JSON is fine — means we left the pressure fuse path + } + } else if (result?.status === 503) { + assert.notEqual(result?.error, "resource_pressure"); + } + } finally { + globalThis.fetch = originalFetch; + } +}); diff --git a/tests/unit/observability-payloads.test.ts b/tests/unit/observability-payloads.test.ts index 656830fcdf..a56d1554c4 100644 --- a/tests/unit/observability-payloads.test.ts +++ b/tests/unit/observability-payloads.test.ts @@ -5,6 +5,7 @@ import { buildHealthPayload, buildSessionsSummary, buildTelemetryPayload, + projectAdaptiveAdmissionSummary, } from "../../src/lib/monitoring/observability.ts"; test("buildSessionsSummary returns sticky counts and ordered top sessions", () => { @@ -162,4 +163,113 @@ test("buildHealthPayload keeps legacy aliases and adds session/quota observabili assert.equal(payload.quotaMonitor.active, 1); assert.equal(payload.quotaMonitor.monitors[0].provider, "codex"); assert.equal(payload.setupComplete, true); + assert.equal(payload.adaptiveAdmission, null); +}); + +test("buildHealthPayload projects allowlisted adaptiveAdmission aggregates only", () => { + const snapshot = { + mode: "enforce", + currentLimit: 4, + minLimit: 1, + maxLimit: 8, + activeCost: 2, + activeCount: 1, + queuedCost: 3, + queuedCount: 1, + virtualActiveCost: 99, + virtualActiveCount: 99, + virtualQueuedCost: 99, + virtualQueuedCount: 99, + admittedCount: 10, + rejectedCount: 2, + wouldAdmitCount: 7, + wouldQueueCount: 1, + wouldRejectCount: 3, + shortLatencyEwma: 12.5, + longLatencyEwma: 40.1, + utilization: 0.42, + pressure: "high", + resourceSeverity: "normal", + resourceReason: "none", + resourceObservedAtMs: 1_700_000_000_000, + pressureGuardRejectCount: 5, + shutdown: false, + // Malicious / high-card sentinels that must never appear in the public payload. + tenantId: "tenant-SECRET-should-not-leak", + apiKey: "sk-live-SHOULD-NOT-LEAK", + model: "openai/gpt-secret-model", + sessionId: "sess-secret", + requestId: "req-secret", + body: { messages: [{ role: "user", content: "PII-body-secret" }] }, + queueItems: [{ tenantKey: "t-secret", cost: 9 }], + resourcePath: "/sys/fs/cgroup/memory.current", + } as unknown as import("../../open-sse/services/admission/runtime.ts").AdaptiveAdmissionPublicSnapshot; + + const payload = buildHealthPayload({ + appVersion: "9.9.9", + settings: { setupComplete: false }, + connections: [], + circuitBreakers: [], + rateLimitStatus: {}, + learnedLimits: {}, + lockouts: {}, + localProviders: {}, + inflightRequests: 0, + quotaMonitorSummary: { + active: 0, + alerting: 0, + exhausted: 0, + errors: 0, + statusCounts: { + starting: 0, + idle: 0, + healthy: 0, + warning: 0, + exhausted: 0, + error: 0, + }, + byProvider: {}, + }, + quotaMonitorMonitors: [], + activeSessions: [], + adaptiveAdmission: snapshot, + }); + + assert.deepEqual(payload.adaptiveAdmission, { + mode: "enforce", + currentLimit: 4, + minLimit: 1, + maxLimit: 8, + activeCost: 2, + activeCount: 1, + queuedCost: 3, + queuedCount: 1, + admittedCount: 10, + rejectedCount: 2, + wouldAdmitCount: 7, + wouldQueueCount: 1, + wouldRejectCount: 3, + utilization: 0.42, + pressure: "high", + resourceSeverity: "normal", + resourceReason: "none", + resourceObservedAtMs: 1_700_000_000_000, + pressureGuardRejectCount: 5, + shutdown: false, + }); + + const json = JSON.stringify(payload); + assert.equal(json.includes("tenant-SECRET"), false); + assert.equal(json.includes("sk-live-SHOULD-NOT-LEAK"), false); + assert.equal(json.includes("gpt-secret-model"), false); + assert.equal(json.includes("PII-body-secret"), false); + assert.equal(json.includes("t-secret"), false); + assert.equal(json.includes("memory.current"), false); + assert.equal(json.includes("queueItems"), false); + assert.equal(json.includes("virtualActiveCost"), false); + assert.equal(json.includes("shortLatencyEwma"), false); + + // Direct projector also null-safe. + assert.equal(projectAdaptiveAdmissionSummary(null), null); + assert.equal(projectAdaptiveAdmissionSummary(undefined), null); }); diff --git a/tests/unit/ollama-transform.test.ts b/tests/unit/ollama-transform.test.ts index bee6f095c2..22cab56898 100644 --- a/tests/unit/ollama-transform.test.ts +++ b/tests/unit/ollama-transform.test.ts @@ -173,6 +173,50 @@ test("transformToOllama prefers reasoning_content without duplicating aliases", ); }); +test("transformToOllama passes through non-ok shared responses without rewriting status or body", async () => { + const errorBody = { + error: { + message: "Request too large for current capacity", + type: "server_error", + code: "admission_oversized", + }, + }; + const upstream = new Response(JSON.stringify(errorBody), { + status: 503, + headers: { + "Content-Type": "application/json", + "Retry-After": "1", + }, + }); + + const result = transformToOllama(upstream, "llama3.2"); + assert.equal(result.status, 503); + assert.equal(result.headers.get("Retry-After"), "1"); + assert.match(String(result.headers.get("Content-Type") || ""), /application\/json/i); + + const payload = await result.json(); + assert.equal(payload.error?.code, "admission_oversized"); + assert.equal(payload.error?.type, "server_error"); + assert.equal(payload.error?.message, "Request too large for current capacity"); +}); + +test("transformToOllama leaves successful non-SSE responses untouched", async () => { + const body = { choices: [{ message: { role: "assistant", content: "hello" } }] }; + const upstream = new Response(JSON.stringify(body), { + status: 200, + headers: { + "Content-Type": "application/json", + "X-Sentinel": "preserved", + }, + }); + + const result = transformToOllama(upstream, "llama3.2"); + assert.equal(result, upstream); + assert.equal(result.status, 200); + assert.equal(result.headers.get("X-Sentinel"), "preserved"); + assert.deepEqual(await result.json(), body); +}); + test("transformToOllama merges multi-chunk numeric tool_call id", async () => { const inputSSE = [ `data: ${JSON.stringify({