From c9926897ba0f816ae0829bcb66a4458a9a53def8 Mon Sep 17 00:00:00 2001 From: joachimBrindeau Date: Thu, 16 Jul 2026 11:18:27 +0700 Subject: [PATCH] feat(rtk): add X-9Router-Token-Saver header to bypass token savers per request --- README.md | 2 ++ open-sse/config/runtimeConfig.js | 2 ++ open-sse/handlers/chatCore.js | 15 +++++---- tests/unit/headroom-chat-core.test.js | 44 +++++++++++++++++++++++++++ 4 files changed, 57 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index a9e169a9..a70294e3 100644 --- a/README.md +++ b/README.md @@ -425,6 +425,8 @@ Default URLs: | 📊 **Usage Analytics** | Track tokens, cost, trends over time | Optimize spending | | 🌐 **Deploy Anywhere** | Localhost, VPS, Docker, Cloudflare Workers | Flexible deployment options | +Set `X-9Router-Token-Saver: off` to bypass all token savers for one chat request. +
📖 Feature Details diff --git a/open-sse/config/runtimeConfig.js b/open-sse/config/runtimeConfig.js index de199233..648281cc 100644 --- a/open-sse/config/runtimeConfig.js +++ b/open-sse/config/runtimeConfig.js @@ -56,6 +56,8 @@ export const GEMINI_NATIVE_TTS_FETCH_TIMEOUT_MS = envMs("GEMINI_NATIVE_TTS_FETCH export const DEFAULT_MAX_TOKENS = 64000; export const DEFAULT_MIN_TOKENS = 32000; +export const TOKEN_SAVER_HEADER = "x-9router-token-saver"; + // Retry config for 429 responses (legacy - kept for backward compatibility) export const RETRY_CONFIG = { maxAttempts: 2, diff --git a/open-sse/handlers/chatCore.js b/open-sse/handlers/chatCore.js index b5cf8a84..be58866e 100644 --- a/open-sse/handlers/chatCore.js +++ b/open-sse/handlers/chatCore.js @@ -9,7 +9,7 @@ import { createRequestLogger } from "../utils/requestLogger.js"; import { getModelTargetFormat, getModelStrip, getModelUpstreamId, getModelType, PROVIDER_ID_TO_ALIAS } from "../config/providerModels.js"; import { PROVIDERS } from "../config/providers.js"; import { createErrorResult, parseUpstreamError, formatProviderError } from "../utils/error.js"; -import { HTTP_STATUS } from "../config/runtimeConfig.js"; +import { HTTP_STATUS, TOKEN_SAVER_HEADER } from "../config/runtimeConfig.js"; import { handleBypassRequest } from "../utils/bypassHandler.js"; import { trackPendingRequest, appendRequestLog, saveRequestDetail } from "@/lib/usageDb.js"; import { getExecutor } from "../executors/index.js"; @@ -156,14 +156,17 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred delete translatedBody.tools; } + // Per-request opt-out: client can bypass all token savers via header + const tokenSaverEnabled = clientRawRequest?.headers?.[TOKEN_SAVER_HEADER]?.toLowerCase() !== "off"; + // RTK: compress tool_result content - const rtkStats = compressMessages(translatedBody, rtkEnabled); + const rtkStats = compressMessages(translatedBody, tokenSaverEnabled && rtkEnabled); const rtkLine = formatRtkLog(rtkStats); if (rtkLine) console.log(rtkLine); // Headroom: optional external proxy compression; fail open if proxy is absent. const headroomDiagnostics = {}; - const headroomStats = await compressWithHeadroom(translatedBody, { enabled: headroomEnabled, url: headroomUrl, model: upstreamModel, format: finalFormat, compressUserMessages: headroomCompressUserMessages, diagnostics: headroomDiagnostics }); + const headroomStats = await compressWithHeadroom(translatedBody, { enabled: tokenSaverEnabled && headroomEnabled, url: headroomUrl, model: upstreamModel, format: finalFormat, compressUserMessages: headroomCompressUserMessages, diagnostics: headroomDiagnostics }); const headroomLine = formatHeadroomLog(headroomStats); const headroomSizeLine = formatHeadroomSizeLog(headroomDiagnostics); if (headroomLine) { @@ -171,16 +174,16 @@ export async function handleChatCore({ body, modelInfo, credentials, log, onCred if (isHeadroomPhantomSavings(headroomStats, headroomDiagnostics)) { log?.warn?.("HEADROOM", `reported token delta, but outbound JSON shrank <5%; provider may bill near-original payload | ${headroomSizeLine}`); } - } else if (headroomEnabled) log?.warn?.("HEADROOM", `skipped: ${headroomDiagnostics.reason || "compression unavailable"}${headroomDiagnostics.endpoint ? ` (${headroomDiagnostics.endpoint})` : ""}`); + } else if (tokenSaverEnabled && headroomEnabled) log?.warn?.("HEADROOM", `skipped: ${headroomDiagnostics.reason || "compression unavailable"}${headroomDiagnostics.endpoint ? ` (${headroomDiagnostics.endpoint})` : ""}`); // Caveman: inject terse-style system prompt - if (cavemanEnabled && cavemanLevel) { + if (tokenSaverEnabled && cavemanEnabled && cavemanLevel) { injectCaveman(translatedBody, finalFormat, cavemanLevel); log?.debug?.("CAVEMAN", `${cavemanLevel} | ${finalFormat}`); } // Ponytail: inject lazy-senior-dev system prompt - if (ponytailEnabled && ponytailLevel) { + if (tokenSaverEnabled && ponytailEnabled && ponytailLevel) { injectPonytail(translatedBody, finalFormat, ponytailLevel); log?.debug?.("PONYTAIL", `${ponytailLevel} | ${finalFormat}`); } diff --git a/tests/unit/headroom-chat-core.test.js b/tests/unit/headroom-chat-core.test.js index 5c552ab9..e61ab898 100644 --- a/tests/unit/headroom-chat-core.test.js +++ b/tests/unit/headroom-chat-core.test.js @@ -250,4 +250,48 @@ describe("handleChatCore Headroom diagnostics", () => { expect.stringContaining("reported token delta, but outbound JSON shrank <5%; provider may bill near-original payload") ); }); + + it("bypasses token savers when requested by the client", async () => { + const log = { debug: vi.fn(), info: vi.fn(), warn: vi.fn() }; + const pxpipeTransform = vi.fn(); + const messages = [{ role: "user", content: "Write polished prose." }]; + + global.fetch = vi.fn(async (url) => { + throw new Error(`unexpected fetch: ${url}`); + }); + + await handleChatCore({ + body: { model: "gpt-4o", stream: false, messages }, + modelInfo: { provider: "openai", model: "gpt-4o" }, + credentials: { apiKey: "test-key", providerSpecificData: {} }, + log, + connectionId: "test-conn", + headroomEnabled: true, + headroomUrl: "http://localhost:8787", + headroomCompressUserMessages: true, + rtkEnabled: true, + cavemanEnabled: true, + cavemanLevel: "full", + ponytailEnabled: true, + ponytailLevel: "full", + pxpipeEnabled: true, + pxpipeTransform, + clientRawRequest: { + endpoint: "/v1/chat/completions", + body: {}, + headers: { + accept: "application/json", + "x-9router-token-saver": "off", + }, + }, + }); + + expect(global.fetch).not.toHaveBeenCalled(); + expect(pxpipeTransform).not.toHaveBeenCalled(); + expect(executeMock).toHaveBeenCalledWith(expect.objectContaining({ + body: expect.objectContaining({ + messages: [{ role: "user", content: "Write polished prose." }], + }), + })); + }); });