diff --git a/client/src/hooks/Chat/useTokenUsage.ts b/client/src/hooks/Chat/useTokenUsage.ts index e9f91d68bc..190c477893 100644 --- a/client/src/hooks/Chat/useTokenUsage.ts +++ b/client/src/hooks/Chat/useTokenUsage.ts @@ -67,16 +67,21 @@ export default function useTokenUsage({ } let total = 0; for (const bucket of Object.values(usageTotals.byRate)) { - /** Endpoint-specific config wins (a custom endpoint may price a known - * model name differently); the adapter provider is the fallback, which - * is what resolves agent runs keyed under the underlying provider */ + /** Resolution order, most to least specific: + * 1. the usage event's own endpoint (a custom endpoint pricing a known + * model name differently) + * 2. the agent's resolved backing endpoint — an Agents conversation + * reports endpoint `agents` and adapter provider `openAI`, neither of + * which keys the custom endpoint's tokenConfig + * 3. the adapter provider, as a last resort */ const rates = tokenConfig?.[bucket.endpoint]?.[bucket.model] ?? + (limits.endpoint != null ? tokenConfig?.[limits.endpoint]?.[bucket.model] : undefined) ?? (bucket.provider != null ? tokenConfig?.[bucket.provider]?.[bucket.model] : undefined); total += costFromUnits(bucket, rates); } return total; - }, [usageTotals, tokenConfig]); + }, [usageTotals, tokenConfig, limits.endpoint]); const isSubmittingRef = useRef(isSubmitting); isSubmittingRef.current = isSubmitting; diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 3e8ade6718..893017afad 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -22,7 +22,6 @@ import type { TSubmission, TConversation, EventSubmission, - TTokenUsageEvent, } from 'librechat-data-provider'; import type { EventHandlerParams } from './useEventHandlers'; import type { ActiveJobsResponse } from '~/data-provider'; @@ -656,22 +655,10 @@ export default function useResumableSSE( * via multiset matching; events that raced past the snapshot read * still fold so multi-call runs don't undercount. */ - const backfilledUsage = data.resumeState?.collectedUsage; - backfillUsage(backfilledUsage ?? [], resumeSubmission); - const backfillCounts = new Map(); - for (const entry of backfilledUsage ?? []) { - const key = JSON.stringify(entry); - backfillCounts.set(key, (backfillCounts.get(key) ?? 0) + 1); - } - const foldIfNotBackfilled = (usageData: TTokenUsageEvent) => { - const key = JSON.stringify(usageData); - const count = backfillCounts.get(key) ?? 0; - if (count > 0) { - backfillCounts.set(key, count - 1); - return; - } - usageHandler(usageData, resumeSubmission); - }; + /** Fold the run's persisted usage; events also replayed below are + * deduped by (runId, seq) inside the handler, so prior prompts in + * this conversation keep their usage and gap events still count */ + backfillUsage(data.resumeState?.collectedUsage ?? [], resumeSubmission); if (data.resumeState?.contextUsage) { contextHandler(data.resumeState.contextUsage, resumeSubmission); } @@ -765,8 +752,14 @@ export default function useResumableSSE( if (replayEvent.event === UsageEvents.ON_CONTEXT_USAGE) { contextHandler(replayEvent.data, resumeSubmission); } else if (replayEvent.event === UsageEvents.ON_TOKEN_USAGE) { - foldIfNotBackfilled(replayEvent.data); + usageHandler(replayEvent.data, resumeSubmission); } else if (replayEvent.event != null) { + if ( + replayEvent.event === StepEvents.ON_MESSAGE_DELTA || + replayEvent.event === StepEvents.ON_REASONING_DELTA + ) { + tapStream(replayEvent.data, resumeSubmission); + } stepHandler(replayEvent, resumeSubmission); } } @@ -780,10 +773,19 @@ export default function useResumableSSE( } else if (pendingEvent.event === UsageEvents.ON_CONTEXT_USAGE) { contextHandler(pendingEvent.data, resumeSubmission); } else if (pendingEvent.event === UsageEvents.ON_TOKEN_USAGE) { - foldIfNotBackfilled(pendingEvent.data); + usageHandler(pendingEvent.data, resumeSubmission); } else if (pendingEvent.event != null) { + if ( + pendingEvent.event === StepEvents.ON_MESSAGE_DELTA || + pendingEvent.event === StepEvents.ON_REASONING_DELTA + ) { + tapStream(pendingEvent.data, resumeSubmission); + } stepHandler(pendingEvent, resumeSubmission); } else if (pendingEvent.type != null) { + /** Gap output streamed past the resume snapshot must reach the + * live estimate too, not just the message UI */ + tapContent(pendingEvent.text, resumeSubmission); contentHandler({ data: pendingEvent, submission: resumeSubmission }); } } diff --git a/client/src/hooks/SSE/useUsageHandler.ts b/client/src/hooks/SSE/useUsageHandler.ts index bdace63922..9d5eb9a71b 100644 --- a/client/src/hooks/SSE/useUsageHandler.ts +++ b/client/src/hooks/SSE/useUsageHandler.ts @@ -8,12 +8,13 @@ import type { TContextUsageEvent, } from 'librechat-data-provider'; import { - EMPTY_USAGE_TOTALS, + markUsageFolded, liveTokensFamily, + removeUsageAtoms, calibrationFamily, usageTotalsFamily, branchTotalsFamily, - removeUsageAtoms, + migrateUsageFolded, contextSnapshotFamily, } from '~/store/usage'; import { @@ -46,7 +47,7 @@ export interface UsageHandlers { tapContent: (text: unknown, submission: UsageSubmissionLike) => void; finalizeUsage: (data: FinalDataLike, submission: UsageSubmissionLike) => void; resetLive: (submission: UsageSubmissionLike) => void; - /** Replaces accumulated totals with the run's collected usage on resume */ + /** Idempotently folds the resumed run's collected usage into the totals */ backfillUsage: (entries: TTokenUsageEvent[], submission: UsageSubmissionLike) => void; /** Seeds the live estimate from already-streamed output chars on resume */ seedLive: (chars: number, submission: UsageSubmissionLike) => void; @@ -115,8 +116,19 @@ export default function useUsageHandler(): UsageHandlers { setLive(convoKey, 0); }; - const foldUsage = (data: TTokenUsageEvent, submission: UsageSubmissionLike) => { + /** Folds one usage event into the totals exactly once per conversation. + * Returns false when the event was already counted (live then replayed + * on resume), so callers can skip the live-estimate bump too. */ + const foldUsage = (data: TTokenUsageEvent, submission: UsageSubmissionLike): boolean => { const convoKey = getConvoKey(submission); + /** runId+seq is unique per model call; fall back to the payload when a + * source predates the sequence tag */ + const usageKey = + data.runId != null && data.seq != null ? `${data.runId}:${data.seq}` : JSON.stringify(data); + if (!markUsageFolded(convoKey, usageKey)) { + return false; + } + const endpoint = submission.conversation?.endpoint ?? ''; const model = data.model ?? submission.conversation?.model ?? ''; const units = normalizeUsageUnits(data); @@ -148,14 +160,16 @@ export default function useUsageHandler(): UsageHandlers { : { provider: data.provider, endpoint, model, ...units }, }, }); + return true; }; const usageHandler: UsageHandlers['usageHandler'] = (data, submission) => { - foldUsage(data, submission); + const folded = foldUsage(data, submission); /** Only primary-call usage drives the live context estimate; tagged - * buckets (summarization, subagent) fold into totals/cost only */ - if (data.usage_type != null) { + * buckets (summarization, subagent) fold into totals/cost only. Skip + * the bump for an already-counted event replayed on resume. */ + if (!folded || data.usage_type != null) { return; } confirmedRef.current += data.output_tokens ?? 0; @@ -203,7 +217,9 @@ export default function useUsageHandler(): UsageHandlers { }; const backfillUsage: UsageHandlers['backfillUsage'] = (entries, submission) => { - jotai.set(usageTotalsFamily(getConvoKey(submission)), EMPTY_USAGE_TOTALS); + /** Fold the resumed run's persisted events idempotently — never reset + * the conversation totals, or a reconnect mid-stream would drop the + * usage of prompts already completed earlier in the session */ for (const entry of entries) { foldUsage(entry, submission); } @@ -231,6 +247,7 @@ export default function useUsageHandler(): UsageHandlers { if (realId !== fromKey) { migrateIndex(fromKey, realId); + migrateUsageFolded(fromKey, realId); jotai.set(contextSnapshotFamily(realId), jotai.get(contextSnapshotFamily(fromKey))); jotai.set(usageTotalsFamily(realId), jotai.get(usageTotalsFamily(fromKey))); jotai.set(calibrationFamily(realId), jotai.get(calibrationFamily(fromKey))); diff --git a/client/src/store/usage.ts b/client/src/store/usage.ts index 280c1975da..51a2feeebf 100644 --- a/client/src/store/usage.ts +++ b/client/src/store/usage.ts @@ -60,6 +60,42 @@ export const liveTokensFamily = atomFamily((_conversationId: string) => atom atom(1)); +/** + * Per-conversation set of usage-event identities already folded into the + * totals. Lets the live path and resume backfill fold each model call exactly + * once, so reconnect mid-stream preserves earlier prompts' usage instead of + * resetting the session totals. + */ +const foldedUsageKeys = new Map>(); + +/** Records an event identity; returns false if it was already folded. */ +export function markUsageFolded(conversationId: string, key: string): boolean { + let seen = foldedUsageKeys.get(conversationId); + if (seen == null) { + seen = new Set(); + foldedUsageKeys.set(conversationId, seen); + } + if (seen.has(key)) { + return false; + } + seen.add(key); + return true; +} + +/** Moves folded-key tracking from a temporary convo id to the persisted one. */ +export function migrateUsageFolded(fromId: string, toId: string): void { + const seen = foldedUsageKeys.get(fromId); + if (seen == null) { + return; + } + const target = foldedUsageKeys.get(toId) ?? new Set(); + for (const key of seen) { + target.add(key); + } + foldedUsageKeys.set(toId, target); + foldedUsageKeys.delete(fromId); +} + /** Jotai atomFamily entries are never GC'd — call on conversation switch/cleanup */ export function removeUsageAtoms(conversationId: string): void { branchTotalsFamily.remove(conversationId); @@ -67,4 +103,5 @@ export function removeUsageAtoms(conversationId: string): void { usageTotalsFamily.remove(conversationId); liveTokensFamily.remove(conversationId); calibrationFamily.remove(conversationId); + foldedUsageKeys.delete(conversationId); } diff --git a/packages/data-schemas/src/methods/transaction.spec.ts b/packages/data-schemas/src/methods/transaction.spec.ts index 176e7faef8..badfdaf014 100644 --- a/packages/data-schemas/src/methods/transaction.spec.ts +++ b/packages/data-schemas/src/methods/transaction.spec.ts @@ -1,9 +1,9 @@ import mongoose from 'mongoose'; import { MongoMemoryServer } from 'mongodb-memory-server'; +import type { IBalance } from '..'; import type { ITransaction } from '~/schema/transaction'; import type { TxData } from './transaction'; -import type { IBalance } from '..'; -import { createTxMethods, tokenValues, premiumTokenValues } from './tx'; +import { createTxMethods, tokenValues, premiumTokenValues, defaultRate } from './tx'; import { matchModelName, findMatchingPattern } from './test-helpers'; import { createSpendTokensMethods } from './spendTokens'; import { createTransactionMethods } from './transaction'; @@ -616,6 +616,36 @@ describe('Transactions Config Tests', () => { }); }); +describe('Partial endpointTokenConfig fallback', () => { + const endpointTokenConfig = { + 'custom-model': { prompt: 1.5, completion: 4.5, read: 0.3 }, + }; + + test('uses override rates for a listed model', () => { + expect(getMultiplier({ model: 'custom-model', tokenType: 'prompt', endpointTokenConfig })).toBe( + 1.5, + ); + expect( + getCacheMultiplier({ model: 'custom-model', cacheType: 'read', endpointTokenConfig }), + ).toBe(0.3); + }); + + test('falls back to standard tables for a model absent from the override', () => { + const fallbackPrompt = getMultiplier({ model: 'gpt-4', tokenType: 'prompt' }); + expect(getMultiplier({ model: 'gpt-4', tokenType: 'prompt', endpointTokenConfig })).toBe( + fallbackPrompt, + ); + expect(getMultiplier({ model: 'gpt-4', tokenType: 'prompt', endpointTokenConfig })).not.toBe( + defaultRate, + ); + + const fallbackCacheRead = getCacheMultiplier({ model: 'claude-3-5-sonnet', cacheType: 'read' }); + expect( + getCacheMultiplier({ model: 'claude-3-5-sonnet', cacheType: 'read', endpointTokenConfig }), + ).toBe(fallbackCacheRead); + }); +}); + describe('calculateTokenValue Edge Cases', () => { test('should derive multiplier from model when valueKey is not provided', async () => { const userId = new mongoose.Types.ObjectId(); diff --git a/packages/data-schemas/src/methods/tx.ts b/packages/data-schemas/src/methods/tx.ts index dd05312f0d..3e77007b9f 100644 --- a/packages/data-schemas/src/methods/tx.ts +++ b/packages/data-schemas/src/methods/tx.ts @@ -510,7 +510,13 @@ export function createTxMethods( endpointTokenConfig?: Record>; }): number { if (endpointTokenConfig && model) { - return endpointTokenConfig?.[model]?.[tokenType as string] ?? defaultRate; + const modelConfig = endpointTokenConfig[model]; + /** A partial override only prices the models it lists; others fall + * through to the standard tables so billing matches the advertised + * token config instead of charging defaultRate */ + if (modelConfig) { + return modelConfig[tokenType as string] ?? defaultRate; + } } if (valueKey && tokenType) { @@ -555,7 +561,12 @@ export function createTxMethods( endpointTokenConfig?: Record>; }): number | null { if (endpointTokenConfig && model) { - return endpointTokenConfig?.[model]?.[cacheType as string] ?? null; + const modelConfig = endpointTokenConfig[model]; + /** Models absent from a partial override fall through to standard + * cache rates rather than reporting no cache pricing */ + if (modelConfig) { + return modelConfig[cacheType as string] ?? null; + } } if (valueKey && cacheType) {