From 376370d61022df87ce65dc53aa83c6accd1b0cee Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 25 Jun 2026 15:29:31 -0400 Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20Compute=20Cont?= =?UTF-8?q?ext=20Gauge=20Client-Side,=20Drop=20Projection=20Endpoint=20(#1?= =?UTF-8?q?3953)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ♻️ refactor: Compute Context Gauge Client-Side, Drop Projection Endpoint The /api/endpoints/context-projection endpoint re-fetched a conversation's messages from Mongo and re-tokenized them to project the context gauge for snapshot-less branches. The browser already holds those messages and their per-message tokenCounts, so this duplicated work on the request path (an unbounded read + server-side BPE tokenization until it was later capped). Move the snapshot-less estimate fully client-side, from the in-memory index: - sumBranch accumulates an uncalibrated char/4 estimate (estTokens) for count-less messages (imports / pre-feature) under the same summary cutoff - useTokenUsage folds estTokens (calibrated via the existing calibrationFamily ratio) into the existing fallback; known per-message counts render unchanged - delete the endpoint, controller, rate limiter, route, the getMessageTextStats data-schemas method, and the data-provider surface (endpoint/key/type/service/query) No DB read, no server tokenization, no rate-limit knobs; the gauge recomputes reactively from the index. Net -793 lines. * 🩹 fix: Count quotes and object-form content in client context estimate Address Codex review on the client-side context estimate: - messageChars now reads object-form content text (part.text.value), not only string text/think, so imported / pre-feature messages whose body lives in content parts are no longer estimated as zero. - Count-less user messages include their merged quote excerpts in the estimate, mirroring what the send path prepends into the prompt. * 🩹 fix: Cap over-window estimate and surface estimated tokens in breakdown Address remaining Codex review on the client-side context estimate: - Clamp the snapshot-less estimate's displayed usedTokens to maxTokens. The send path prunes an over-window branch before calling the model, so the gauge never actually exceeds the window; this avoids impossible values (e.g. 50k / 8k) without re-introducing client-side pruning. - Surface the calibrated count-less estimate as its own "Estimated" row in the breakdown popover, so a branch of only count-less imported / pre-feature messages is no longer shown as Input 0 / Output 0 under a non-zero header. * 🩹 fix: Refine client context estimate per Codex re-review - Drop calibration from the snapshot-less estimate. The removed projection never actually calibrated (the client never sent a ratio), and a ratio inflated by provider-injected context over-estimates visible imported text. - Exclude reasoning (think) / error parts from the estimate; the send path strips them, so they are not part of the next call's context. - Fold quote text into the estimate even when a tokenCount is present, since the edit route recounts tokenCount from text only and drops the merged quote. * 🩹 fix: Recount quoted user turns instead of topping up the stored count The previous round added quote chars on top of a quoted message's stored tokenCount, which double-counts the common (unedited) case where the count already includes the merged quote prompt. Match the removed projection instead: for quoted user turns, ignore the stored count and estimate the full merged text. This both avoids the double-count and still corrects the stale text-only count an edit leaves behind. * 🩹 fix: Trust stored counts for quoted turns; count tool-call parts - Quoted user turns: revert to trusting a present tokenCount. The send path's stored count already includes the merged quote (and any calibration), and the client's char/4 path is coarser, so recounting regressed normal turns. Only count-less messages estimate quotes from text. - Count tool-call name/args/output for count-less assistant messages; the formatter sends them back as context, so omitting them under-reported imported branches with tool history. * 🩹 fix: Exclude in-flight tail from estimate to avoid resume double-count On resume the live path seeds liveTokens from the partial response and also writes that content into the messages cache, where the count-less response is estimated into estTokens too β€” double-counting the in-flight output on the snapshot-less estimate path. sumBranch now exposes the tail message's own estimate (tailEstTokens); the estimate path drops it while a stream is live, so the in-flight response is counted once (via liveTokens). The breakdown's Estimated row uses the same in-flight-adjusted value. * 🩹 fix: Recount quoted user turns in context estimate (match send path) A quoted user turn's stored tokenCount is unreliable for the gauge: a text-only Save edit recomputes it from text alone, and the send path (needsCanonicalTokenCount in agents/client.js) recounts the quote-merged prompt every turn regardless of the stored value. Mirror that on the client β€” estimate quoted turns from the merged text+quotes and ignore the stored count β€” so snapshot-less branches don't under-report by the quote block. Reverts the earlier "trust the count" assumption, which the server disproves. * 🧹 chore: Route useResumableSSE diagnostics through the frontend logger Convert the [ResumableSSE]/[Debug] console.log and console.error diagnostics to the gated frontend `logger` (client/src/utils/logger), splitting the tag from the message so object arguments are passed through as real args (logged expandably, not stringified) and the logs stay tag-filterable and off the production console unless explicitly enabled. All log statements preserved; nothing removed. * 🩹 fix: Prefer content over text when estimating count-less messages A stopped agent response is saved with both a `text` field and a structured `content` array, and the send path formats from content. messageChars early-returned on `text`, dropping the content array (and the tool-call tokens it carries) from the snapshot-less estimate β€” also making the tool_call handling dead for such messages. Prefer content when present, fall back to text. --- .../ContextProjectionController.js | 35 -- .../limiters/contextProjectionLimiter.js | 19 - api/server/middleware/limiters/index.js | 2 - api/server/routes/endpoints.js | 9 - .../Chat/Input/TokenUsage/Breakdown.tsx | 3 + client/src/data-provider/Endpoints/queries.ts | 39 --- client/src/hooks/Chat/useTokenUsage.ts | 118 +++---- client/src/hooks/SSE/useResumableSSE.ts | 76 ++-- client/src/locales/en/translation.json | 1 + client/src/utils/tokens.spec.ts | 130 +++++++ client/src/utils/tokens.ts | 117 ++++++- packages/api/src/endpoints/index.ts | 1 - packages/api/src/endpoints/projection.spec.ts | 206 ----------- packages/api/src/endpoints/projection.ts | 326 ------------------ packages/data-provider/src/api-endpoints.ts | 2 - packages/data-provider/src/data-service.ts | 7 - packages/data-provider/src/keys.ts | 1 - packages/data-provider/src/types/runs.ts | 23 -- .../data-schemas/src/methods/message.spec.ts | 27 -- packages/data-schemas/src/methods/message.ts | 107 +----- 20 files changed, 331 insertions(+), 918 deletions(-) delete mode 100644 api/server/controllers/ContextProjectionController.js delete mode 100644 api/server/middleware/limiters/contextProjectionLimiter.js delete mode 100644 packages/api/src/endpoints/projection.spec.ts delete mode 100644 packages/api/src/endpoints/projection.ts diff --git a/api/server/controllers/ContextProjectionController.js b/api/server/controllers/ContextProjectionController.js deleted file mode 100644 index 9c56b2ae34..0000000000 --- a/api/server/controllers/ContextProjectionController.js +++ /dev/null @@ -1,35 +0,0 @@ -const { logger } = require('@librechat/data-schemas'); -const { resolveContextProjection } = require('@librechat/api'); -const db = require('~/models'); - -/** - * Returns a server-side context-usage projection for the viewed branch + config - * (agents SDK, no model call) β€” powers the gauge for snapshot-less branches and - * after a model/window switch. Resolution lives in `@librechat/api`; this - * controller only injects request-scoped model accessors. - * @param {ServerRequest} req - * @param {ServerResponse} res - */ -async function contextProjectionController(req, res) { - try { - const params = req.body ?? {}; - if (!params.conversationId || !params.messageId) { - res.json(null); - return; - } - const projection = await resolveContextProjection( - { - userId: req.user?.id, - getMessages: db.getMessages, - getMessageTextStats: db.getMessageTextStats, - }, - params, - ); - res.json(projection ?? null); - } catch (error) { - logger.error('[contextProjectionController]', error); - res.status(500).json({ error: 'Failed to resolve context projection' }); - } -} - -module.exports = contextProjectionController; diff --git a/api/server/middleware/limiters/contextProjectionLimiter.js b/api/server/middleware/limiters/contextProjectionLimiter.js deleted file mode 100644 index 1f70c7ea8e..0000000000 --- a/api/server/middleware/limiters/contextProjectionLimiter.js +++ /dev/null @@ -1,19 +0,0 @@ -const rateLimit = require('express-rate-limit'); -const { limiterCache } = require('@librechat/api'); - -const { CONTEXT_PROJECTION_WINDOW = 1, CONTEXT_PROJECTION_MAX = 20 } = process.env; - -const windowMs = (parseInt(CONTEXT_PROJECTION_WINDOW, 10) || 1) * 60 * 1000; -const max = parseInt(CONTEXT_PROJECTION_MAX, 10) || 20; - -const contextProjectionLimiter = rateLimit({ - windowMs, - max, - handler: (_req, res) => { - res.status(429).json({ message: 'Too many context projection requests. Try again later' }); - }, - keyGenerator: (req) => req.user?.id, - store: limiterCache('context_projection_limiter'), -}); - -module.exports = contextProjectionLimiter; diff --git a/api/server/middleware/limiters/index.js b/api/server/middleware/limiters/index.js index 19f246d039..4a569e2698 100644 --- a/api/server/middleware/limiters/index.js +++ b/api/server/middleware/limiters/index.js @@ -9,7 +9,6 @@ const registerLimiter = require('./registerLimiter'); const toolCallLimiter = require('./toolCallLimiter'); const messageLimiters = require('./messageLimiters'); const promptUsageLimiter = require('./promptUsageLimiter'); -const contextProjectionLimiter = require('./contextProjectionLimiter'); const verifyEmailLimiter = require('./verifyEmailLimiter'); const resetPasswordLimiter = require('./resetPasswordLimiter'); const twoFactorTempLimiter = require('./twoFactorTempLimiter'); @@ -23,7 +22,6 @@ module.exports = { loginLimiter, registerLimiter, toolCallLimiter, - contextProjectionLimiter, createTTSLimiters, createSTTLimiters, verifyEmailLimiter, diff --git a/api/server/routes/endpoints.js b/api/server/routes/endpoints.js index ea55a9e54a..8b1fceccc4 100644 --- a/api/server/routes/endpoints.js +++ b/api/server/routes/endpoints.js @@ -3,19 +3,10 @@ const requireJwtAuth = require('~/server/middleware/requireJwtAuth'); const configMiddleware = require('~/server/middleware/config/app'); const endpointController = require('~/server/controllers/EndpointController'); const tokenConfigController = require('~/server/controllers/TokenConfigController'); -const contextProjectionController = require('~/server/controllers/ContextProjectionController'); -const { contextProjectionLimiter } = require('~/server/middleware/limiters'); const router = express.Router(); /** Auth required for role/tenant-scoped endpoint config resolution. */ router.get('/', requireJwtAuth, endpointController); router.get('/token-config', requireJwtAuth, configMiddleware, tokenConfigController); -router.post( - '/context-projection', - requireJwtAuth, - contextProjectionLimiter, - configMiddleware, - contextProjectionController, -); module.exports = router; diff --git a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx b/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx index f8c479b68a..fe62305b62 100644 --- a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx +++ b/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx @@ -150,6 +150,9 @@ export default function Breakdown({ view, showCost, currency }: BreakdownProps) label={localize('com_ui_output')} value={view.branchTotals.output + view.liveTokens} /> + {view.estimatedTokens > 0 && ( + + )} {maxTokens == null && (

{localize('com_ui_context_unknown')}

)} diff --git a/client/src/data-provider/Endpoints/queries.ts b/client/src/data-provider/Endpoints/queries.ts index 389a9c44cf..d3887e90c6 100644 --- a/client/src/data-provider/Endpoints/queries.ts +++ b/client/src/data-provider/Endpoints/queries.ts @@ -41,45 +41,6 @@ export const useTokenConfigQuery = ( }); }; -/** - * Server-side context-usage projection for the viewed branch + resolved config - * (agents SDK, no model call). Keyed on the fields that change what the next - * call would send — branch tail, endpoint/model/agent, window — so a branch or - * model/window switch refetches. Disabled until a branch tail is known. - */ -export const useContextProjectionQuery = ( - params: t.TContextProjectionRequest | null, - config?: UseQueryOptions, -): QueryObserverResult => { - const queriesEnabled = useRecoilValue(store.queriesEnabled); - return useQuery( - [ - QueryKeys.contextProjection, - params?.conversationId, - params?.messageId, - params?.endpoint, - params?.model, - params?.agentId, - params?.maxContextTokens, - params?.revision, - ], - () => dataService.getContextProjection(params as t.TContextProjectionRequest), - { - staleTime: Infinity, - refetchOnWindowFocus: false, - refetchOnReconnect: false, - refetchOnMount: false, - ...config, - enabled: - (config?.enabled ?? true) === true && - queriesEnabled && - params != null && - (params.conversationId?.length ?? 0) > 0 && - (params.messageId?.length ?? 0) > 0, - }, - ); -}; - /** * Auth-aware query key so unauthenticated (login page) and authenticated * (chat page) configs are cached independently, preventing stale diff --git a/client/src/hooks/Chat/useTokenUsage.ts b/client/src/hooks/Chat/useTokenUsage.ts index 6a6d4ff845..39b6df9220 100644 --- a/client/src/hooks/Chat/useTokenUsage.ts +++ b/client/src/hooks/Chat/useTokenUsage.ts @@ -2,13 +2,7 @@ import { useEffect, useMemo, useRef } from 'react'; import { useAtomValue, useSetAtom } from 'jotai'; import { useQueryClient } from '@tanstack/react-query'; import { Constants, QueryKeys } from 'librechat-data-provider'; -import type { - TMessage, - TConversation, - TModelTokenomics, - TContextUsageEvent, - TContextProjectionRequest, -} from 'librechat-data-provider'; +import type { TMessage, TConversation, TModelTokenomics } from 'librechat-data-provider'; import type { BranchTotals, BranchUsage } from '~/utils/tokens'; import type { ContextSnapshot } from '~/store/usage'; import { @@ -30,7 +24,6 @@ import { findBranchSnapshotAnchor, } from '~/utils'; import { useLatestMessageId } from '~/hooks/Messages/useLatestMessage'; -import { useContextProjectionQuery } from '~/data-provider'; import useTokenLimits from './useTokenLimits'; export interface TokenUsageParams { @@ -60,6 +53,9 @@ export interface TokenUsageView { /** Authoritative cost across all branches (shown when it differs from branch) */ totalCost: number; liveTokens: number; + /** Estimated tokens for count-less messages (in-flight tail excluded while + * streaming); 0 on snapshots. Rendered as its own breakdown row. */ + estimatedTokens: number; rates?: TModelTokenomics; } @@ -101,40 +97,6 @@ export default function useTokenUsage({ return anchor != null ? (snapshotsByAnchor.get(anchor) ?? null) : null; }, [conversationKey, branchTotals.tailId, snapshotsByAnchor]); - const resolvedMax = limits.maxContextTokens; - - /** Project the branch (agents SDK, no model call) ONLY when no persisted/live - * snapshot covers it — snapshot-less branches (G2: pre-feature history, - * imports, never-generated branches). A present snapshot stays authoritative; - * reliable window-switch (G1) detection needs the snapshot to carry its - * model/window (deferred to the fidelity follow-up), and the SDK window - * (reserve-derived) doesn't equal the client-resolved raw window, so we must - * NOT mis-flag a valid snapshot as stale here. Cached + refetched by branch/ - * endpoint/model/window/revision. */ - const projectionParams: TContextProjectionRequest | null = - !isSubmitting && - branchSnapshot == null && - conversation?.conversationId != null && - conversation.conversationId !== Constants.NEW_CONVO && - branchTotals.tailId != null && - conversation.endpoint != null - ? { - conversationId: conversation.conversationId, - messageId: branchTotals.tailId, - /** Resolved provider/model (e.g. an agent's actual provider, not the - * `agents` endpoint) so the server picks the right tokenizer. */ - endpoint: limits.endpoint || conversation.endpoint, - model: limits.model || conversation.model || undefined, - agentId: conversation.agent_id ?? undefined, - spec: conversation.spec ?? undefined, - maxContextTokens: resolvedMax, - /** Content revision so an in-place message edit (same tail id) refetches. */ - revision: branchTotals.input + branchTotals.output, - } - : null; - const { data: projectionData } = useContextProjectionQuery(projectionParams); - const projection = projectionData ?? null; - /** Branch/total provider usage is index-derived; the in-flight response is * the only live add (the pending holder), counted into both — it sits on the * active branch tail and inside the conversation. The backend prices each @@ -242,25 +204,11 @@ export default function useTokenUsage({ snapshot != null && (isSubmitting || (snapshot.anchorMessageId != null && branchTotals.containsAnchor)); - /** Precedence: live/active snapshot → persisted branch snapshot → server - * projection (snapshot-less branches, G2) → per-message estimate. The first - * two preserve the pre-projection behavior exactly; the projection only - * slots in ahead of the estimate when no snapshot exists. Snapshot and - * projection share the render-relevant fields, so they render uniformly. */ - let effective: ContextSnapshot | TContextUsageEvent | null = null; - /** A server projection is the SDK's windowing but, in this first cut, omits - * instruction/tool overhead — so it's surfaced as an ESTIMATE (a better one - * than sumBranch), never a false-authoritative number. Real snapshots stay - * authoritative. */ - let projected = false; - if (currentActive) { - effective = snapshot; - } else if (branchSnapshot != null) { - effective = branchSnapshot; - } else if (projection != null) { - effective = projection; - projected = true; - } + /** Precedence: live/active snapshot → persisted branch snapshot → + * per-message estimate. The first two are authoritative (real runs with the + * feature on); the estimate covers snapshot-less branches (pre-feature + * history, imports, never-generated branches) entirely client-side. */ + const effective: ContextSnapshot | null = currentActive ? snapshot : branchSnapshot; if (effective != null) { const breakdown = effective.breakdown; @@ -270,18 +218,18 @@ export default function useTokenUsage({ effective.remainingContextTokens != null ? maxTokens - effective.remainingContextTokens : instructionTokens + breakdown.messageTokens; - /** The snapshot/projection is pre-invoke: in-flight output rides on - * `liveTokens` (0 unless streaming this branch), the last call's finalized - * output on `completedOutputTokens` (absent on a projection → 0). */ + /** The snapshot is pre-invoke: in-flight output rides on `liveTokens` (0 + * unless streaming this branch), the last call's finalized output on + * `completedOutputTokens`. */ const usedTokens = Math.max(0, baseUsed) + liveTokens + (effective.completedOutputTokens ?? 0); return { usedTokens, maxTokens, percent: maxTokens > 0 ? Math.min((usedTokens / maxTokens) * 100, 100) : 0, - isEstimate: projected, - snapshot: projected ? null : (effective as ContextSnapshot), - snapshotActive: !projected, + isEstimate: false, + snapshot: effective, + snapshotActive: true, branchTotals, branchUsage, totalUsage, @@ -289,18 +237,38 @@ export default function useTokenUsage({ branchCost: branchUsage.cost, totalCost: totalUsage.cost, liveTokens, + estimatedTokens: 0, rates: limits.rates, }; } - /** `summaryBaseline` is the compacted-context size from the deepest - * summarized response on the branch (0 if none). The branch walk stops - * there, so input/output are post-summary only — adding the baseline keeps - * the estimate from re-summing the discarded pre-summary history (which - * otherwise pins the gauge at 100% forever after a compaction). */ - const usedTokens = - branchTotals.input + branchTotals.output + branchTotals.summaryBaseline + liveTokens; + /** Snapshot-less estimate, computed from the in-memory message index — no + * server round-trip. All terms are local per-message counts / char estimates + * (uncalibrated): the learned calibration ratio reconciles provider-injected + * context that isn't present in this visible text, so applying it here would + * over-inflate. `summaryBaseline` is the compacted-context size from the + * deepest summarized response on the branch (0 if none); the walk stops + * there, so input/output are post-summary only — adding it keeps the estimate + * from re-summing the discarded pre-summary history (which otherwise pins the + * gauge at 100% after a compaction). */ const maxTokens = limits.maxContextTokens; + /** When a stream is live the tail is the in-flight response, already counted + * by `liveTokens`; drop its static estimate so a resumed/partial response + * isn't double-counted on the estimate path. */ + const estimatedTokens = Math.max( + 0, + branchTotals.estTokens - (liveTokens > 0 ? branchTotals.tailEstTokens : 0), + ); + const rawUsed = + branchTotals.input + + branchTotals.output + + estimatedTokens + + branchTotals.summaryBaseline + + liveTokens; + /** The send path prunes an over-window branch before calling the model, so the + * live gauge never actually exceeds the window; clamp the display to the + * window rather than show impossible values (e.g. 50k / 8k). */ + const usedTokens = maxTokens != null && maxTokens > 0 ? Math.min(rawUsed, maxTokens) : rawUsed; return { usedTokens, maxTokens, @@ -316,6 +284,7 @@ export default function useTokenUsage({ branchCost: branchUsage.cost, totalCost: totalUsage.cost, liveTokens, + estimatedTokens, rates: limits.rates, }; }, [ @@ -328,6 +297,5 @@ export default function useTokenUsage({ liveTokens, limits, branchSnapshot, - projection, ]); } diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 1eb7267044..501f8a73c5 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -27,6 +27,7 @@ import type { EventHandlerParams } from './useEventHandlers'; import type { ActiveJobsResponse } from '~/data-provider'; import type { TResData } from '~/common'; import { + logger, clearAllDrafts, removeConvoFromAllQueries, upsertConvoInAllQueries, @@ -510,7 +511,7 @@ export default function useResumableSSE( const baseUrl = `${apiBaseUrl()}/api/agents/chat/stream/${encodeURIComponent(currentStreamId)}`; const url = isResume ? `${baseUrl}?resume=true` : baseUrl; - console.log('[ResumableSSE] Subscribing to stream:', url, { isResume }); + logger.log('ResumableSSE', 'Subscribing to stream:', url, { isResume }); const sse = new SSE(url, { headers: { Authorization: `Bearer ${token}` }, @@ -519,7 +520,7 @@ export default function useResumableSSE( sseRef.current = sse; sse.addEventListener('open', () => { - console.log('[ResumableSSE] Stream connected'); + logger.log('ResumableSSE', 'Stream connected'); setAbortScroll(false); // Restore UI state on successful connection (including reconnection) setIsSubmitting(true); @@ -532,7 +533,7 @@ export default function useResumableSSE( const data = JSON.parse(e.data); if (data.final != null) { - console.log('[ResumableSSE] Received FINAL event', { + logger.log('ResumableSSE', 'Received FINAL event', { aborted: data.aborted, conversationId: data.conversation?.conversationId, hasResponseMessage: !!data.responseMessage, @@ -545,7 +546,7 @@ export default function useResumableSSE( finalHandler(data, currentSubmission as EventSubmission); finalizeUsage(data, { ...currentSubmission, userMessage }); } catch (error) { - console.error('[ResumableSSE] Error in finalHandler:', error); + logger.error('ResumableSSE', 'Error in finalHandler:', error); setIsSubmitting(false); setShowStopButton(false); } @@ -562,7 +563,7 @@ export default function useResumableSSE( } if (data.created != null) { - console.log('[ResumableSSE] Received CREATED event', { + logger.log('ResumableSSE', 'Received CREATED event', { messageId: data.message?.messageId, conversationId: data.message?.conversationId, }); @@ -634,7 +635,7 @@ export default function useResumableSSE( } if (data.sync != null) { - console.log('[ResumableSSE] SYNC received', { + logger.log('ResumableSSE', 'SYNC received', { runSteps: data.resumeState?.runSteps?.length ?? 0, pendingEvents: data.pendingEvents?.length ?? 0, }); @@ -700,7 +701,7 @@ export default function useResumableSSE( ); } - console.log('[ResumableSSE] SYNC update', { + logger.log('ResumableSSE', 'SYNC update', { userMsgId, serverResponseId, responseIdx, @@ -721,7 +722,7 @@ export default function useResumableSSE( model: preferDefinedString(messages[responseIdx]?.model, data.resumeState.model), } as TMessage; const updated = mergeResumeMessages(messages, userMessage, responseMessage); - console.log('[ResumableSSE] SYNC updating message', { + logger.log('ResumableSSE', 'SYNC updating message', { messageId: responseMessage.messageId, oldContentLength: Array.isArray(oldContent) ? oldContent.length : 0, newContentLength: data.resumeState.aggregatedContent?.length, @@ -729,7 +730,7 @@ export default function useResumableSSE( setMessages(updated); resetContentHandler(); syncStepMessage(responseMessage); - console.log('[ResumableSSE] SYNC complete, handlers synced'); + logger.log('ResumableSSE', 'SYNC complete, handlers synced'); } else { const responseId = serverResponseId ?? `${userMsgId}_`; const newMessage = { @@ -753,8 +754,9 @@ export default function useResumableSSE( } if (data.resumeState?.replayEvents?.length > 0) { - console.log( - `[ResumableSSE] Replaying ${data.resumeState.replayEvents.length} resume events`, + logger.log( + 'ResumableSSE', + `Replaying ${data.resumeState.replayEvents.length} resume events`, ); for (const replayEvent of data.resumeState.replayEvents) { if (replayEvent.event === UsageEvents.ON_CONTEXT_USAGE) { @@ -774,7 +776,7 @@ export default function useResumableSSE( } if (data.pendingEvents?.length > 0) { - console.log(`[ResumableSSE] Replaying ${data.pendingEvents.length} pending events`); + logger.log('ResumableSSE', `Replaying ${data.pendingEvents.length} pending events`); for (const pendingEvent of data.pendingEvents) { if (pendingEvent.event === 'title') { titleHandler(pendingEvent); @@ -827,7 +829,7 @@ export default function useResumableSSE( messageHandler(text, { ...currentSubmission, userMessage, initialResponse }); } } catch (error) { - console.error('[ResumableSSE] Error processing message:', error); + logger.error('ResumableSSE', 'Error processing message:', error); } }); @@ -848,7 +850,7 @@ export default function useResumableSSE( // Invalidate cache once so react-query refetches instead of showing an error. if (responseCode === 404) { const convoId = currentSubmission.conversation?.conversationId; - console.log('[ResumableSSE] Stream 404, invalidating messages for:', convoId); + logger.log('ResumableSSE', 'Stream 404, invalidating messages for:', convoId); sse.close(); removeActiveJob(currentStreamId); /** Terminal: drop any in-flight live estimate so the gauge doesn't @@ -893,7 +895,7 @@ export default function useResumableSSE( sse.stream(); return; } catch (error) { - console.log('[ResumableSSE] Token refresh failed:', error); + logger.log('ResumableSSE', 'Token refresh failed:', error); } } @@ -905,7 +907,7 @@ export default function useResumableSSE( * not a server-sent error payload. Use `== null` to only match undefined/null (no HTTP status). */ if (responseCode == null && e.data) { - console.log('[ResumableSSE] Server-sent error event received:', e.data); + logger.log('ResumableSSE', 'Server-sent error event received:', e.data); sse.close(); removeActiveJob(currentStreamId); resetLive({ ...currentSubmission, userMessage }); @@ -936,7 +938,7 @@ export default function useResumableSSE( // Not JSON or parsing failed - treat as generic error } - console.log('[ResumableSSE] Error type check:', { isKnownError, errorString }); + logger.log('ResumableSSE', 'Error type check:', { isKnownError, errorString }); // Display the error to user via errorHandler errorHandler({ @@ -944,7 +946,7 @@ export default function useResumableSSE( submission: currentSubmission as EventSubmission, }); } catch (parseError) { - console.error('[ResumableSSE] Failed to parse server error:', parseError); + logger.error('ResumableSSE', 'Failed to parse server error:', parseError); errorHandler({ data: { text: e.data } as unknown as Parameters[0]['data'], submission: currentSubmission as EventSubmission, @@ -961,7 +963,7 @@ export default function useResumableSSE( } // Network failure or unknown HTTP error - attempt reconnection with backoff - console.log('[ResumableSSE] Stream error (network failure) - will attempt reconnect', { + logger.log('ResumableSSE', 'Stream error (network failure) - will attempt reconnect', { responseCode, hasData: !!e.data, }); @@ -971,8 +973,9 @@ export default function useResumableSSE( reconnectAttemptRef.current++; const delay = Math.min(1000 * Math.pow(2, reconnectAttemptRef.current - 1), 30000); - console.log( - `[ResumableSSE] Reconnecting in ${delay}ms (attempt ${reconnectAttemptRef.current}/${MAX_RETRIES})`, + logger.log( + 'ResumableSSE', + `Reconnecting in ${delay}ms (attempt ${reconnectAttemptRef.current}/${MAX_RETRIES})`, ); sse.close(); @@ -989,7 +992,7 @@ export default function useResumableSSE( setIsSubmitting(true); setShowStopButton(true); } else { - console.error('[ResumableSSE] Max reconnect attempts reached'); + logger.error('ResumableSSE', 'Max reconnect attempts reached'); sse.close(); errorHandler({ data: undefined, submission: currentSubmission as EventSubmission }); /** Terminal: clear the in-flight live estimate like the other @@ -1020,11 +1023,11 @@ export default function useResumableSSE( // If we're in a reconnection cycle, don't reset state // (error handler will set up the reconnect timeout) if (reconnectAttemptRef.current > 0) { - console.log('[ResumableSSE] Stream closed for reconnect - preserving state'); + logger.log('ResumableSSE', 'Stream closed for reconnect - preserving state'); return; } - console.log('[ResumableSSE] Stream aborted (intentional close) - no reconnect'); + logger.log('ResumableSSE', 'Stream aborted (intentional close) - no reconnect'); // Clear any pending reconnect attempts if (reconnectTimeoutRef.current) { clearTimeout(reconnectTimeoutRef.current); @@ -1055,14 +1058,14 @@ export default function useResumableSSE( /** Simulate network drop - triggers error event → reconnection */ debugWindow.__killNetwork = () => { - console.log('[Debug] Simulating network drop...'); + logger.log('Debug', 'Simulating network drop...'); // @ts-ignore - sse.js types are incorrect, dispatchEvent actually takes Event sse.dispatchEvent(new Event('error')); }; /** Simulate clean close (navigation away) - triggers abort event → no reconnection */ debugWindow.__closeClean = () => { - console.log('[Debug] Simulating clean close (navigation away)...'); + logger.log('Debug', 'Simulating clean close (navigation away)...'); sse.close(); }; } @@ -1131,7 +1134,7 @@ export default function useResumableSSE( if (signal?.aborted) { return null; } - console.log('[ResumableSSE] Generation started:', { streamId: data.streamId }); + logger.log('ResumableSSE', 'Generation started:', { streamId: data.streamId }); return data.streamId; } catch (error) { if (signal?.aborted) { @@ -1158,8 +1161,9 @@ export default function useResumableSSE( const limit = isServerNotReady ? `${Math.ceil(START_GENERATION_READINESS_TIMEOUT_MS / 1000)}s readiness window` : `${START_GENERATION_NETWORK_RETRIES}`; - console.log( - `[ResumableSSE] ${reason} starting generation, retrying in ${retryDelay}ms (attempt ${attempt}/${limit})`, + logger.log( + 'ResumableSSE', + `${reason} starting generation, retrying in ${retryDelay}ms (attempt ${attempt}/${limit})`, ); const shouldContinue = await waitForRetryDelay(retryDelay, signal); if (!shouldContinue) { @@ -1177,7 +1181,7 @@ export default function useResumableSSE( return null; } - console.error('[ResumableSSE] Error starting generation:', lastError); + logger.error('ResumableSSE', 'Error starting generation:', lastError); const axiosError = lastError as { response?: { data?: Record } }; const errorData = axiosError?.response?.data; @@ -1195,7 +1199,7 @@ export default function useResumableSSE( useEffect(() => { if (!submission || Object.keys(submission).length === 0) { - console.log('[ResumableSSE] No submission, cleaning up'); + logger.log('ResumableSSE', 'No submission, cleaning up'); // Clear reconnect timeout if submission is cleared if (reconnectTimeoutRef.current) { clearTimeout(reconnectTimeoutRef.current); @@ -1213,7 +1217,7 @@ export default function useResumableSSE( } const resumeStreamId = (submission as TSubmission & { resumeStreamId?: string }).resumeStreamId; - console.log('[ResumableSSE] Effect triggered', { + logger.log('ResumableSSE', 'Effect triggered', { conversationId: submission.conversation?.conversationId, hasResumeStreamId: !!resumeStreamId, resumeStreamId, @@ -1237,14 +1241,14 @@ export default function useResumableSSE( return; } // Resume: just subscribe to existing stream, don't start new generation - console.log('[ResumableSSE] Resuming existing stream:', resumeStreamId); + logger.log('ResumableSSE', 'Resuming existing stream:', resumeStreamId); setStreamId(resumeStreamId); // Optimistically add to active jobs (in case it's not already there) addActiveJob(resumeStreamId); subscribeToStream(resumeStreamId, submission, true); // isResume=true } else { // New generation: start and then subscribe - console.log('[ResumableSSE] Starting NEW generation'); + logger.log('ResumableSSE', 'Starting NEW generation'); const newStreamId = await startGeneration(submission, signal); if (signal.aborted) { return; @@ -1268,7 +1272,7 @@ export default function useResumableSSE( submissionRef.current = streamSubmission; subscribeToStream(newStreamId, streamSubmission); } else { - console.error('[ResumableSSE] Failed to get streamId from startGeneration'); + logger.error('ResumableSSE', 'Failed to get streamId from startGeneration'); } } }; @@ -1276,7 +1280,7 @@ export default function useResumableSSE( initStream(); return () => { - console.log('[ResumableSSE] Cleanup - closing SSE, resetting UI state'); + logger.log('ResumableSSE', 'Cleanup - closing SSE, resetting UI state'); startController.abort(); // Cleanup on unmount/navigation - close connection but DO NOT abort backend // Reset UI state so it doesn't leak to other conversations diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 8190bf2de3..f602a7c352 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -949,6 +949,7 @@ "com_ui_context_cost": "Cost", "com_ui_context_cost_branch": "Cost (this branch)", "com_ui_context_cost_total": "All branches", + "com_ui_context_estimated": "Estimated", "com_ui_context_filter_sort": "Filter and Sort by Context", "com_ui_context_free": "Free space", "com_ui_context_messages": "Messages", diff --git a/client/src/utils/tokens.spec.ts b/client/src/utils/tokens.spec.ts index 725fdab224..eddc7a2c89 100644 --- a/client/src/utils/tokens.spec.ts +++ b/client/src/utils/tokens.spec.ts @@ -88,6 +88,136 @@ describe('token index', () => { expect(altTotals.output).toBe(1019); }); + it('estimates count-less messages by text length without inflating counted totals', () => { + buildIndex(CONVO, [ + msg('u1', Constants.NO_PARENT, true, 12), + /** Imported message with no `tokenCount`: 40 chars of text → ~10 est tokens. */ + { + messageId: 'a1', + parentMessageId: 'u1', + isCreatedByUser: false, + conversationId: CONVO, + text: 'x'.repeat(40), + } as TMessage, + ]); + + const totals = sumBranch(CONVO, 'a1'); + /** Known counts feed input/output/counted; the count-less message stays out + * of those and lands in the separate (uncalibrated) estimate bucket. */ + expect(totals.input).toBe(12); + expect(totals.output).toBe(0); + expect(totals.counted).toBe(1); + expect(totals.total).toBe(2); + expect(totals.estTokens).toBe(10); + }); + + it('estimates object-form content text and merged quote excerpts', () => { + buildIndex(CONVO, [ + /** Assistant body lives only in object-form content (`text.value`). */ + { + messageId: 'a1', + parentMessageId: Constants.NO_PARENT, + isCreatedByUser: false, + conversationId: CONVO, + content: [{ type: 'text', text: { value: 'y'.repeat(20) } }], + } as unknown as TMessage, + /** User turn whose quotes are merged into the prompt at send time. */ + { + messageId: 'u1', + parentMessageId: 'a1', + isCreatedByUser: true, + conversationId: CONVO, + text: 'z'.repeat(16), + quotes: ['q'.repeat(8)], + } as TMessage, + ]); + + /** a1: 20 content chars / 4 = 5; u1: (16 text + 8 quote) / 4 = 6. */ + const totals = sumBranch(CONVO, 'u1'); + expect(totals.counted).toBe(0); + expect(totals.estTokens).toBe(11); + }); + + it('recounts quoted user turns (ignoring stale counts), counts tool calls, skips reasoning', () => { + buildIndex(CONVO, [ + /** Quoted user turn with a stale text-only stored count: the send path + * recounts the merged prompt every turn, so the estimate ignores the count + * and recounts from text+quotes. */ + { + messageId: 'u1', + parentMessageId: Constants.NO_PARENT, + isCreatedByUser: true, + conversationId: CONVO, + tokenCount: 999, + text: 'hi', + quotes: ['q'.repeat(38)], + } as TMessage, + /** Count-less assistant turn: tool-call name/args/output count toward the + * estimate (sent back as context); reasoning does not. */ + { + messageId: 'a1', + parentMessageId: 'u1', + isCreatedByUser: false, + conversationId: CONVO, + content: [ + { type: 'think', think: 'r'.repeat(40) }, + { type: 'tool_call', tool_call: { name: 'sub', args: 'aa', output: 'o'.repeat(11) } }, + ], + } as unknown as TMessage, + ]); + + const totals = sumBranch(CONVO, 'a1'); + /** u1 quoted: stored 999 ignored; (2 text + 38 quote) / 4 = 10. a1 tool_call + * name 3 + args 2 + output 11 = 16 / 4 = 4 (think skipped). */ + expect(totals.input).toBe(0); + expect(totals.counted).toBe(0); + expect(totals.estTokens).toBe(14); + }); + + it('prefers content over text for count-less messages carrying both', () => { + buildIndex(CONVO, [ + msg('u1', Constants.NO_PARENT, true, 8), + /** Stopped agent response: saved with both a short `text` and structured + * `content` (a tool call). The send path formats from content, so the + * estimate must use content (tool tokens), not the shorter text. */ + { + messageId: 'a1', + parentMessageId: 'u1', + isCreatedByUser: false, + conversationId: CONVO, + text: 'hi', + content: [ + { type: 'tool_call', tool_call: { name: 'run', args: 'aa', output: 'o'.repeat(13) } }, + ], + } as unknown as TMessage, + ]); + + const totals = sumBranch(CONVO, 'a1'); + /** a1 uses content (name 3 + args 2 + output 13 = 18 / 4 = 5), not text 'hi'. */ + expect(totals.input).toBe(8); + expect(totals.estTokens).toBe(5); + }); + + it('exposes the count-less tail estimate so live output is not double-counted', () => { + buildIndex(CONVO, [ + msg('u1', Constants.NO_PARENT, true, 12), + /** In-flight / resumed response: count-less, so it lands in estTokens; it is + * also covered by liveTokens, so the estimate path drops tailEstTokens. */ + { + messageId: 'a1', + parentMessageId: 'u1', + isCreatedByUser: false, + conversationId: CONVO, + text: 'o'.repeat(20), + } as TMessage, + ]); + + const totals = sumBranch(CONVO, 'a1'); + /** a1 is the tail: 20 / 4 = 5, surfaced both in estTokens and tailEstTokens. */ + expect(totals.estTokens).toBe(5); + expect(totals.tailEstTokens).toBe(5); + }); + it('caps the branch at a summary marker instead of re-summing compacted history', () => { const summarized = { messageId: 'a2', diff --git a/client/src/utils/tokens.ts b/client/src/utils/tokens.ts index adbf65a097..1372aad111 100644 --- a/client/src/utils/tokens.ts +++ b/client/src/utils/tokens.ts @@ -27,6 +27,11 @@ export const EMPTY_USAGE: BranchUsage = { export interface TokenEntry { tokenCount: number; + /** Char/4 token estimate used in place of `tokenCount`: count-less (imported / + * pre-feature) message bodies, plus quoted user turns (recounted from merged + * text+quotes, since the send path ignores their stored count). Includes + * tool-call name/args/output. Mutually exclusive with a counted `tokenCount`. */ + estTokens: number; isCreatedByUser: boolean; parentMessageId: string | null; /** Per-response provider usage from `metadata.usage` (response messages only) */ @@ -46,6 +51,13 @@ export interface BranchTotals { counted: number; /** Total messages on the branch */ total: number; + /** Uncalibrated estimate sum for count-less branch messages (imports / + * pre-feature). Kept separate so known counts aren't re-estimated. */ + estTokens: number; + /** The tail (latest) message's own `estTokens`. When a stream is live the tail + * is the in-flight response, already covered by `liveTokens`, so the estimate + * path excludes this to avoid double-counting a resumed/partial response. */ + tailEstTokens: number; tailId: string | null; /** Whether the latest run's anchor message is on this branch */ containsAnchor: boolean; @@ -63,6 +75,8 @@ export const EMPTY_BRANCH: BranchTotals = { output: 0, counted: 0, total: 0, + estTokens: 0, + tailEstTokens: 0, tailId: null, containsAnchor: false, usage: EMPTY_USAGE, @@ -137,11 +151,101 @@ function addUsage(target: BranchUsage, usage?: BranchUsage): void { } } +/** Chars of a content part's text, handling both the string and `{ value }` forms. + * Reasoning (`think`) and error parts are excluded — the send path strips them + * before counting, so they aren't part of the next call's context. */ +function partTextChars(part: unknown): number { + if (part == null || typeof part !== 'object') { + return 0; + } + const type = (part as { type?: unknown }).type; + if (type === 'think' || type === 'error') { + return 0; + } + if (type === 'tool_call') { + const call = (part as { tool_call?: { name?: unknown; args?: unknown; output?: unknown } }) + .tool_call; + if (call == null) { + return 0; + } + let chars = typeof call.name === 'string' ? call.name.length : 0; + if (typeof call.args === 'string') { + chars += call.args.length; + } else if (call.args != null) { + chars += JSON.stringify(call.args).length; + } + if (typeof call.output === 'string') { + chars += call.output.length; + } + return chars; + } + const text = (part as { text?: unknown }).text; + if (typeof text === 'string') { + return text.length; + } + if ( + text != null && + typeof text === 'object' && + typeof (text as { value?: unknown }).value === 'string' + ) { + return (text as { value: string }).value.length; + } + return 0; +} + +/** Char length of a message's rendered text, for estimating count-less messages. + * Prefer structured `content` when present — the send path formats from it (incl. + * tool calls), so a message carrying both `text` and `content` (e.g. a stopped + * agent response) would otherwise drop its content/tool-call tokens. */ +function messageChars(message: Partial): number { + if (Array.isArray(message.content) && message.content.length > 0) { + let chars = 0; + for (const part of message.content) { + chars += partTextChars(part); + } + return chars; + } + if (typeof message.text === 'string') { + return message.text.length; + } + return 0; +} + +/** Quoted excerpts the send path merges into a user message's prompt. */ +function quoteChars(message: Partial): number { + if (!Array.isArray(message.quotes)) { + return 0; + } + let chars = 0; + for (const quote of message.quotes) { + if (typeof quote === 'string') { + chars += quote.length; + } + } + return chars; +} + function toEntry(message: Partial): TokenEntry { const summaryUsedTokens = message.metadata?.summaryUsedTokens; + const tokenCount = typeof message.tokenCount === 'number' ? message.tokenCount : 0; + const isCreatedByUser = message.isCreatedByUser === true; + const quoted = isCreatedByUser && Array.isArray(message.quotes) && message.quotes.length > 0; + /** A quoted user turn's stored `tokenCount` is unreliable: a text-only Save edit + * recomputes it from `text` alone, and the send path recounts the quote-merged + * prompt every turn regardless (`needsCanonicalTokenCount` in agents/client.js). + * So mirror the server — estimate quoted turns from the merged text+quotes and + * ignore the stored count. Other count-less imports / pre-feature messages + * estimate from text. */ + let estTokens = 0; + if (quoted) { + estTokens = Math.round((messageChars(message) + quoteChars(message)) / 4); + } else if (tokenCount === 0) { + estTokens = Math.round(messageChars(message) / 4); + } return { - tokenCount: typeof message.tokenCount === 'number' ? message.tokenCount : 0, - isCreatedByUser: message.isCreatedByUser === true, + tokenCount: quoted ? 0 : tokenCount, + estTokens, + isCreatedByUser, parentMessageId: message.parentMessageId ?? null, usage: readPersistedUsage(message), summaryUsedTokens: @@ -231,7 +335,10 @@ export function sumBranch( return EMPTY_BRANCH; } - const totals = { input: 0, output: 0, counted: 0, total: 0, containsAnchor: false }; + const totals = { input: 0, output: 0, counted: 0, total: 0, estTokens: 0, containsAnchor: false }; + /** The in-flight response, when streaming, is the branch tail and is covered by + * `liveTokens`; expose its estimate so the estimate path can drop it. */ + const tailEstTokens = index.get(tailId)?.estTokens ?? 0; const usage: BranchUsage = { ...EMPTY_USAGE }; let summaryBaseline = 0; /** Once a summary marker is crossed, older turns are out of the CONTEXT WINDOW @@ -262,6 +369,8 @@ export function sumBranch( } else { totals.output += entry.tokenCount; } + } else if (!contextCapped && entry.estTokens > 0) { + totals.estTokens += entry.estTokens; } /** Cost/usage is cumulative spend — never truncated at the summary boundary. */ addUsage(usage, entry.usage); @@ -275,7 +384,7 @@ export function sumBranch( currentId = entry.parentMessageId; } - return { ...totals, tailId, usage, summaryBaseline }; + return { ...totals, tailEstTokens, tailId, usage, summaryBaseline }; } /** diff --git a/packages/api/src/endpoints/index.ts b/packages/api/src/endpoints/index.ts index 4be03df1e3..9e6e9dbac0 100644 --- a/packages/api/src/endpoints/index.ts +++ b/packages/api/src/endpoints/index.ts @@ -6,5 +6,4 @@ export * from './google'; export * from './models'; export * from './openai'; export * from './pricing'; -export * from './projection'; export * from './tokenConfig'; diff --git a/packages/api/src/endpoints/projection.spec.ts b/packages/api/src/endpoints/projection.spec.ts deleted file mode 100644 index 1fb7f91a4c..0000000000 --- a/packages/api/src/endpoints/projection.spec.ts +++ /dev/null @@ -1,206 +0,0 @@ -import { resolveContextProjection } from './projection'; -import { QUOTE_MAX_COUNT } from '~/utils/quotes'; - -jest.mock('@librechat/agents', () => ({ - Providers: { OPENAI: 'openai' }, - createTokenCounter: jest.fn(async () => jest.fn(() => 1)), - projectAgentContextUsage: jest.fn(() => ({ tokenCount: 1, maxContextTokens: 1000 })), -})); - -const GRAPH_SELECT = 'messageId parentMessageId metadata.summaryUsedTokens'; -const BODY_SELECT = 'messageId parentMessageId tokenCount isCreatedByUser text quotes'; - -function textStats(messageId: string, textBytes = 5) { - return { - messageId, - textBytes, - quoteCount: 0, - quoteBytes: 0, - quoteLineCount: 0, - nonStringQuoteCount: 0, - }; -} - -describe('resolveContextProjection', () => { - const baseParams = { - conversationId: 'conversation-1', - messageId: 'message-1', - endpoint: 'openai', - maxContextTokens: 1000, - model: 'gpt-4o', - }; - - beforeEach(() => { - jest.clearAllMocks(); - }); - - it('returns null before tokenization when the conversation is too large', async () => { - const { createTokenCounter } = jest.requireMock('@librechat/agents'); - const messages = Array.from({ length: 513 }, (_, index) => ({ - messageId: `message-${index}`, - parentMessageId: index === 0 ? null : `message-${index - 1}`, - isCreatedByUser: true, - text: 'hello', - })); - const getMessages = jest.fn(async () => messages); - const getMessageTextStats = jest.fn(); - - const result = await resolveContextProjection( - { userId: 'user-1', getMessages, getMessageTextStats }, - { ...baseParams, messageId: 'message-512' }, - ); - - expect(result).toBeNull(); - expect(getMessages).toHaveBeenCalledTimes(1); - expect(getMessages).toHaveBeenCalledWith( - { conversationId: 'conversation-1', user: 'user-1' }, - GRAPH_SELECT, - { limit: 513, sort: false }, - ); - expect(getMessageTextStats).not.toHaveBeenCalled(); - expect(createTokenCounter).not.toHaveBeenCalled(); - }); - - it('returns null before tokenization when the branch is too long', async () => { - const { createTokenCounter } = jest.requireMock('@librechat/agents'); - const messages = Array.from({ length: 257 }, (_, index) => ({ - messageId: `message-${index}`, - parentMessageId: index === 0 ? null : `message-${index - 1}`, - isCreatedByUser: true, - text: 'hello', - })); - const getMessages = jest.fn(async () => messages); - const getMessageTextStats = jest.fn(); - - const result = await resolveContextProjection( - { userId: 'user-1', getMessages, getMessageTextStats }, - { ...baseParams, messageId: 'message-256' }, - ); - - expect(result).toBeNull(); - expect(getMessages).toHaveBeenCalledTimes(1); - expect(getMessageTextStats).not.toHaveBeenCalled(); - expect(createTokenCounter).not.toHaveBeenCalled(); - }); - - it('returns null before loading bodies when the branch text is too large', async () => { - const { createTokenCounter } = jest.requireMock('@librechat/agents'); - const getMessages = jest.fn(async () => [ - { - messageId: 'message-1', - parentMessageId: null, - }, - ]); - const getMessageTextStats = jest.fn(async () => [textStats('message-1', 512 * 1024 + 1)]); - const result = await resolveContextProjection( - { - userId: 'user-1', - getMessages, - getMessageTextStats, - }, - baseParams, - ); - - expect(result).toBeNull(); - expect(getMessages).toHaveBeenCalledTimes(1); - expect(getMessageTextStats).toHaveBeenCalledWith( - { - conversationId: 'conversation-1', - user: 'user-1', - messageId: { $in: ['message-1'] }, - }, - { limit: 1 }, - ); - expect(createTokenCounter).not.toHaveBeenCalled(); - }); - - it('loads only branch message bodies after resolving the graph', async () => { - const graph = [ - { messageId: 'message-1', parentMessageId: null }, - { messageId: 'message-2', parentMessageId: 'message-1' }, - { messageId: 'off-branch', parentMessageId: null }, - ]; - const bodies = [ - { - messageId: 'message-1', - parentMessageId: null, - isCreatedByUser: true, - text: 'first', - tokenCount: 5, - }, - { - messageId: 'message-2', - parentMessageId: 'message-1', - isCreatedByUser: false, - text: 'second', - tokenCount: 6, - }, - ]; - const getMessages = jest.fn(async (_filter: object, select?: string) => - select === GRAPH_SELECT ? graph : bodies, - ); - const getMessageTextStats = jest.fn(async () => [ - textStats('message-1', 5), - textStats('message-2', 6), - ]); - - const result = await resolveContextProjection( - { userId: 'user-1', getMessages, getMessageTextStats }, - { ...baseParams, messageId: 'message-2' }, - ); - - expect(result).toEqual({ tokenCount: 1, maxContextTokens: 1000 }); - expect(getMessages).toHaveBeenNthCalledWith( - 1, - { conversationId: 'conversation-1', user: 'user-1' }, - GRAPH_SELECT, - { limit: 513, sort: false }, - ); - expect(getMessageTextStats).toHaveBeenCalledWith( - { - conversationId: 'conversation-1', - user: 'user-1', - messageId: { $in: ['message-1', 'message-2'] }, - }, - { limit: 2 }, - ); - expect(getMessages).toHaveBeenNthCalledWith( - 2, - { - conversationId: 'conversation-1', - user: 'user-1', - messageId: { $in: ['message-1', 'message-2'] }, - }, - BODY_SELECT, - { limit: 2, sort: false }, - ); - }); - - it('returns null before loading bodies when a branch message has too many quotes', async () => { - const { createTokenCounter } = jest.requireMock('@librechat/agents'); - const getMessages = jest.fn(async () => [ - { - messageId: 'message-1', - parentMessageId: null, - }, - ]); - const getMessageTextStats = jest.fn(async () => [ - { - ...textStats('message-1'), - quoteCount: QUOTE_MAX_COUNT + 1, - quoteBytes: 10, - quoteLineCount: QUOTE_MAX_COUNT + 1, - }, - ]); - - const result = await resolveContextProjection( - { userId: 'user-1', getMessages, getMessageTextStats }, - baseParams, - ); - - expect(result).toBeNull(); - expect(getMessages).toHaveBeenCalledTimes(1); - expect(getMessageTextStats).toHaveBeenCalledTimes(1); - expect(createTokenCounter).not.toHaveBeenCalled(); - }); -}); diff --git a/packages/api/src/endpoints/projection.ts b/packages/api/src/endpoints/projection.ts deleted file mode 100644 index 8622e4b981..0000000000 --- a/packages/api/src/endpoints/projection.ts +++ /dev/null @@ -1,326 +0,0 @@ -import { HumanMessage, AIMessage } from '@langchain/core/messages'; -import { Providers, createTokenCounter, projectAgentContextUsage } from '@librechat/agents'; -import type { TContextProjectionRequest, TContextUsageEvent } from 'librechat-data-provider'; -import type { BaseMessage } from '@langchain/core/messages'; -import { QUOTE_MAX_COUNT, mergeQuotedText } from '~/utils/quotes'; - -const MAX_PROJECTION_MESSAGES = 512; -const MAX_PROJECTION_BRANCH_MESSAGES = 256; -const MAX_PROJECTION_BRANCH_TEXT_BYTES = 512 * 1024; -const PROJECTION_GRAPH_SELECT = 'messageId parentMessageId metadata.summaryUsedTokens'; -const PROJECTION_BODY_SELECT = 'messageId parentMessageId tokenCount isCreatedByUser text quotes'; - -interface ProjectionMessage { - messageId: string; - parentMessageId?: string | null; - tokenCount?: number; - isCreatedByUser?: boolean; - text?: string; - /** Quoted excerpts merged into the model-facing text by the live path; must be - * included here so the context gauge counts the same prompt the model sees. */ - quotes?: string[]; - /** Compaction marker written by the live path (`agents/usage.ts`); its - * presence means the next call sends the summary + tail, not this raw chain. */ - metadata?: { summaryUsedTokens?: number }; -} - -interface ProjectionMessageFilter { - conversationId: string; - user?: string; - messageId?: string | { $in: string[] }; -} - -interface ProjectionMessageQueryOptions { - limit?: number; - sort?: false; -} - -interface ProjectionMessageTextStats { - messageId: string; - textBytes: number; - quoteCount: number; - quoteBytes: number; - quoteLineCount: number; - nonStringQuoteCount: number; -} - -interface ProjectionMessageTextStatsOptions { - limit?: number; -} - -export interface ContextProjectionDeps { - /** Authenticated requester — branch lookups are scoped to this user. */ - userId?: string; - getMessages: ( - filter: ProjectionMessageFilter, - select?: string, - options?: ProjectionMessageQueryOptions, - ) => Promise; - getMessageTextStats: ( - filter: ProjectionMessageFilter, - options?: ProjectionMessageTextStatsOptions, - ) => Promise; -} - -/** - * Walks the parent chain from `tailId` to root and returns the branch messages - * oldest→newest. The visited set guards against cycles / self-referential links. - */ -function resolveBranch(messages: ProjectionMessage[], tailId: string): ProjectionMessage[] { - const byId = new Map(); - for (const message of messages) { - byId.set(message.messageId, message); - } - const branch: ProjectionMessage[] = []; - const seen = new Set(); - let currentId: string | null | undefined = tailId; - while (currentId != null && !seen.has(currentId)) { - const message = byId.get(currentId); - if (message == null) { - break; - } - seen.add(currentId); - branch.push(message); - currentId = message.parentMessageId; - } - return branch.reverse(); -} - -function hasValidProjectionIds(params: TContextProjectionRequest): boolean { - return typeof params.conversationId === 'string' && typeof params.messageId === 'string'; -} - -function getProjectionText(message: ProjectionMessage): string | null { - const hasQuotes = - message.isCreatedByUser === true && Array.isArray(message.quotes) && message.quotes.length > 0; - if (!hasQuotes) { - return message.text ?? ''; - } - if (message.quotes == null || message.quotes.length > QUOTE_MAX_COUNT) { - return null; - } - for (const quote of message.quotes) { - if (typeof quote !== 'string') { - return null; - } - } - return mergeQuotedText(message.text ?? '', message.quotes); -} - -function hasExceededBranchTextLimit(branch: ProjectionMessage[]): boolean { - let bytes = 0; - for (const message of branch) { - const text = getProjectionText(message); - if (text == null) { - return true; - } - bytes += Buffer.byteLength(text, 'utf8'); - if (bytes > MAX_PROJECTION_BRANCH_TEXT_BYTES) { - return true; - } - } - return false; -} - -function getEstimatedMergedTextBytes(stats: ProjectionMessageTextStats): number | null { - if ( - stats.nonStringQuoteCount > 0 || - stats.quoteCount > QUOTE_MAX_COUNT || - stats.quoteLineCount < stats.quoteCount - ) { - return null; - } - if (stats.quoteCount === 0) { - return stats.textBytes; - } - - const quotePrefixBytes = stats.quoteLineCount * 2; - const quoteLineBreakBytes = stats.quoteLineCount - stats.quoteCount; - const quoteSeparatorBytes = (stats.quoteCount - 1) * 2; - const bodySeparatorBytes = stats.textBytes > 0 ? 2 : 0; - return ( - stats.textBytes + - stats.quoteBytes + - quotePrefixBytes + - quoteLineBreakBytes + - quoteSeparatorBytes + - bodySeparatorBytes - ); -} - -function hasExceededBranchTextStatsLimit(stats: ProjectionMessageTextStats[]): boolean { - let bytes = 0; - for (const messageStats of stats) { - const messageBytes = getEstimatedMergedTextBytes(messageStats); - if (messageBytes == null) { - return true; - } - bytes += messageBytes; - if (bytes > MAX_PROJECTION_BRANCH_TEXT_BYTES) { - return true; - } - } - return false; -} - -/** Maps an endpoint/provider string to the agents `Providers` enum. */ -function resolveProvider(value?: string): Providers { - if (value == null || value === '') { - return Providers.OPENAI; - } - const lower = value.toLowerCase(); - for (const provider of Object.values(Providers)) { - if (provider.toLowerCase() === lower) { - return provider; - } - } - if (lower.includes('anthropic') || lower.includes('claude')) { - return Providers.ANTHROPIC; - } - if (lower.includes('google') || lower.includes('gemini') || lower.includes('vertex')) { - return Providers.GOOGLE; - } - if (lower.includes('bedrock')) { - return Providers.BEDROCK; - } - return Providers.OPENAI; -} - -async function getBranchMessages( - deps: ContextProjectionDeps, - baseFilter: ProjectionMessageFilter, - branch: ProjectionMessage[], -): Promise { - const branchIds = branch.map((message) => message.messageId); - const stats = await deps.getMessageTextStats( - { ...baseFilter, messageId: { $in: branchIds } }, - { limit: branchIds.length }, - ); - if (stats.length !== branchIds.length || hasExceededBranchTextStatsLimit(stats)) { - return null; - } - - const stored = await deps.getMessages( - { ...baseFilter, messageId: { $in: branchIds } }, - PROJECTION_BODY_SELECT, - { limit: branchIds.length, sort: false }, - ); - if (stored.length !== branchIds.length) { - return null; - } - const byId = new Map(); - for (const message of stored) { - byId.set(message.messageId, message); - } - const ordered: ProjectionMessage[] = []; - for (const messageId of branchIds) { - const message = byId.get(messageId); - if (message == null) { - return null; - } - ordered.push(message); - } - return ordered; -} - -/** - * Server-side context-usage projection: reconstructs the viewed branch and asks - * the agents SDK what the next call's context would be, WITHOUT invoking the - * model. Provider/model/window come from the (client-resolved) request — no - * agent or model-spec config is loaded here, so there is no cross-user config - * exposure. Reuses LibreChat's already-calibrated per-message `tokenCount`s (no - * re-tokenizing). Returns null when there is no resolvable context window. - * NOTE: this first cut targets message-windowing accuracy — instruction and - * tool-schema tokens (agent instructions, `promptPrefix`, model-spec presets, - * tool schemas) are NOT yet included; a follow-up will reuse the full - * `initializeAgent`/send path for exact overhead and proper access control. - */ -export async function resolveContextProjection( - deps: ContextProjectionDeps, - params: TContextProjectionRequest, -): Promise { - if (!hasValidProjectionIds(params)) { - return null; - } - - const maxContextTokens = params.maxContextTokens; - if (maxContextTokens == null || maxContextTokens <= 0) { - return null; - } - - const baseFilter = { conversationId: params.conversationId, user: deps.userId }; - const stored = await deps.getMessages(baseFilter, PROJECTION_GRAPH_SELECT, { - limit: MAX_PROJECTION_MESSAGES + 1, - sort: false, - }); - if (stored.length > MAX_PROJECTION_MESSAGES) { - return null; - } - - const branch = resolveBranch(stored, params.messageId); - if (branch.length === 0) { - return null; - } - if (branch.length > MAX_PROJECTION_BRANCH_MESSAGES) { - return null; - } - - /** A summarized/compacted branch's next call sends the saved summary + the - * post-summary tail, NOT this raw parent chain — projecting from the full - * history would prune/count the wrong context and omit the summary. Detect it - * via the live path's `metadata.summaryUsedTokens` marker and fall back (null) - * so the client's summary-baseline-aware estimate handles these branches until - * a follow-up replays the summary boundary. */ - if (branch.some((message) => (message.metadata?.summaryUsedTokens ?? 0) > 0)) { - return null; - } - - const bodyBranch = await getBranchMessages(deps, baseFilter, branch); - if (bodyBranch == null || hasExceededBranchTextLimit(bodyBranch)) { - return null; - } - - const model = params.model; - const encoding = (model ?? '').toLowerCase().includes('claude') ? 'claude' : 'o200k_base'; - const tokenCounter = await createTokenCounter(encoding); - - const messages: BaseMessage[] = []; - const indexTokenCountMap: Record = {}; - for (let i = 0; i < bodyBranch.length; i++) { - const message = bodyBranch[i]; - /** Mirror the live path: prepend quoted excerpts into the user text the model - * receives so the gauge counts the same prompt. */ - const hasQuotes = - message.isCreatedByUser === true && - Array.isArray(message.quotes) && - message.quotes.length > 0; - const text = getProjectionText(message); - if (text == null) { - return null; - } - const lcMessage = - message.isCreatedByUser === true ? new HumanMessage(text) : new AIMessage(text); - messages.push(lcMessage); - /** Recount messages with no stored count (imported / pre-feature) rather - * than charging 0 — a real 0 and "unknown" must not collapse, or the - * snapshot-less histories this endpoint targets would under-report. Also - * recount quoted messages: a text-only Save edit leaves a stale text-only - * `tokenCount` that omits the quote block, so trust the merged recount. */ - indexTokenCountMap[String(i)] = - !hasQuotes && message.tokenCount != null && message.tokenCount > 0 - ? message.tokenCount - : tokenCounter(lcMessage); - } - - return projectAgentContextUsage({ - agent: { - agentId: params.agentId ?? 'projection', - provider: resolveProvider(params.endpoint), - maxContextTokens, - }, - messages, - tokenCounter, - indexTokenCountMap, - calibrationRatio: params.calibrationRatio, - }); -} diff --git a/packages/data-provider/src/api-endpoints.ts b/packages/data-provider/src/api-endpoints.ts index 5661fc5166..4cc2fd1c47 100644 --- a/packages/data-provider/src/api-endpoints.ts +++ b/packages/data-provider/src/api-endpoints.ts @@ -160,8 +160,6 @@ export const aiEndpoints = () => `${BASE_URL}/api/endpoints`; export const tokenConfig = () => `${BASE_URL}/api/endpoints/token-config`; -export const contextProjection = () => `${BASE_URL}/api/endpoints/context-projection`; - export const models = () => `${BASE_URL}/api/models`; export const tokenizer = () => `${BASE_URL}/api/tokenizer`; diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index 1a1eea87f8..4a4721d186 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -1,5 +1,4 @@ import type { AxiosResponse } from 'axios'; -import type { TContextProjectionRequest, TContextUsageEvent } from './types/runs'; import type { TFileConfig } from './file-config'; import type * as t from './types'; import * as permissions from './accessPermissions'; @@ -259,12 +258,6 @@ export const getTokenConfig = (): Promise => { return request.get(endpoints.tokenConfig()); }; -export const getContextProjection = ( - payload: TContextProjectionRequest, -): Promise => { - return request.post(endpoints.contextProjection(), payload); -}; - export const getModels = async (): Promise => { return request.get(endpoints.models()); }; diff --git a/packages/data-provider/src/keys.ts b/packages/data-provider/src/keys.ts index dde54005a6..26eed7462a 100644 --- a/packages/data-provider/src/keys.ts +++ b/packages/data-provider/src/keys.ts @@ -14,7 +14,6 @@ export enum QueryKeys { balance = 'balance', endpoints = 'endpoints', tokenConfig = 'tokenConfig', - contextProjection = 'contextProjection', presets = 'presets', searchResults = 'searchResults', tokenCount = 'tokenCount', diff --git a/packages/data-provider/src/types/runs.ts b/packages/data-provider/src/types/runs.ts index 8e1033c07f..1fdc07cd81 100644 --- a/packages/data-provider/src/types/runs.ts +++ b/packages/data-provider/src/types/runs.ts @@ -88,29 +88,6 @@ export type TContextUsageEvent = { completedOutputTokens?: number; }; -/** - * Request payload for a server-side context-usage projection: "what context - * would the next call send for this branch under this config", computed by the - * agents SDK without invoking the model. Powers the gauge in states the live - * snapshot can't cover (page load of a snapshot-less branch, window/model - * switch). `messageId` is the viewed branch's tail; the server walks its parent - * chain. - */ -export type TContextProjectionRequest = { - conversationId: string; - messageId: string; - endpoint: string; - model?: string; - agentId?: string; - spec?: string; - maxContextTokens?: number; - /** Provider-calibrated ratio from a prior snapshot, applied as a static seed. */ - calibrationRatio?: number; - /** Client-only cache-bust: a branch content revision so a message edit - * (which keeps the same tail id) refetches. The server ignores it. */ - revision?: number; -}; - /** * Per-response usage rollup persisted on `responseMessage.metadata.usage`, in * display units (input excludes cache; output includes repaired completion). diff --git a/packages/data-schemas/src/methods/message.spec.ts b/packages/data-schemas/src/methods/message.spec.ts index 7e3747d12c..4b4f6b015b 100644 --- a/packages/data-schemas/src/methods/message.spec.ts +++ b/packages/data-schemas/src/methods/message.spec.ts @@ -21,7 +21,6 @@ let mongoServer: InstanceType; let Message: mongoose.Model; let saveMessage: ReturnType['saveMessage']; let getMessages: ReturnType['getMessages']; -let getMessageTextStats: ReturnType['getMessageTextStats']; let updateMessage: ReturnType['updateMessage']; let deleteMessages: ReturnType['deleteMessages']; let bulkSaveMessages: ReturnType['bulkSaveMessages']; @@ -40,7 +39,6 @@ beforeAll(async () => { const methods = createMessageMethods(mongoose); saveMessage = methods.saveMessage; getMessages = methods.getMessages; - getMessageTextStats = methods.getMessageTextStats; updateMessage = methods.updateMessage; deleteMessages = methods.deleteMessages; bulkSaveMessages = methods.bulkSaveMessages; @@ -273,31 +271,6 @@ describe('Message Operations', () => { expect(messages[0].text).toBe('First message'); expect(messages[1].text).toBe('Second message'); }); - - it('should retrieve message text stats without returning message bodies', async () => { - const conversationId = uuidv4(); - - await saveMessage(mockCtx, { - messageId: 'msg1', - conversationId, - text: 'hello', - quotes: ['a\nb', ''], - user: 'user123', - }); - - const stats = await getMessageTextStats({ conversationId, user: 'user123' }, { limit: 1 }); - - expect(stats).toEqual([ - { - messageId: 'msg1', - textBytes: 5, - quoteCount: 2, - quoteBytes: 3, - quoteLineCount: 3, - nonStringQuoteCount: 0, - }, - ]); - }); }); describe('deleteMessages', () => { diff --git a/packages/data-schemas/src/methods/message.ts b/packages/data-schemas/src/methods/message.ts index 42273cbf5f..a66666f5c6 100644 --- a/packages/data-schemas/src/methods/message.ts +++ b/packages/data-schemas/src/methods/message.ts @@ -1,5 +1,5 @@ import { RetentionMode } from 'librechat-data-provider'; -import type { DeleteResult, FilterQuery, Model, PipelineStage } from 'mongoose'; +import type { DeleteResult, FilterQuery, Model } from 'mongoose'; import type { AppConfig, IMessage } from '~/types'; import { createTempChatExpirationDate } from '~/utils/tempChatRetention'; import { createFallbackRetentionDate } from '~/utils/retention'; @@ -14,19 +14,6 @@ interface MessageQueryOptions { sort?: Record | false; } -interface MessageTextStatsOptions { - limit?: number; -} - -export interface MessageTextStats { - messageId: string; - textBytes: number; - quoteCount: number; - quoteBytes: number; - quoteLineCount: number; - nonStringQuoteCount: number; -} - export interface MessageMethods { saveMessage( ctx: { userId: string; isTemporary?: boolean; interfaceConfig?: AppConfig['interfaceConfig'] }, @@ -60,10 +47,6 @@ export interface MessageMethods { select?: string, options?: MessageQueryOptions, ): Promise; - getMessageTextStats( - filter: FilterQuery, - options?: MessageTextStatsOptions, - ): Promise; getMessage(params: { user: string; messageId: string }): Promise; getMessagesByCursor( filter: FilterQuery, @@ -374,93 +357,6 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa } } - async function getMessageTextStats( - filter: FilterQuery, - options: MessageTextStatsOptions = {}, - ) { - try { - const Message = mongoose.models.Message as Model; - const pipeline: PipelineStage[] = [{ $match: filter }]; - if (options.limit != null && options.limit > 0) { - pipeline.push({ $limit: options.limit }); - } - pipeline.push({ - $project: { - _id: 0, - messageId: 1, - textBytes: { - $cond: [{ $eq: [{ $type: '$text' }, 'string'] }, { $strLenBytes: '$text' }, 0], - }, - quoteCount: { - $cond: [{ $isArray: '$quotes' }, { $size: '$quotes' }, 0], - }, - quoteBytes: { - $cond: [ - { $isArray: '$quotes' }, - { - $sum: { - $map: { - input: '$quotes', - as: 'quote', - in: { - $cond: [ - { $eq: [{ $type: '$$quote' }, 'string'] }, - { $strLenBytes: '$$quote' }, - 0, - ], - }, - }, - }, - }, - 0, - ], - }, - quoteLineCount: { - $cond: [ - { $isArray: '$quotes' }, - { - $sum: { - $map: { - input: '$quotes', - as: 'quote', - in: { - $cond: [ - { $eq: [{ $type: '$$quote' }, 'string'] }, - { $size: { $split: ['$$quote', '\n'] } }, - 0, - ], - }, - }, - }, - }, - 0, - ], - }, - nonStringQuoteCount: { - $cond: [ - { $isArray: '$quotes' }, - { - $size: { - $filter: { - input: '$quotes', - as: 'quote', - cond: { $ne: [{ $type: '$$quote' }, 'string'] }, - }, - }, - }, - 0, - ], - }, - }, - }); - - return await Message.aggregate(pipeline); - } catch (err) { - logger.error('Error getting message text stats:', err); - throw err; - } - } - /** * Retrieves a single message from the database. */ @@ -547,7 +443,6 @@ export function createMessageMethods(mongoose: typeof import('mongoose')): Messa updateMessage, deleteMessagesSince, getMessages, - getMessageTextStats, getMessage, getMessagesByCursor, searchMessages,