⏲️ 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.
This commit is contained in:
Danny Avila 2026-01-08 21:55:33 -05:00
parent 2958fcd0c5
commit 70a218ff82
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
27 changed files with 1016 additions and 82 deletions

View file

@ -10,14 +10,15 @@ import {
} from 'librechat-data-provider';
import type {
AgentToolResources,
AgentToolOptions,
TEndpointOption,
TFile,
Agent,
TUser,
} from 'librechat-data-provider';
import type { GenericTool, LCToolRegistry, ToolMap } from '@librechat/agents';
import type { Response as ServerResponse } from 'express';
import type { IMongoFile } from '@librechat/data-schemas';
import type { GenericTool } from '@librechat/agents';
import type { InitializeResultBase, ServerRequest, EndpointDbMethods } from '~/types';
import { getModelMaxTokens, extractLibreChatParams, optionalChainWithEmptyCheck } from '~/utils';
import { filterFilesByEndpointConfig } from '~/files';
@ -36,6 +37,12 @@ export type InitializedAgent = Agent & {
useLegacyContent: boolean;
resendFiles: boolean;
userMCPAuthMap?: Record<string, Record<string, string>>;
/** Tool map for ToolNode to use when executing tools (required for PTC) */
toolMap?: ToolMap;
/** Tool registry for PTC and tool search (only present when MCP tools with env classification exist) */
toolRegistry?: LCToolRegistry;
/** Precomputed flag indicating if any tools have defer_loading enabled (for efficient runtime checks) */
hasDeferredTools?: boolean;
};
/**
@ -61,11 +68,14 @@ export interface InitializeAgentParams {
agentId: string;
tools: string[];
model: string | null;
tool_options: AgentToolOptions | undefined;
tool_resources: AgentToolResources | undefined;
}) => Promise<{
tools: GenericTool[];
toolContextMap: Record<string, unknown>;
userMCPAuthMap?: Record<string, Record<string, string>>;
toolRegistry?: LCToolRegistry;
hasDeferredTools?: boolean;
} | null>;
/** Endpoint option (contains model_parameters and endpoint info) */
endpointOption?: Partial<TEndpointOption>;
@ -201,6 +211,8 @@ export async function initializeAgent(
tools: structuredTools,
toolContextMap,
userMCPAuthMap,
toolRegistry,
hasDeferredTools,
} = (await loadTools?.({
req,
res,
@ -208,8 +220,15 @@ export async function initializeAgent(
agentId: agent.id,
tools: agent.tools ?? [],
model: agent.model,
tool_options: agent.tool_options,
tool_resources,
})) ?? { tools: [], toolContextMap: {}, userMCPAuthMap: undefined };
})) ?? {
tools: [],
toolContextMap: {},
userMCPAuthMap: undefined,
toolRegistry: undefined,
hasDeferredTools: false,
};
const { getOptions, overrideProvider } = getProviderConfig({
provider,
@ -312,6 +331,8 @@ export async function initializeAgent(
attachments: finalAttachments,
resendFiles,
userMCPAuthMap,
toolRegistry,
hasDeferredTools,
toolContextMap: toolContextMap ?? {},
useLegacyContent: !!options.useLegacyContent,
maxContextTokens: Math.round((agentMaxContextNum - maxOutputTokensNum) * 0.9),

View file

@ -1,9 +1,11 @@
import { Run, Providers } from '@librechat/agents';
import { providerEndpointMap, KnownEndpoints } from 'librechat-data-provider';
import type { BaseMessage } from '@langchain/core/messages';
import type {
MultiAgentGraphConfig,
OpenAIClientOptions,
StandardGraphConfig,
LCToolRegistry,
AgentInputs,
GenericTool,
RunConfig,
@ -14,6 +16,121 @@ import type { Agent } from 'librechat-data-provider';
import type * as t from '~/types';
import { resolveHeaders, createSafeUser } from '~/utils/env';
/** Tool search tool name constant */
const TOOL_SEARCH_NAME = 'tool_search';
/** Expected shape of JSON tool search results */
interface ToolSearchJsonResult {
found?: number;
tools?: Array<{ name: string }>;
}
/**
* Parses tool names from JSON-formatted tool_search output.
* Format: { "found": N, "tools": [{ "name": "tool_name", ... }], ... }
*
* @param content - The JSON string content
* @param discoveredTools - Set to add discovered tool names to
* @returns true if parsing succeeded, false otherwise
*/
function parseToolSearchJson(content: string, discoveredTools: Set<string>): boolean {
try {
const parsed = JSON.parse(content) as ToolSearchJsonResult;
if (!parsed.tools || !Array.isArray(parsed.tools)) {
return false;
}
for (const tool of parsed.tools) {
if (tool.name && typeof tool.name === 'string') {
discoveredTools.add(tool.name);
}
}
return parsed.tools.length > 0;
} catch {
return false;
}
}
/**
* Parses tool names from legacy text-formatted tool_search output.
* Format: "- tool_name (score: X.XX)"
*
* @param content - The text content
* @param discoveredTools - Set to add discovered tool names to
*/
function parseToolSearchLegacy(content: string, discoveredTools: Set<string>): void {
const toolNameRegex = /^- ([^\s(]+)\s*\(score:/gm;
let match: RegExpExecArray | null;
while ((match = toolNameRegex.exec(content)) !== null) {
const toolName = match[1];
if (toolName) {
discoveredTools.add(toolName);
}
}
}
/**
* Extracts discovered tool names from message history by parsing tool_search results.
* When the LLM calls tool_search, the result contains tool names that were discovered.
* These tools should have defer_loading overridden to false on subsequent turns.
*
* Supports both:
* - New JSON format: { "tools": [{ "name": "tool_name" }] }
* - Legacy text format: "- tool_name (score: X.XX)"
*
* @param messages - The conversation message history
* @returns Set of tool names that were discovered via tool_search
*/
export function extractDiscoveredToolsFromHistory(messages: BaseMessage[]): Set<string> {
const discoveredTools = new Set<string>();
for (const message of messages) {
const msgType = message._getType?.() ?? message.constructor?.name ?? '';
if (msgType !== 'tool') {
continue;
}
const name = (message as { name?: string }).name;
if (name !== TOOL_SEARCH_NAME) {
continue;
}
const content = message.content;
if (typeof content !== 'string') {
continue;
}
/** Try JSON format first (new), fall back to regex (legacy) */
if (!parseToolSearchJson(content, discoveredTools)) {
parseToolSearchLegacy(content, discoveredTools);
}
}
return discoveredTools;
}
/**
* Overrides defer_loading to false for tools that were already discovered via tool_search.
* This prevents the LLM from having to re-discover tools on every turn.
*
* @param toolRegistry - The tool registry to modify (mutated in place)
* @param discoveredTools - Set of tool names that were previously discovered
* @returns Number of tools that had defer_loading overridden
*/
export function overrideDeferLoadingForDiscoveredTools(
toolRegistry: LCToolRegistry,
discoveredTools: Set<string>,
): number {
let overrideCount = 0;
for (const toolName of discoveredTools) {
const toolDef = toolRegistry.get(toolName);
if (toolDef && toolDef.defer_loading === true) {
toolDef.defer_loading = false;
overrideCount++;
}
}
return overrideCount;
}
const customProviders = new Set([
Providers.XAI,
Providers.DEEPSEEK,
@ -48,6 +165,9 @@ type RunAgent = Omit<Agent, 'tools'> & {
maxContextTokens?: number;
useLegacyContent?: boolean;
toolContextMap?: Record<string, string>;
toolRegistry?: LCToolRegistry;
/** Precomputed flag indicating if any tools have defer_loading enabled */
hasDeferredTools?: boolean;
};
/**
@ -60,12 +180,16 @@ type RunAgent = Omit<Agent, 'tools'> & {
* @param options.customHandlers - Custom event handlers.
* @param options.streaming - Whether to use streaming.
* @param options.streamUsage - Whether to stream usage information.
* @param options.messages - Optional message history to extract discovered tools from.
* When provided, tools that were previously discovered via tool_search will have
* their defer_loading overridden to false, preventing redundant re-discovery.
* @returns {Promise<Run<IState>>} A promise that resolves to a new Run instance.
*/
export async function createRun({
runId,
signal,
agents,
messages,
requestBody,
user,
tokenCounter,
@ -81,9 +205,26 @@ export async function createRun({
streamUsage?: boolean;
requestBody?: t.RequestBody;
user?: IUser;
/** Message history for extracting previously discovered tools */
messages?: BaseMessage[];
} & Pick<RunConfig, 'tokenCounter' | 'customHandlers' | 'indexTokenCountMap'>): Promise<
Run<IState>
> {
/**
* Only extract discovered tools if:
* 1. We have message history to parse
* 2. At least one agent has deferred tools (using precomputed flag)
*
* This optimization avoids iterating through messages in the ~95% of cases
* where no agent uses deferred tool loading.
*/
const hasAnyDeferredTools = agents.some((agent) => agent.hasDeferredTools === true);
const discoveredTools =
hasAnyDeferredTools && messages?.length
? extractDiscoveredToolsFromHistory(messages)
: new Set<string>();
const agentInputs: AgentInputs[] = [];
const buildAgentContext = (agent: RunAgent) => {
const provider =
@ -135,6 +276,14 @@ export async function createRun({
llmConfig.usage = true;
}
/**
* Override defer_loading for tools that were discovered in previous turns.
* This prevents the LLM from having to re-discover tools via tool_search.
*/
if (discoveredTools.size > 0 && agent.toolRegistry) {
overrideDeferLoadingForDiscoveredTools(agent.toolRegistry, discoveredTools);
}
const reasoningKey = getReasoningKey(provider, llmConfig, agent.endpoint);
const agentInput: AgentInputs = {
provider,
@ -144,6 +293,7 @@ export async function createRun({
tools: agent.tools,
clientOptions: llmConfig,
instructions: systemContent,
toolRegistry: agent.toolRegistry,
maxContextTokens: agent.maxContextTokens,
useLegacyContent: agent.useLegacyContent ?? false,
};

View file

@ -51,6 +51,15 @@ export const graphEdgeSchema = z.object({
promptKey: z.string().optional(),
});
/** Per-tool options schema (defer_loading, allowed_callers) */
export const toolOptionsSchema = z.object({
defer_loading: z.boolean().optional(),
allowed_callers: z.array(z.enum(['direct', 'code_execution'])).optional(),
});
/** Agent tool options - map of tool_id to tool options */
export const agentToolOptionsSchema = z.record(z.string(), toolOptionsSchema).optional();
/** Base agent schema with all common fields */
export const agentBaseSchema = z.object({
name: z.string().nullable().optional(),
@ -68,6 +77,7 @@ export const agentBaseSchema = z.object({
recursion_limit: z.number().optional(),
conversation_starters: z.array(z.string()).optional(),
tool_resources: agentToolResourcesSchema,
tool_options: agentToolOptionsSchema,
support_contact: agentSupportContactSchema,
category: z.string().optional(),
});

View file

@ -17,7 +17,8 @@ import type {
Tool,
} from '@modelcontextprotocol/sdk/types.js';
import type { SearchResultData, UIResource, TPlugin } from 'librechat-data-provider';
import type { TokenMethods, JsonSchemaType, IUser } from '@librechat/data-schemas';
import type { TokenMethods, IUser } from '@librechat/data-schemas';
import type { LCTool } from '@librechat/agents';
import type { FlowStateManager } from '~/flow/manager';
import type { RequestBody } from '~/types/http';
import type * as o from '~/mcp/oauth/types';
@ -42,11 +43,6 @@ export interface MCPResource {
description?: string;
mimeType?: string;
}
export interface LCTool {
name: string;
description?: string;
parameters: JsonSchemaType;
}
export interface LCFunctionTool {
type: 'function';

View file

@ -0,0 +1,491 @@
/**
* @fileoverview Utility functions for building tool registries from environment variables.
* This is a temporary solution for tool classification until UI-based configuration is available.
*
* Environment Variables:
* - TOOL_PROGRAMMATIC_ONLY: Comma-separated tool names or server patterns (sys__all__sys_mcp_ServerName)
* - TOOL_PROGRAMMATIC_ONLY_EXCLUDE: Comma-separated tool names to exclude from programmatic only
* - TOOL_DUAL_CONTEXT: Comma-separated tool names or server patterns callable BOTH by LLM and PTC
* - TOOL_DUAL_CONTEXT_EXCLUDE: Comma-separated tool names to exclude from dual context
* - TOOL_DEFERRED: Comma-separated tool names or server patterns for deferred tools
* - TOOL_DEFERRED_EXCLUDE: Comma-separated tool names to exclude from deferred
* - TOOL_CLASSIFICATION_AGENT_IDS: Optional comma-separated agent IDs to restrict classification features
*
* Server patterns: Use `sys__all__sys_mcp_ServerName` to match all tools from an MCP server.
* Example: `sys__all__sys_mcp_Google-Workspace` matches all Google Workspace tools.
*
* Agent restriction: If TOOL_CLASSIFICATION_AGENT_IDS is set, only those agents will get
* PTC and tool search tools. If not set, all agents with matching tools get them.
*
* Smart enablement: PTC/tool search are only created if the agent has tools that actually
* match the classification patterns. An agent with no programmatic/deferred tools won't
* get PTC/tool search even if the env vars are set.
*
* @module packages/api/src/tools/classification
*/
import { logger } from '@librechat/data-schemas';
import { Constants } from 'librechat-data-provider';
import { EnvVar, createProgrammaticToolCallingTool, createToolSearch } from '@librechat/agents';
import type { AgentToolOptions } from 'librechat-data-provider';
import type {
LCToolRegistry,
JsonSchemaType,
AllowedCaller,
GenericTool,
LCTool,
} from '@librechat/agents';
export type { LCTool, LCToolRegistry, AllowedCaller, JsonSchemaType };
/** Pattern prefix for matching all tools from an MCP server */
const MCP_ALL_PATTERN = `${Constants.mcp_all}${Constants.mcp_delimiter}`;
export interface ToolDefinition {
name: string;
description?: string;
parameters?: JsonSchemaType;
}
/**
* Parses a comma-separated tool list from an environment variable.
* @param envValue - The environment variable value
* @returns Set of tool names or server patterns
*/
export function parseToolList(envValue: string | undefined): Set<string> {
if (!envValue || envValue.trim() === '') {
return new Set();
}
return new Set(
envValue
.split(',')
.map((s) => s.trim())
.filter((s) => s.length > 0),
);
}
/**
* Extracts the MCP server name from a tool name.
* Tool names follow the pattern: toolName_mcp_ServerName
* @param toolName - The full tool name
* @returns The server name or undefined if not an MCP tool
*/
export function getServerNameFromTool(toolName: string): string | undefined {
const parts = toolName.split(Constants.mcp_delimiter);
if (parts.length >= 2) {
return parts[parts.length - 1];
}
return undefined;
}
/**
* Checks if a tool matches a set of patterns (tool names or server patterns).
* Supports both exact tool name matches and server-wide patterns like `mcp_all_mcp_ServerName`.
*
* @param toolName - The tool name to check
* @param patterns - Set of patterns (tool names or mcp_all_mcp_ServerName patterns)
* @param excludes - Set of tool names to exclude (takes precedence over patterns)
* @returns Whether the tool matches any pattern and is not excluded
*/
export function toolMatchesPatterns(
toolName: string,
patterns: Set<string>,
excludes: Set<string>,
): boolean {
if (excludes.has(toolName)) {
return false;
}
if (patterns.has(toolName)) {
return true;
}
const serverName = getServerNameFromTool(toolName);
if (serverName) {
const serverPattern = `${MCP_ALL_PATTERN}${serverName}`;
if (patterns.has(serverPattern)) {
return true;
}
}
return false;
}
/**
* Builds a tool registry from environment variables for the given tools.
* This is a temporary solution while UI-based configuration is being developed.
*
* Supports server-wide patterns using `mcp_all_mcp_ServerName` syntax.
* Exclusion env vars take precedence over inclusion patterns.
*
* Default behavior (if tool not listed in any env var):
* - allowed_callers: ['direct']
* - defer_loading: false
*
* @param tools - Array of tool definitions
* @returns Map of tool name to tool definition with classification
*
* @example
* // Environment for server-wide configuration:
* // TOOL_PROGRAMMATIC_ONLY=mcp_all_mcp_Google-Workspace
* // TOOL_DEFERRED=mcp_all_mcp_Google-Workspace
* // TOOL_DEFERRED_EXCLUDE=list_spreadsheets_mcp_Google-Workspace,read_sheet_values_mcp_Google-Workspace
*
* @example
* // Environment for individual tools:
* // TOOL_PROGRAMMATIC_ONLY=get_expenses,get_team_members
* // TOOL_DUAL_CONTEXT=get_weather
* // TOOL_DEFERRED=generate_report
*/
export function buildToolRegistryFromEnv(tools: ToolDefinition[]): LCToolRegistry {
const programmaticOnly = parseToolList(process.env.TOOL_PROGRAMMATIC_ONLY);
const programmaticOnlyExclude = parseToolList(process.env.TOOL_PROGRAMMATIC_ONLY_EXCLUDE);
const dualContext = parseToolList(process.env.TOOL_DUAL_CONTEXT);
const dualContextExclude = parseToolList(process.env.TOOL_DUAL_CONTEXT_EXCLUDE);
const deferred = parseToolList(process.env.TOOL_DEFERRED);
const deferredExclude = parseToolList(process.env.TOOL_DEFERRED_EXCLUDE);
const registry: LCToolRegistry = new Map();
for (const tool of tools) {
const { name, description, parameters } = tool;
let allowed_callers: AllowedCaller[];
if (toolMatchesPatterns(name, programmaticOnly, programmaticOnlyExclude)) {
allowed_callers = ['code_execution'];
} else if (toolMatchesPatterns(name, dualContext, dualContextExclude)) {
allowed_callers = ['direct', 'code_execution'];
} else {
// Default: direct only (LLM can call, PTC cannot)
allowed_callers = ['direct'];
}
const toolDef: LCTool = {
name,
allowed_callers,
defer_loading: toolMatchesPatterns(name, deferred, deferredExclude),
};
// Include description and parameters if available (needed for tool search and PTC stub generation)
if (description) {
toolDef.description = description;
}
if (parameters) {
toolDef.parameters = parameters;
}
registry.set(name, toolDef);
}
return registry;
}
/**
* Builds a tool registry from agent-level tool_options.
* This takes precedence over environment variable configuration when provided.
*
* @param tools - Array of tool definitions
* @param agentToolOptions - Per-tool configuration from the agent
* @returns Map of tool name to tool definition with classification
*/
export function buildToolRegistryFromAgentOptions(
tools: ToolDefinition[],
agentToolOptions: AgentToolOptions,
): LCToolRegistry {
/** Fall back to env vars for tools not configured at agent level */
const programmaticOnly = parseToolList(process.env.TOOL_PROGRAMMATIC_ONLY);
const programmaticOnlyExclude = parseToolList(process.env.TOOL_PROGRAMMATIC_ONLY_EXCLUDE);
const dualContext = parseToolList(process.env.TOOL_DUAL_CONTEXT);
const dualContextExclude = parseToolList(process.env.TOOL_DUAL_CONTEXT_EXCLUDE);
const registry: LCToolRegistry = new Map();
for (const tool of tools) {
const { name, description, parameters } = tool;
const agentOptions = agentToolOptions[name];
/** Determine allowed_callers: agent options take precedence, then env vars, then default */
let allowed_callers: AllowedCaller[];
if (agentOptions?.allowed_callers && agentOptions.allowed_callers.length > 0) {
allowed_callers = agentOptions.allowed_callers;
} else if (toolMatchesPatterns(name, programmaticOnly, programmaticOnlyExclude)) {
allowed_callers = ['code_execution'];
} else if (toolMatchesPatterns(name, dualContext, dualContextExclude)) {
allowed_callers = ['direct', 'code_execution'];
} else {
allowed_callers = ['direct'];
}
/** Determine defer_loading: agent options take precedence (explicit true/false) */
const defer_loading = agentOptions?.defer_loading === true;
const toolDef: LCTool = {
name,
allowed_callers,
defer_loading,
};
if (description) {
toolDef.description = description;
}
if (parameters) {
toolDef.parameters = parameters;
}
registry.set(name, toolDef);
}
return registry;
}
/**
* Checks if PTC (Programmatic Tool Calling) should be enabled based on environment configuration.
* PTC is enabled if any tools or server patterns are configured for programmatic calling.
* @returns Whether PTC should be enabled
*/
export function shouldEnablePTC(): boolean {
const programmaticOnly = parseToolList(process.env.TOOL_PROGRAMMATIC_ONLY);
const dualContext = parseToolList(process.env.TOOL_DUAL_CONTEXT);
return programmaticOnly.size > 0 || dualContext.size > 0;
}
/**
* Checks if tool search should be enabled based on environment configuration.
* Tool search is enabled if any tools or server patterns are configured as deferred.
* @returns Whether tool search should be enabled
*/
export function shouldEnableToolSearch(): boolean {
const deferred = parseToolList(process.env.TOOL_DEFERRED);
return deferred.size > 0;
}
interface MCPToolInstance {
name: string;
description?: string;
mcp?: boolean;
/** Original JSON schema attached at MCP tool creation time */
mcpJsonSchema?: JsonSchemaType;
}
/**
* Extracts MCP tool definition from a loaded tool instance.
* MCP tools have the original JSON schema attached as `mcpJsonSchema` property.
*
* @param tool - The loaded tool instance
* @returns Tool definition
*/
export function extractMCPToolDefinition(tool: MCPToolInstance): ToolDefinition {
const def: ToolDefinition = { name: tool.name };
if (tool.description) {
def.description = tool.description;
}
if (tool.mcpJsonSchema) {
def.parameters = tool.mcpJsonSchema;
}
return def;
}
/**
* Checks if a tool is an MCP tool based on its properties.
* @param tool - The tool to check (can be any object with potential mcp property)
* @returns Whether the tool is an MCP tool
*/
export function isMCPTool(tool: unknown): tool is MCPToolInstance {
return typeof tool === 'object' && tool !== null && (tool as MCPToolInstance).mcp === true;
}
/**
* Cleans up the temporary mcpJsonSchema property from MCP tools after registry is populated.
* This property is only needed during registry building and can be safely removed afterward.
*
* @param tools - Array of tools to clean up
*/
export function cleanupMCPToolSchemas(tools: MCPToolInstance[]): void {
for (const tool of tools) {
if (tool.mcpJsonSchema !== undefined) {
delete tool.mcpJsonSchema;
}
}
}
/** Parameters for building tool classification and creating PTC/tool search tools */
export interface BuildToolClassificationParams {
/** All loaded tools (will be filtered for MCP tools) */
loadedTools: GenericTool[];
/** User ID for auth lookup */
userId: string;
/** Agent ID (used to check if this agent should have classification features) */
agentId?: string;
/** Per-tool configuration from the agent (takes precedence over env vars) */
agentToolOptions?: AgentToolOptions;
/** Function to load auth values (dependency injection) */
loadAuthValues: (params: {
userId: string;
authFields: string[];
}) => Promise<Record<string, string>>;
}
/** Result from building tool classification */
export interface BuildToolClassificationResult {
/** Tool registry built from MCP tools (undefined if no MCP tools) */
toolRegistry?: LCToolRegistry;
/** Additional tools created (PTC and/or tool search) */
additionalTools: GenericTool[];
/** Whether any tools have defer_loading enabled (precomputed for efficiency) */
hasDeferredTools: boolean;
}
/**
* Checks if an agent is allowed to have classification features based on TOOL_CLASSIFICATION_AGENT_IDS.
* If TOOL_CLASSIFICATION_AGENT_IDS is not set, all agents are allowed (including when no agentId).
* If set, requires agentId to be in the list.
* @param agentId - The agent ID to check
* @returns Whether the agent is allowed
*/
export function isAgentAllowedForClassification(agentId?: string): boolean {
const allowedAgentIds = parseToolList(process.env.TOOL_CLASSIFICATION_AGENT_IDS);
if (allowedAgentIds.size === 0) {
return true;
}
if (!agentId) {
return false;
}
return allowedAgentIds.has(agentId);
}
/**
* Checks if an agent's tools have any that match PTC patterns (programmatic only or dual context).
* @param toolRegistry - The tool registry to check
* @returns Whether any tools are configured for programmatic calling
*/
export function agentHasProgrammaticTools(toolRegistry: LCToolRegistry): boolean {
for (const toolDef of toolRegistry.values()) {
if (toolDef.allowed_callers?.includes('code_execution')) {
return true;
}
}
return false;
}
/**
* Checks if an agent's tools have any that are deferred.
* @param toolRegistry - The tool registry to check
* @returns Whether any tools are configured as deferred
*/
export function agentHasDeferredTools(toolRegistry: LCToolRegistry): boolean {
for (const toolDef of toolRegistry.values()) {
if (toolDef.defer_loading === true) {
return true;
}
}
return false;
}
/**
* Builds the tool registry from MCP tools and conditionally creates PTC and tool search tools.
*
* This function:
* 1. Checks if the agent is allowed for classification features (via TOOL_CLASSIFICATION_AGENT_IDS)
* 2. Filters loaded tools for MCP tools
* 3. Extracts tool definitions and builds the registry
* - Uses agent's tool_options if provided (UI-based configuration)
* - Falls back to env vars for tools not configured at agent level
* 4. Cleans up temporary mcpJsonSchema properties
* 5. Creates PTC tool only if agent has tools configured for programmatic calling
* 6. Creates tool search tool only if agent has deferred tools
*
* @param params - Parameters including loaded tools, userId, agentId, agentToolOptions, and dependencies
* @returns Tool registry and any additional tools created
*/
export async function buildToolClassification(
params: BuildToolClassificationParams,
): Promise<BuildToolClassificationResult> {
const { loadedTools, userId, agentId, agentToolOptions, loadAuthValues } = params;
const additionalTools: GenericTool[] = [];
/** Check if this agent is allowed to have classification features (requires agentId) */
if (!isAgentAllowedForClassification(agentId)) {
logger.debug(
`[buildToolClassification] Agent ${agentId ?? 'undefined'} not allowed for classification, skipping`,
);
return { toolRegistry: undefined, additionalTools, hasDeferredTools: false };
}
const mcpTools = loadedTools.filter(isMCPTool);
if (mcpTools.length === 0) {
return { toolRegistry: undefined, additionalTools, hasDeferredTools: false };
}
const mcpToolDefs = mcpTools.map(extractMCPToolDefinition);
/**
* Build registry from agent's tool_options if provided (UI config).
* Environment variable-based classification is only used as fallback
* when TOOL_CLASSIFICATION_FROM_ENV=true is explicitly set.
*/
let toolRegistry: LCToolRegistry | undefined;
if (agentToolOptions && Object.keys(agentToolOptions).length > 0) {
toolRegistry = buildToolRegistryFromAgentOptions(mcpToolDefs, agentToolOptions);
} else if (process.env.TOOL_CLASSIFICATION_FROM_ENV === 'true') {
toolRegistry = buildToolRegistryFromEnv(mcpToolDefs);
} else {
/** No agent-level config and env-based classification not enabled */
return { toolRegistry: undefined, additionalTools, hasDeferredTools: false };
}
/** Clean up temporary mcpJsonSchema property from tools now that registry is populated */
cleanupMCPToolSchemas(mcpTools);
/**
* Check if this agent actually has tools that match the patterns.
* Only enable PTC if the agent has programmatic tools.
* Only enable tool search if the agent has deferred tools.
*/
const hasProgrammaticTools = agentHasProgrammaticTools(toolRegistry);
const hasDeferredTools = agentHasDeferredTools(toolRegistry);
if (!hasProgrammaticTools && !hasDeferredTools) {
logger.debug(
`[buildToolClassification] Agent ${agentId} has no programmatic or deferred tools, skipping PTC/ToolSearch`,
);
return { toolRegistry, additionalTools, hasDeferredTools: false };
}
/** Tool search uses local mode (no API key needed) */
if (hasDeferredTools) {
const toolSearchTool = createToolSearch({
mode: 'local',
toolRegistry,
});
additionalTools.push(toolSearchTool);
logger.debug(`[buildToolClassification] Tool Search enabled for agent ${agentId}`);
}
/** PTC requires CODE_API_KEY for sandbox execution */
if (hasProgrammaticTools) {
try {
const authValues = await loadAuthValues({
userId,
authFields: [EnvVar.CODE_API_KEY],
});
const codeApiKey = authValues[EnvVar.CODE_API_KEY];
if (!codeApiKey) {
logger.warn('[buildToolClassification] PTC configured but CODE_API_KEY not available');
} else {
const ptcTool = createProgrammaticToolCallingTool({ apiKey: codeApiKey });
additionalTools.push(ptcTool);
logger.debug(`[buildToolClassification] PTC tool enabled for agent ${agentId}`);
}
} catch (error) {
logger.error('[buildToolClassification] Error creating PTC tool:', error);
}
}
return { toolRegistry, additionalTools, hasDeferredTools };
}

View file

@ -1,2 +1,3 @@
export * from './format';
export * from './toolkits';
export * from './classification';