⏲️ feat: Defer Loading MCP Tools (#11270)

* 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/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(),

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) {

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,

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<string, unknown>,
-   * userMCPAuthMap?: Record<string, Record<string, string>>
+   *   tools: StructuredTool[],
+   *   toolContextMap: Record<string, unknown>,
+   *   userMCPAuthMap?: Record<string, Record<string, string>>,
+   *   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,

api/server/services/MCP.js

@@ -548,6 +548,7 @@ function createToolInstance({
   });
   toolInstance.mcp = true;
   toolInstance.mcpRawServerName = serverName;
+  toolInstance.mcpJsonSchema = parameters;
   return toolInstance;
 }
 

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<Agent, 'id' | 'provider' | 'model' | 'tools'} params.agent - The agent to load tools for.
  * @param {string | undefined} [params.openAIApiKey] - The OpenAI API key.
- * @returns {Promise<{ tools?: StructuredTool[]; userMCPAuthMap?: Record<string, Record<string, string>> }>} The agent tools.
+ * @returns {Promise<{
+ *   tools?: StructuredTool[];
+ *   toolContextMap?: Record<string, unknown>;
+ *   userMCPAuthMap?: Record<string, Record<string, string>>;
+ *   toolRegistry?: Map<string, import('~/utils/toolClassification').LCTool>;
+ *   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,
   };
 }