From 70a218ff82f669db9b90dc06a1a0651322c2a536 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 8 Jan 2026 21:55:33 -0500 Subject: [PATCH] =?UTF-8?q?=E2=8F=B2=EF=B8=8F=20feat:=20Defer=20Loading=20?= =?UTF-8?q?MCP=20Tools=20(#11270)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * WIP: code ptc * refactor: tool classification and calling logic * 🔧 fix: Update @librechat/agents dependency to version 3.0.68 * chore: import order and correct renamed tool name for tool search * refactor: streamline tool classification logic for local and programmatic tools * feat: add per-tool configuration options for agents, including deferred loading and allowed callers - Introduced `tool_options` in agent forms to manage tool behavior. - Updated tool classification logic to prioritize agent-level configurations. - Enhanced UI components to support tool deferral functionality. - Added localization strings for new tool options and actions. * feat: enhance agent schema with per-tool options for configuration - Added `tool_options` schema to support per-tool configurations, including `defer_loading` and `allowed_callers`. - Updated agent data model to incorporate new tool options, ensuring flexibility in tool behavior management. - Modified type definitions to reflect the new `tool_options` structure for agents. * feat: add tool_options parameter to loadTools and initializeAgent for enhanced agent configuration * chore: update @librechat/agents dependency to version 3.0.71 and enhance agent tool loading logic - Updated the @librechat/agents package to version 3.0.71 across multiple files. - Added support for handling deferred loading of tools in agent initialization and execution processes. - Improved the extraction of discovered tools from message history to optimize tool loading behavior. * chore: update @librechat/agents dependency to version 3.0.72 * chore: update @librechat/agents dependency to version 3.0.75 * refactor: simplify tool defer loading logic in MCPTool component - Removed local state management for deferred tools, relying on form state instead. - Updated related functions to directly use form values for checking and toggling defer loading. - Cleaned up code by eliminating unnecessary optimistic updates and local state dependencies. * chore: remove deprecated localization strings for tool deferral in translation.json - Eliminated unused strings related to deferred loading descriptions in the English translation file. - Streamlined localization to reflect recent changes in tool loading logic. * refactor: improve tool defer loading handling in MCPTool component - Enhanced the logic for managing deferred loading of tools by simplifying the update process for tool options. - Ensured that the state reflects the correct loading behavior based on the new deferred loading conditions. - Cleaned up the code to remove unnecessary complexity in handling tool options. * refactor: update agent mocks in callbacks test to use actual implementations - Modified the agent mocks in the callbacks test to include actual implementations from the @librechat/agents module. - This change enhances the accuracy of the tests by ensuring they reflect the real behavior of the agent functions. --- api/package.json | 2 +- .../agents/__tests__/callbacks.spec.js | 4 +- api/server/controllers/agents/callbacks.js | 11 +- api/server/controllers/agents/client.js | 1 + .../services/Endpoints/agents/initialize.js | 20 +- api/server/services/MCP.js | 1 + api/server/services/ToolService.js | 26 +- client/src/common/agents-types.ts | 3 + .../components/Chat/Messages/Content/Part.tsx | 6 +- .../Messages/Content/Parts/ExecuteCode.tsx | 2 +- .../SidePanel/Agents/AgentPanel.tsx | 9 +- .../SidePanel/Agents/AgentSelect.tsx | 5 + .../components/SidePanel/Agents/MCPTool.tsx | 247 +++++++-- client/src/locales/en/translation.json | 9 +- package-lock.json | 17 +- packages/api/package.json | 2 +- packages/api/src/agents/initialize.ts | 25 +- packages/api/src/agents/run.ts | 150 ++++++ packages/api/src/agents/validation.ts | 10 + packages/api/src/mcp/types/index.ts | 8 +- packages/api/src/tools/classification.ts | 491 ++++++++++++++++++ packages/api/src/tools/index.ts | 1 + packages/data-provider/src/config.ts | 2 + packages/data-provider/src/schemas.ts | 1 + .../data-provider/src/types/assistants.ts | 36 ++ packages/data-schemas/src/schema/agent.ts | 5 + packages/data-schemas/src/types/agent.ts | 4 +- 27 files changed, 1016 insertions(+), 82 deletions(-) create mode 100644 packages/api/src/tools/classification.ts diff --git a/api/package.json b/api/package.json index 9e134bd32a..5dcf531b99 100644 --- a/api/package.json +++ b/api/package.json @@ -46,7 +46,7 @@ "@googleapis/youtube": "^20.0.0", "@keyv/redis": "^4.3.3", "@langchain/core": "^0.3.80", - "@librechat/agents": "^3.0.66", + "@librechat/agents": "^3.0.75", "@librechat/api": "*", "@librechat/data-schemas": "*", "@microsoft/microsoft-graph-client": "^3.0.7", diff --git a/api/server/controllers/agents/__tests__/callbacks.spec.js b/api/server/controllers/agents/__tests__/callbacks.spec.js index 7922c31efa..103f9f3236 100644 --- a/api/server/controllers/agents/__tests__/callbacks.spec.js +++ b/api/server/controllers/agents/__tests__/callbacks.spec.js @@ -16,9 +16,7 @@ jest.mock('@librechat/data-schemas', () => ({ })); jest.mock('@librechat/agents', () => ({ - EnvVar: { CODE_API_KEY: 'CODE_API_KEY' }, - Providers: { GOOGLE: 'google' }, - GraphEvents: {}, + ...jest.requireActual('@librechat/agents'), getMessageId: jest.fn(), ToolEndHandler: jest.fn(), handleToolCalls: jest.fn(), diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index aee419577a..ff70d13d37 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -1,6 +1,7 @@ const { nanoid } = require('nanoid'); -const { sendEvent, GenerationJobManager } = require('@librechat/api'); +const { Constants } = require('@librechat/agents'); const { logger } = require('@librechat/data-schemas'); +const { sendEvent, GenerationJobManager } = require('@librechat/api'); const { Tools, StepTypes, FileContext, ErrorTypes } = require('librechat-data-provider'); const { EnvVar, @@ -441,10 +442,10 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null }) return; } - { - if (output.name !== Tools.execute_code) { - return; - } + const isCodeTool = + output.name === Tools.execute_code || output.name === Constants.PROGRAMMATIC_TOOL_CALLING; + if (!isCodeTool) { + return; } if (!output.artifact.files) { diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 79e63d1c7f..f7f1f5161f 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -1019,6 +1019,7 @@ class AgentClient extends BaseClient { run = await createRun({ agents, + messages, indexTokenCountMap, runId: this.responseMessageId, signal: abortController.signal, diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 626beed153..fc054c1e6f 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -43,13 +43,23 @@ function createToolLoader(signal, streamId = null) { * @param {string} params.model * @param {AgentToolResources} params.tool_resources * @returns {Promise<{ - * tools: StructuredTool[], - * toolContextMap: Record, - * userMCPAuthMap?: Record> + * tools: StructuredTool[], + * toolContextMap: Record, + * userMCPAuthMap?: Record>, + * toolRegistry?: import('@librechat/agents').LCToolRegistry * } | undefined>} */ - return async function loadTools({ req, res, agentId, tools, provider, model, tool_resources }) { - const agent = { id: agentId, tools, provider, model }; + return async function loadTools({ + req, + res, + tools, + model, + agentId, + provider, + tool_options, + tool_resources, + }) { + const agent = { id: agentId, tools, provider, model, tool_options }; try { return await loadAgentTools({ req, diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index 81d7107de4..220c1f7d7e 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -548,6 +548,7 @@ function createToolInstance({ }); toolInstance.mcp = true; toolInstance.mcpRawServerName = serverName; + toolInstance.mcpJsonSchema = parameters; return toolInstance; } diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index 1e2074cdf4..8fcb261b41 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -6,6 +6,7 @@ const { hasCustomUserVars, getUserMCPAuthMap, isActionDomainAllowed, + buildToolClassification, } = require('@librechat/api'); const { Tools, @@ -36,6 +37,7 @@ const { recordUsage } = require('~/server/services/Threads'); const { loadTools } = require('~/app/clients/tools/util'); const { redactMessage } = require('~/config/parsers'); const { findPluginAuthsByKeys } = require('~/models'); +const { loadAuthValues } = require('~/server/services/Tools/credentials'); /** * Processes the required actions by calling the appropriate tools and returning the outputs. * @param {OpenAIClient} client - OpenAI or StreamRunManager Client. @@ -367,7 +369,13 @@ async function processRequiredActions(client, requiredActions) { * @param {AbortSignal} params.signal * @param {Pick> }>} The agent tools. + * @returns {Promise<{ + * tools?: StructuredTool[]; + * toolContextMap?: Record; + * userMCPAuthMap?: Record>; + * toolRegistry?: Map; + * hasDeferredTools?: boolean; + * }>} The agent tools and registry. */ async function loadAgentTools({ req, @@ -510,11 +518,23 @@ async function loadAgentTools({ return map; }, {}); + /** Build tool registry from MCP tools and create PTC/tool search tools if configured */ + const { toolRegistry, additionalTools, hasDeferredTools } = await buildToolClassification({ + loadedTools, + userId: req.user.id, + agentId: agent.id, + agentToolOptions: agent.tool_options, + loadAuthValues, + }); + agentTools.push(...additionalTools); + if (!checkCapability(AgentCapabilities.actions)) { return { tools: agentTools, userMCPAuthMap, toolContextMap, + toolRegistry, + hasDeferredTools, }; } @@ -527,6 +547,8 @@ async function loadAgentTools({ tools: agentTools, userMCPAuthMap, toolContextMap, + toolRegistry, + hasDeferredTools, }; } @@ -654,6 +676,8 @@ async function loadAgentTools({ tools: agentTools, toolContextMap, userMCPAuthMap, + toolRegistry, + hasDeferredTools, }; } diff --git a/client/src/common/agents-types.ts b/client/src/common/agents-types.ts index 9ac6b440a3..c3832b7ff8 100644 --- a/client/src/common/agents-types.ts +++ b/client/src/common/agents-types.ts @@ -1,6 +1,7 @@ import { AgentCapabilities, ArtifactModes } from 'librechat-data-provider'; import type { AgentModelParameters, + AgentToolOptions, SupportContact, AgentProvider, GraphEdge, @@ -33,6 +34,8 @@ export type AgentForm = { model: string | null; model_parameters: AgentModelParameters; tools?: string[]; + /** Per-tool configuration options (deferred loading, allowed callers, etc.) */ + tool_options?: AgentToolOptions; provider?: AgentProvider | OptionWithIcon; /** @deprecated Use edges instead */ agent_ids?: string[]; diff --git a/client/src/components/Chat/Messages/Content/Part.tsx b/client/src/components/Chat/Messages/Content/Part.tsx index bfa2b28fac..4a74e3606f 100644 --- a/client/src/components/Chat/Messages/Content/Part.tsx +++ b/client/src/components/Chat/Messages/Content/Part.tsx @@ -91,7 +91,11 @@ const Part = memo( const isToolCall = 'args' in toolCall && (!toolCall.type || toolCall.type === ToolCallTypes.TOOL_CALL); - if (isToolCall && toolCall.name === Tools.execute_code) { + if ( + isToolCall && + (toolCall.name === Tools.execute_code || + toolCall.name === Constants.PROGRAMMATIC_TOOL_CALLING) + ) { return ( (0); const prevShowCodeRef = useRef(showCode); - const { lang, code } = useParseArgs(args) ?? ({} as ParsedArgs); + const { lang = 'py', code } = useParseArgs(args) ?? ({} as ParsedArgs); const progress = useProgress(initialProgress); useEffect(() => { diff --git a/client/src/components/SidePanel/Agents/AgentPanel.tsx b/client/src/components/SidePanel/Agents/AgentPanel.tsx index 86ec27dc5e..f74dcfddcc 100644 --- a/client/src/components/SidePanel/Agents/AgentPanel.tsx +++ b/client/src/components/SidePanel/Agents/AgentPanel.tsx @@ -23,6 +23,7 @@ import { import { createProviderOption, getDefaultAgentFormValues } from '~/utils'; import { useResourcePermissions } from '~/hooks/useResourcePermissions'; import { useSelectAgent, useLocalize, useAuthContext } from '~/hooks'; +import type { TranslationKeys } from '~/hooks/useLocalize'; import { useAgentPanelContext } from '~/Providers/AgentPanelContext'; import AgentPanelSkeleton from './AgentPanelSkeleton'; import AdvancedPanel from './Advanced/AdvancedPanel'; @@ -36,8 +37,8 @@ import ModelPanel from './ModelPanel'; function getUpdateToastMessage( noVersionChange: boolean, avatarActionState: AgentForm['avatar_action'], - name: string | undefined, - localize: (key: string, vars?: Record | Array) => string, + name: string | null | undefined, + localize: (key: TranslationKeys, vars?: Record) => string, ): string | null { // If only avatar upload is pending (separate endpoint), suppress the no-changes toast. if (noVersionChange && avatarActionState === 'upload') { @@ -72,6 +73,7 @@ export function composeAgentUpdatePayload(data: AgentForm, agent_id?: string | n recursion_limit, category, support_contact, + tool_options, avatar_action: avatarActionState, } = data; @@ -97,6 +99,7 @@ export function composeAgentUpdatePayload(data: AgentForm, agent_id?: string | n recursion_limit, category, support_contact, + tool_options, ...(shouldResetAvatar ? { avatar: null } : {}), }, provider, @@ -545,7 +548,7 @@ export default function AgentPanel() { (); + const { getValues, setValue, control } = useFormContext(); const { getServerStatusIconProps, getConfigDialogProps } = useMCPServerManager(); const [isFocused, setIsFocused] = useState(false); const [isHovering, setIsHovering] = useState(false); const [accordionValue, setAccordionValue] = useState(''); + const formToolOptions = useWatch({ control, name: 'tool_options' }); + + /** Check if a specific tool has defer_loading enabled */ + const isToolDeferred = useCallback( + (toolId: string): boolean => formToolOptions?.[toolId]?.defer_loading === true, + [formToolOptions], + ); + + /** Toggle defer_loading for a specific tool */ + const toggleToolDefer = useCallback( + (toolId: string) => { + const currentOptions = getValues('tool_options') || {}; + const currentToolOptions = currentOptions[toolId] || {}; + const newDeferred = !currentToolOptions.defer_loading; + + const updatedOptions: AgentToolOptions = { ...currentOptions }; + + if (newDeferred) { + updatedOptions[toolId] = { + ...currentToolOptions, + defer_loading: true, + }; + } else { + const { defer_loading: _, ...restOptions } = currentToolOptions; + if (Object.keys(restOptions).length === 0) { + delete updatedOptions[toolId]; + } else { + updatedOptions[toolId] = restOptions; + } + } + + setValue('tool_options', updatedOptions, { shouldDirty: true }); + }, + [getValues, setValue], + ); + + /** Check if all server tools are deferred */ + const areAllToolsDeferred = + serverInfo?.tools && + serverInfo.tools.length > 0 && + serverInfo.tools.every((tool) => formToolOptions?.[tool.tool_id]?.defer_loading === true); + + /** Toggle defer_loading for all tools from this server */ + const toggleDeferAll = useCallback(() => { + if (!serverInfo?.tools) return; + + const shouldDefer = !areAllToolsDeferred; + const currentOptions = getValues('tool_options') || {}; + const updatedOptions: AgentToolOptions = { ...currentOptions }; + + for (const tool of serverInfo.tools) { + if (shouldDefer) { + updatedOptions[tool.tool_id] = { + ...(updatedOptions[tool.tool_id] || {}), + defer_loading: true, + }; + } else { + if (updatedOptions[tool.tool_id]) { + delete updatedOptions[tool.tool_id].defer_loading; + if (Object.keys(updatedOptions[tool.tool_id]).length === 0) { + delete updatedOptions[tool.tool_id]; + } + } + } + } + + setValue('tool_options', updatedOptions, { shouldDirty: true }); + }, [serverInfo?.tools, getValues, setValue, areAllToolsDeferred]); + if (!serverInfo) { return null; } @@ -170,6 +241,47 @@ export default function MCPTool({ serverInfo }: { serverInfo?: MCPServerInfo }) /> + {/* Defer All toggle - icon only with tooltip */} + { + e.stopPropagation(); + toggleDeferAll(); + }} + onKeyDown={(e) => { + if (e.key === 'Enter' || e.key === ' ') { + e.preventDefault(); + e.stopPropagation(); + toggleDeferAll(); + } + }} + > + + +
{/* Caret button for accordion */} @@ -230,52 +342,95 @@ export default function MCPTool({ serverInfo }: { serverInfo?: MCPServerInfo })
- {serverInfo.tools?.map((subTool) => ( -