diff --git a/open-sse/services/admission/requestFeatures.ts b/open-sse/services/admission/requestFeatures.ts index 0116a1e43b..87fdf07ab0 100644 --- a/open-sse/services/admission/requestFeatures.ts +++ b/open-sse/services/admission/requestFeatures.ts @@ -4,11 +4,13 @@ */ import { estimateSizeFast } from "../../utils/estimateSize.ts"; -import type { AdmissionCostFeatures } from "./types.ts"; +import { resolveCostConfig } from "./cost.ts"; +import type { AdmissionCostConfig, AdmissionCostFeatures } from "./types.ts"; export type AdmissionFeatureExtractionContext = { /** When set, wins over any body/wrapped stream field. */ streaming?: boolean; + cost?: Partial; }; /** @@ -159,7 +161,8 @@ export function extractAdmissionCostFeatures( body: unknown, context?: AdmissionFeatureExtractionContext ): AdmissionCostFeatures { - const bodyBytes = estimateSizeFast(body); + const cost = resolveCostConfig(context?.cost); + const bodyBytes = estimateSizeFast(body, cost.bodyBytesPerUnit * cost.maxRequestCost); const layers = featureLayers(body); const draft: FeatureDraft = { messageCount: 0, diff --git a/open-sse/services/admission/runtime.ts b/open-sse/services/admission/runtime.ts index 919509ce60..fb04c3549b 100644 --- a/open-sse/services/admission/runtime.ts +++ b/open-sse/services/admission/runtime.ts @@ -5,6 +5,7 @@ import { AdaptiveAdmissionController } from "./controller.ts"; import { validateConfig } from "./config.ts"; +import { resolveCostConfig } from "./cost.ts"; import { extractAdmissionCostFeatures } from "./requestFeatures.ts"; import { type AdaptiveAdmissionConfig, @@ -323,6 +324,7 @@ function classifyHttpOutcome(status: number, signal?: AbortSignal): AdmissionRel class AdaptiveAdmissionRuntimeImpl implements AdaptiveAdmissionRuntime { private readonly controller: AdaptiveAdmissionController; + private readonly costConfig: ReturnType; private readonly checkResourcePressure: () => ResourcePressureGuardResult | null; private readonly getResourcePressureObservation: () => ResourcePressureObservation; private readonly onPressureObserved?: (pressure: AdmissionPressure) => void; @@ -338,6 +340,7 @@ class AdaptiveAdmissionRuntimeImpl implements AdaptiveAdmissionRuntime { constructor(options: AdaptiveAdmissionRuntimeOptions, config: AdaptiveAdmissionConfig) { this.controller = new AdaptiveAdmissionController(config, options.clock); + this.costConfig = resolveCostConfig(config.cost); this.checkResourcePressure = options.checkResourcePressure ?? checkResourcePressureGuard; this.getResourcePressureObservation = options.getResourcePressureObservation ?? getResourcePressureObservation; @@ -372,7 +375,10 @@ class AdaptiveAdmissionRuntimeImpl implements AdaptiveAdmissionRuntime { const features = extractAdmissionCostFeatures( input.body, - input.streaming === undefined ? undefined : { streaming: input.streaming } + { + streaming: input.streaming, + cost: this.costConfig, + } ); let result: AdmissionAcquireResult; try { diff --git a/tests/unit/adaptive-admission-features.test.ts b/tests/unit/adaptive-admission-features.test.ts index 30ebaa8ff7..fa004b90b3 100644 --- a/tests/unit/adaptive-admission-features.test.ts +++ b/tests/unit/adaptive-admission-features.test.ts @@ -37,6 +37,25 @@ describe("bounded request feature extraction", () => { } }); + it("measures admission bodies beyond the default size-estimator limit", () => { + const body = { + messages: Array.from({ length: 200 }, () => ({ role: "user", content: "x".repeat(10_000) })), + }; + const features = extractAdmissionCostFeatures(body); + + assert.ok((features.bodyBytes ?? 0) > 1_000_000); + assert.ok((features.estimatedInputTokens ?? 0) > 250_000); + }); + + it("uses the active admission cost budget when measuring bodies", () => { + const body = { payload: "x".repeat(17_000_000) }; + const features = extractAdmissionCostFeatures(body, { + cost: { bodyBytesPerUnit: 1_000_000, maxRequestCost: 20 }, + }); + + assert.ok((features.bodyBytes ?? 0) > 16_384_000); + }); + it("extracts production-realistic Chat, Responses, Gemini, and Antigravity shapes", () => { // OpenAI Chat Completions — stream omitted defaults false (higher non-stream class). const chat = extractAdmissionCostFeatures({ @@ -146,7 +165,7 @@ describe("bounded request feature extraction", () => { it("bounds tool scans and never touches entries beyond the budget (conservative count)", () => { // Huge leading string makes estimateSizeFast byte-exit before walking tools, // so only countTools can touch the tools proxy — proving its scan bound alone. - const sizePad = "x".repeat(300_000); + const sizePad = "x".repeat(17_000_000); let accesses = 0; const tools = new Proxy([] as unknown[], {