From 96312aa4fd2425a5e3d0a8c75bd635957b4140ce Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 8 Apr 2026 15:07:38 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=AF=20fix:=20Use=20Resolved=20Provider?= =?UTF-8?q?=20for=20Agent=20Token=20Lookup=20on=20Custom=20Endpoints=20(#1?= =?UTF-8?q?2574)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: Use resolved provider for agent token lookup on custom endpoints The providerEndpointMap lookup in initializeAgent used the original provider name (e.g. "EduGPT") instead of the resolved overrideProvider ("openai"). Since providerEndpointMap only contains 4 built-in providers, custom providers resolved to undefined, causing getModelMaxTokens to miss the token map and fall back to 18000 tokens. With agent instructions + tool schemas consuming most of that budget, createPruneMessages would strip all messages on the first turn. * fix: Use correct EndpointTokenConfig type in test * refactor: Unify test factory, remove non-discriminating test Address review findings: - Remove Test 2 ("uses the model real context window") which passed with and without the fix due to getModelMaxTokens defaulting to openAI when endpoint is undefined (JS default parameter semantics) - Merge createCustomProviderMocks into createMocks via provider, overrideProvider, and useRealTokenLookup parameters - Hoist jest.requireActual to file scope for shared access * refactor: Address followup review findings - Replace loose `maxContextTokens > 18000` assertion with precise computed value `Math.round((65536 - 4096) * 0.95)` so the outcome assertion is meaningful and self-documenting - Hoist `customProvider` to describe-level constant `CUSTOM_PROVIDER` - Document `overrideProvider` semantics and `useRealTokenLookup` in factory JSDoc - Add comment on real `optionalChainWithEmptyCheck` noting its zero-handling semantics are load-bearing for the maxContextTokens=0 test * style: Use // for inline comment, clarify pipeline assertion role --- .../src/agents/__tests__/initialize.test.ts | 142 ++++++++++++++---- packages/api/src/agents/initialize.ts | 2 +- 2 files changed, 117 insertions(+), 27 deletions(-) diff --git a/packages/api/src/agents/__tests__/initialize.test.ts b/packages/api/src/agents/__tests__/initialize.test.ts index f9982a6e46..fe1493628a 100644 --- a/packages/api/src/agents/__tests__/initialize.test.ts +++ b/packages/api/src/agents/__tests__/initialize.test.ts @@ -1,7 +1,7 @@ import { Providers } from '@librechat/agents'; import { EModelEndpoint } from 'librechat-data-provider'; import type { Agent } from 'librechat-data-provider'; -import type { ServerRequest, InitializeResultBase } from '~/types'; +import type { ServerRequest, InitializeResultBase, EndpointTokenConfig } from '~/types'; import type { InitializeAgentDbMethods } from '../initialize'; // Mock logger @@ -55,22 +55,48 @@ jest.mock('../resources', () => ({ import { initializeAgent } from '../initialize'; +const realUtils = jest.requireActual('~/utils'); + /** * Creates minimal mock objects for initializeAgent tests. + * + * @param overrides.overrideProvider - Simulates the value returned by `getProviderConfig`. + * Defaults to `provider` (native endpoint where no remapping occurs). Set to a different + * value (e.g. `Providers.OPENAI`) alongside a custom `provider` to simulate a custom + * endpoint whose provider is resolved to a built-in. + * @param overrides.useRealTokenLookup - When true, `getModelMaxTokens` delegates to the real + * implementation so tests exercise actual token-map resolution. Otherwise a controlled + * `modelDefault` is returned. */ function createMocks(overrides?: { + provider?: string; + overrideProvider?: string; + model?: string; maxContextTokens?: number; modelDefault?: number; maxOutputTokens?: number; + endpointTokenConfig?: EndpointTokenConfig; + useRealTokenLookup?: boolean; }) { - const { maxContextTokens, modelDefault = 200000, maxOutputTokens = 4096 } = overrides ?? {}; + const { + provider = Providers.OPENAI, + overrideProvider, + model = 'test-model', + maxContextTokens, + modelDefault = 200000, + maxOutputTokens = 4096, + endpointTokenConfig, + useRealTokenLookup = false, + } = overrides ?? {}; + + const resolvedOverrideProvider = overrideProvider ?? provider; const agent = { id: 'agent-1', - model: 'test-model', - provider: Providers.OPENAI, + model, + provider, tools: [], - model_parameters: { model: 'test-model' }, + model_parameters: { model }, } as unknown as Agent; const req = { @@ -81,39 +107,29 @@ function createMocks(overrides?: { const res = {} as unknown as import('express').Response; const mockGetOptions = jest.fn().mockResolvedValue({ - llmConfig: { - model: 'test-model', - maxTokens: maxOutputTokens, - }, - endpointTokenConfig: undefined, + llmConfig: { model, maxTokens: maxOutputTokens }, + endpointTokenConfig, } satisfies InitializeResultBase); mockGetProviderConfig.mockReturnValue({ getOptions: mockGetOptions, - overrideProvider: Providers.OPENAI, + overrideProvider: resolvedOverrideProvider, }); - // extractLibreChatParams returns maxContextTokens when provided in model_parameters mockExtractLibreChatParams.mockReturnValue({ resendFiles: false, maxContextTokens, - modelOptions: { model: 'test-model' }, + modelOptions: { model }, }); - // getModelMaxTokens returns the model's default context window - mockGetModelMaxTokens.mockReturnValue(modelDefault); + if (useRealTokenLookup) { + mockGetModelMaxTokens.mockImplementation(realUtils.getModelMaxTokens); + } else { + mockGetModelMaxTokens.mockReturnValue(modelDefault); + } - // Implement real optionalChainWithEmptyCheck behavior - mockOptionalChainWithEmptyCheck.mockImplementation( - (...values: (string | number | undefined)[]) => { - for (const v of values) { - if (v !== undefined && v !== null && v !== '') { - return v; - } - } - return values[values.length - 1]; - }, - ); + // Real implementation: treats 0 as a valid (non-empty) value — load-bearing for the maxContextTokens=0 test + mockOptionalChainWithEmptyCheck.mockImplementation(realUtils.optionalChainWithEmptyCheck); const loadTools = jest.fn().mockResolvedValue({ tools: [], @@ -136,6 +152,80 @@ function createMocks(overrides?: { return { agent, req, res, loadTools, db }; } +describe('initializeAgent — custom provider token lookup', () => { + const CUSTOM_PROVIDER = 'EduGPT'; + + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('passes the resolved provider endpoint to getModelMaxTokens, not the custom name', async () => { + const { agent, req, res, loadTools, db } = createMocks({ + provider: CUSTOM_PROVIDER, + overrideProvider: Providers.OPENAI, + model: 'qwen3-235b-a22b', + useRealTokenLookup: true, + }); + + await initializeAgent( + { + req, + res, + agent, + loadTools, + endpointOption: { endpoint: EModelEndpoint.agents }, + allowedProviders: new Set([CUSTOM_PROVIDER]), + isInitialAgent: true, + }, + db, + ); + + // providerEndpointMap["openAI"] = "openAI" (valid), not providerEndpointMap["EduGPT"] = undefined + expect(mockGetModelMaxTokens).toHaveBeenCalledWith( + 'qwen3-235b-a22b', + EModelEndpoint.openAI, + undefined, + ); + }); + + it('uses endpointTokenConfig from the custom endpoint for unrecognized models', async () => { + const customTokenConfig: EndpointTokenConfig = { + 'my-custom-model-v1': { context: 65536, prompt: 1, completion: 1 }, + }; + const { agent, req, res, loadTools, db } = createMocks({ + provider: CUSTOM_PROVIDER, + overrideProvider: Providers.OPENAI, + model: 'my-custom-model-v1', + endpointTokenConfig: customTokenConfig, + useRealTokenLookup: true, + }); + + const result = await initializeAgent( + { + req, + res, + agent, + loadTools, + endpointOption: { endpoint: EModelEndpoint.agents }, + allowedProviders: new Set([CUSTOM_PROVIDER]), + isInitialAgent: true, + }, + db, + ); + + expect(mockGetModelMaxTokens).toHaveBeenCalledWith( + 'my-custom-model-v1', + EModelEndpoint.openAI, + customTokenConfig, + ); + + // Pipeline check: verifies endpointTokenConfig.context flows through the full + // optionalChainWithEmptyCheck → Math.max formula. The toHaveBeenCalledWith + // assertion above catches the actual provider-resolution regression. + expect(result.maxContextTokens).toBe(Math.round((65536 - 4096) * 0.95)); + }); +}); + describe('initializeAgent — maxContextTokens', () => { beforeEach(() => { jest.clearAllMocks(); diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 81bc89cac4..1ebd20bb01 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -354,7 +354,7 @@ export async function initializeAgent( maxContextTokens, getModelMaxTokens( tokensModel ?? '', - providerEndpointMap[provider as keyof typeof providerEndpointMap], + providerEndpointMap[overrideProvider as keyof typeof providerEndpointMap], options.endpointTokenConfig, ), 18000,