diff --git a/open-sse/utils/usageTracking.ts b/open-sse/utils/usageTracking.ts index 38913a2263..b70dcd631c 100644 --- a/open-sse/utils/usageTracking.ts +++ b/open-sse/utils/usageTracking.ts @@ -93,6 +93,16 @@ export function invalidateBufferTokensCache(): void { _cacheTimestamp = 0; } +/** + * Directly set the cached buffer value — called by runtimeSettings after a + * settings save so the new value is available synchronously on the next request + * (no race window between invalidation and the async DB re-read). + */ +export function setBufferTokensCache(value: number): void { + _cachedBuffer = value; + _cacheTimestamp = Date.now(); +} + // Get HH:MM:SS timestamp function getTimeString() { return new Date().toLocaleTimeString("en-US", { diff --git a/src/lib/config/runtimeSettings.ts b/src/lib/config/runtimeSettings.ts index b22f48c294..c8ad61ad17 100644 --- a/src/lib/config/runtimeSettings.ts +++ b/src/lib/config/runtimeSettings.ts @@ -306,10 +306,16 @@ async function applyCacheControlSection() { invalidateCacheControlSettingsCache(); } -async function applyUsageTrackingSection() { - const { invalidateBufferTokensCache } = +async function applyUsageTrackingSection(newBuffer: number | null) { + const { invalidateBufferTokensCache, setBufferTokensCache } = await import("@omniroute/open-sse/utils/usageTracking.ts"); - invalidateBufferTokensCache(); + if (typeof newBuffer === "number" && newBuffer >= 0) { + // Set the value directly so the first request after a settings save gets the + // correct count synchronously — no race window back to DEFAULT (2000). + setBufferTokensCache(newBuffer); + } else { + invalidateBufferTokensCache(); + } } async function applyThoughtSignatureSection(mode: string) { @@ -465,7 +471,9 @@ export async function applyRuntimeSettings( } if (force || hasChanged(currentSnapshot.usageTokenBuffer, previousSnapshot.usageTokenBuffer)) { - await applyUsageTrackingSection(); + const newBuffer = + typeof currentSnapshot.usageTokenBuffer === "number" ? currentSnapshot.usageTokenBuffer : null; + await applyUsageTrackingSection(newBuffer); markChanged("usageTracking"); } diff --git a/tests/unit/usage-token-buffer.test.ts b/tests/unit/usage-token-buffer.test.ts new file mode 100644 index 0000000000..8872b28675 --- /dev/null +++ b/tests/unit/usage-token-buffer.test.ts @@ -0,0 +1,160 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { + addBufferToUsage, + invalidateBufferTokensCache, + setBufferTokensCache, +} from "../../open-sse/utils/usageTracking.ts"; + +// ─── Helpers ────────────────────────────────────────────────────────────── + +function resetEnv(saved: string | undefined) { + if (saved === undefined) { + delete process.env.USAGE_TOKEN_BUFFER; + } else { + process.env.USAGE_TOKEN_BUFFER = saved; + } +} + +// ─── addBufferToUsage — baseline / env-var path ─────────────────────────── + +test("addBufferToUsage — adds DEFAULT 2000 when no env var and cache is null", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + invalidateBufferTokensCache(); + + const result = addBufferToUsage({ prompt_tokens: 25, completion_tokens: 24, total_tokens: 49 }); + + // Documents the race: after invalidation the sync path falls back to DEFAULT=2000 + assert.equal(result.prompt_tokens, 2025); + assert.equal(result.completion_tokens, 24); + assert.equal(result.total_tokens, 2049); + + resetEnv(saved); +}); + +test("addBufferToUsage — respects USAGE_TOKEN_BUFFER=0 env override", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + process.env.USAGE_TOKEN_BUFFER = "0"; + invalidateBufferTokensCache(); + + const result = addBufferToUsage({ prompt_tokens: 25, completion_tokens: 24, total_tokens: 49 }); + + assert.equal(result.prompt_tokens, 25); + assert.equal(result.total_tokens, 49); + + resetEnv(saved); +}); + +test("addBufferToUsage — respects USAGE_TOKEN_BUFFER=500 env override", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + process.env.USAGE_TOKEN_BUFFER = "500"; + invalidateBufferTokensCache(); + + const result = addBufferToUsage({ prompt_tokens: 86, completion_tokens: 52, total_tokens: 138 }); + + assert.equal(result.prompt_tokens, 586); + assert.equal(result.total_tokens, 638); + + resetEnv(saved); +}); + +test("addBufferToUsage — also adds to Claude-format input_tokens", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + process.env.USAGE_TOKEN_BUFFER = "100"; + invalidateBufferTokensCache(); + + const result = addBufferToUsage({ input_tokens: 40, output_tokens: 20 }); + + assert.equal(result.input_tokens, 140); + + resetEnv(saved); +}); + +test("addBufferToUsage — returns usage unchanged when buffer is 0 via env", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + process.env.USAGE_TOKEN_BUFFER = "0"; + invalidateBufferTokensCache(); + + const usage = { prompt_tokens: 86, completion_tokens: 52, total_tokens: 138 }; + const result = addBufferToUsage(usage); + + assert.equal(result.prompt_tokens, 86); + assert.equal(result.completion_tokens, 52); + assert.equal(result.total_tokens, 138); + + resetEnv(saved); +}); + +// ─── setBufferTokensCache — the fix for the race condition ──────────────── +// +// The race: invalidateBufferTokensCache() sets _cachedBuffer=null; the next +// synchronous call to getBufferTokens() falls back to DEFAULT=2000 before +// _loadBufferFromDb() (async) completes. +// +// The fix: runtimeSettings.ts calls setBufferTokensCache(newValue) instead of +// invalidateBufferTokensCache() so the correct value is available synchronously. + +test("setBufferTokensCache(0) — immediately prevents buffer addition (no race window)", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + + // Simulates what runtimeSettings does after saving usageTokenBuffer=0 in DB + setBufferTokensCache(0); + + const result = addBufferToUsage({ prompt_tokens: 25, completion_tokens: 24, total_tokens: 49 }); + + // With the fix: 0 is applied synchronously — no 2000-token race window + assert.equal(result.prompt_tokens, 25); + assert.equal(result.completion_tokens, 24); + assert.equal(result.total_tokens, 49); + + resetEnv(saved); + invalidateBufferTokensCache(); +}); + +test("setBufferTokensCache(500) — immediately sets custom buffer value", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + + setBufferTokensCache(500); + + const result = addBufferToUsage({ prompt_tokens: 86, completion_tokens: 52, total_tokens: 138 }); + + assert.equal(result.prompt_tokens, 586); + assert.equal(result.total_tokens, 638); + + resetEnv(saved); + invalidateBufferTokensCache(); +}); + +test("setBufferTokensCache(0) — works for Claude-format (input_tokens)", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + + setBufferTokensCache(0); + + const result = addBufferToUsage({ input_tokens: 40, output_tokens: 20 }); + + assert.equal(result.input_tokens, 40); + + resetEnv(saved); + invalidateBufferTokensCache(); +}); + +test("invalidateBufferTokensCache — still resets to null (returns DEFAULT on next sync call)", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + delete process.env.USAGE_TOKEN_BUFFER; + + // First prime the cache with a custom value + setBufferTokensCache(0); + const afterSet = addBufferToUsage({ prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }); + assert.equal(afterSet.prompt_tokens, 10); // 0 buffer + + // Then invalidate — next sync call reverts to DEFAULT (2000) while async reload happens + invalidateBufferTokensCache(); + const afterInvalidate = addBufferToUsage({ prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }); + assert.equal(afterInvalidate.prompt_tokens, 2010); // DEFAULT=2000 applied (race window) + + resetEnv(saved); +});