🎯 fix: Use Resolved Provider for Agent Token Lookup on Custom Endpoints (#12574)

* 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

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<typeof import('~/utils')>('~/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();

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,