diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 8fa9e71f9c..8f8808208f 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -1,6 +1,7 @@ import { injectMemoryAndSkills } from "./chatCore/memorySkillsInjection.ts"; import { resolveChatCoreRequestSetup } from "./chatCore/requestSetup.ts"; import { buildFailureUsageRecord } from "./chatCore/failureUsage.ts"; +import { estimateFinalInputTokens } from "./chatCore/contextEstimation.ts"; import { extractSystemRoleMessages } from "./chatCore/claudeSystemRole.ts"; export { extractSystemRoleMessages } from "./chatCore/claudeSystemRole.ts"; import { checkIdempotencyCache } from "./chatCore/idempotency.ts"; @@ -1822,27 +1823,6 @@ export async function handleChatCore({ // filtering is advisory and may preserve an all-incompatible pool; this is the // hard boundary that prevents a too-large prompt (or a negative token budget) // from reaching an OpenAI-compatible upstream such as NVIDIA NIM. - const estimateFinalInputTokens = (requestBody: Record | null | undefined) => { - const adapted = requestBody - ? adaptBodyForCompression(requestBody as Record).body - : null; - const messages = - adapted?.messages || - requestBody?.contents || - requestBody?.request?.contents || - (Array.isArray(requestBody?.input) - ? requestBody.input - : requestBody?.input && typeof requestBody.input === "object" - ? requestBody.input - : []); - return ( - estimateTokens(messages) + - (Array.isArray(requestBody?.tools) ? estimateTokens(requestBody.tools) : 0) + - estimateTokens(requestBody?.system) + - estimateTokens(requestBody?.instructions) - ); - }; - let finalEstimatedInputTokens = estimateFinalInputTokens(body as Record); // Reuse the already-resolved `contextLimit` (may have been narrowed to the // per-target combo window above, resolveComboContextLimit) instead of a bare diff --git a/open-sse/handlers/chatCore/contextEstimation.ts b/open-sse/handlers/chatCore/contextEstimation.ts new file mode 100644 index 0000000000..f339cffc90 --- /dev/null +++ b/open-sse/handlers/chatCore/contextEstimation.ts @@ -0,0 +1,29 @@ +import { adaptBodyForCompression } from "../../services/compression/bodyAdapter.ts"; +import { estimateTokens } from "../../services/contextManager.ts"; + +type JsonRecord = Record; + +function asJsonRecord(value: unknown): JsonRecord | null { + return value && typeof value === "object" && !Array.isArray(value) ? (value as JsonRecord) : null; +} + +export function estimateFinalInputTokens(requestBody: JsonRecord | null | undefined): number { + const adapted = requestBody ? adaptBodyForCompression(requestBody).body : null; + const nestedRequest = asJsonRecord(requestBody?.request); + const messages = + adapted?.messages || + requestBody?.contents || + nestedRequest?.contents || + (Array.isArray(requestBody?.input) + ? requestBody.input + : requestBody?.input && typeof requestBody.input === "object" + ? requestBody.input + : []); + + return ( + estimateTokens(messages) + + (Array.isArray(requestBody?.tools) ? estimateTokens(requestBody.tools) : 0) + + estimateTokens(requestBody?.system) + + estimateTokens(requestBody?.instructions) + ); +} diff --git a/open-sse/services/contextManager.ts b/open-sse/services/contextManager.ts index c6d51211f7..cb34af4018 100644 --- a/open-sse/services/contextManager.ts +++ b/open-sse/services/contextManager.ts @@ -254,7 +254,7 @@ function extractImageTokens(node: unknown, seen: Set): { node: unknown; * budget instead of measuring its base64 payload as raw text, then the * remainder of the structure is measured normally via the char/4 heuristic. */ -export function estimateTokens(text: string | object | null | undefined): number { +export function estimateTokens(text: unknown): number { if (!text) return 0; if (typeof text === "string") { return Math.ceil(text.length / CHARS_PER_TOKEN); diff --git a/tests/unit/chatcore-context-estimation.test.ts b/tests/unit/chatcore-context-estimation.test.ts new file mode 100644 index 0000000000..65aee71c1e --- /dev/null +++ b/tests/unit/chatcore-context-estimation.test.ts @@ -0,0 +1,40 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { estimateFinalInputTokens } from "../../open-sse/handlers/chatCore/contextEstimation.ts"; +import { estimateTokens } from "../../open-sse/services/contextManager.ts"; + +test("estimateFinalInputTokens counts nested request contents and auxiliary fields", () => { + const contents = [{ role: "user", parts: [{ text: "nested Gemini request" }] }]; + const tools = [{ name: "lookup", description: "find a record" }]; + const body = { + request: { contents }, + tools, + system: "system prompt", + instructions: 42, + }; + + assert.equal( + estimateFinalInputTokens(body), + estimateTokens(contents) + + estimateTokens(tools) + + estimateTokens(body.system) + + estimateTokens(body.instructions) + ); +}); + +test("estimateFinalInputTokens ignores a malformed nested request and uses input", () => { + const input = [{ role: "user", content: "fallback input" }]; + const body = { + request: "not an object", + input, + }; + + assert.equal(estimateFinalInputTokens(body), estimateTokens(input)); +}); + +test("estimateTokens accepts the unknown JSON values it already serializes", () => { + assert.equal(estimateTokens(true), 1); + assert.equal(estimateTokens(42), 1); + assert.equal(estimateTokens(undefined), 0); +});