diff --git a/README.md b/README.md index f01e7687..a269df28 100644 --- a/README.md +++ b/README.md @@ -142,14 +142,17 @@ flowchart LR H --> B ``` -1MCP runs as an aggregated runtime behind `1mcp serve`. Static servers are prepared from startup configuration, template servers are materialized when client context is known, and the runtime can use async loading and lazy loading to reduce startup blocking and tool-surface noise. Instruction aggregation, presets, and notifications sit alongside that runtime rather than outside it. +1MCP runs as an aggregated runtime behind `1mcp serve`. Static servers are prepared from startup configuration, template servers are materialized when client context is known, and async loading can reduce startup blocking. Instruction aggregation, presets, and notifications sit alongside that runtime rather than outside it. + +Lazy loading is an opt-in stable tool-surface compatibility mode. It keeps the backend discovery and invocation surface at `tool_list`, `tool_schema`, and `tool_invoke` so capable agents can discover tools progressively without replacing their MCP tool table. Any explicitly enabled internal management tools remain directly exposed. Lazy loading reduces the initial schema payload, but it does not reduce backend connections or processes, make synchronous startup bind earlier, or repair orphaned proxy processes. See [#392](https://github.com/1mcp-app/agent/issues/392) for the async late-server visibility contract. ## Core Capabilities - Unified runtime for many MCP servers behind one `serve` process - CLI mode for progressive discovery with `1mcp instructions`, `1mcp inspect `, `1mcp inspect /`, and `1mcp run / --args ''` - Template servers for per-client or per-session resolution -- Async loading and lazy loading for faster startup and narrower exposure +- Async loading for earlier HTTP availability while static servers load +- Opt-in lazy loading for a stable progressive-discovery tool surface and smaller initial schemas - Opt-in automatic recovery for owned stdio backends, with health/status visibility and operator restart controls - Instruction aggregation across static and template-backed servers - Presets, filters, and preset change notifications diff --git a/config.toml.example b/config.toml.example index 31b5f45a..1ea01409 100644 --- a/config.toml.example +++ b/config.toml.example @@ -57,12 +57,10 @@ timeout = 30000 # batchDelay = 100 [lazyLoading] -# Enable lazy loading for tools +# Opt in to stable meta-tool exposure for progressive discovery. When false +# (the default), all tools are exposed directly. enabled = false -# Lazy loading mode: "full" | "metatool" | "hybrid" -mode = "full" - # Maximum tool schemas to cache # cacheMaxEntries = 1000 diff --git a/docs/en/commands/serve.md b/docs/en/commands/serve.md index bc8ccc85..4de0966c 100644 --- a/docs/en/commands/serve.md +++ b/docs/en/commands/serve.md @@ -77,7 +77,7 @@ For runtime-wide configuration details, see the **[Configuration Guide](/guide/e ### Runtime behavior - **`--enable-async-loading`**: Start HTTP availability before all static servers finish loading. -- **`--enable-lazy-loading`**: Enable lazy loading behavior for exposed server capabilities. +- **`--enable-lazy-loading`**: Opt into meta-tool exposure for progressive tool discovery. Omit it for full direct exposure. - **`--enable-config-reload`**: Enable config reload handling. - **`--enable-session-persistence`**: Enable HTTP session persistence. @@ -234,10 +234,10 @@ Behavior: 1mcp serve --config-dir ./config ``` -### Start with async and lazy loading +### Start with lazy loading ```bash -1mcp serve --enable-async-loading --enable-lazy-loading +1mcp serve --enable-lazy-loading ``` ### Start with filtered server exposure diff --git a/docs/en/guide/essentials/configuration.md b/docs/en/guide/essentials/configuration.md index a2eb298b..914bc517 100644 --- a/docs/en/guide/essentials/configuration.md +++ b/docs/en/guide/essentials/configuration.md @@ -356,6 +356,27 @@ ONE_MCP_PAGINATION=true \ npx -y @1mcp/agent ``` +### Lazy Loading + +Lazy loading is opt-in stable tool-surface compatibility for clients that support progressive discovery. With it disabled (the default), 1MCP exposes every backend tool directly. With it enabled, the backend tool surface is `tool_list`, `tool_schema`, and `tool_invoke`; clients discover schemas and invoke backend tools through those meta-tools. Internal management tools explicitly enabled by the operator remain directly exposed with the `1mcp` namespace. + +This reduces the initial tool-schema payload. It does not reduce backend connections or processes, make synchronous startup bind earlier, or repair orphaned proxy processes. Do not enable it together with async loading by default: choose each setting for its own runtime behavior. See [#392](https://github.com/1mcp-app/agent/issues/392) for the request-time visibility contract that lets a matching late-ready server appear after async capability publication without reconnecting. + +Enable it for one launch: + +```bash +1mcp serve --enable-lazy-loading +``` + +Or make it persistent in `config.toml` and restart the process: + +```toml +[lazyLoading] +enabled = true +``` + +The CLI flag takes precedence over `config.toml`. Legacy `--lazy-mode`, `[lazyLoading].mode`, and `--lazy-direct-expose` inputs are accepted temporarily but ignored with a deprecation warning; `enabled` is the only working switch. + ### Configuration Reload Control configuration file hot-reload behavior for seamless updates. diff --git a/docs/zh/commands/serve.md b/docs/zh/commands/serve.md index 7f141f90..e15d840d 100644 --- a/docs/zh/commands/serve.md +++ b/docs/zh/commands/serve.md @@ -67,7 +67,7 @@ CLI 模式依赖一个正在运行的 `serve` 实例。 ### 运行时行为 - **`--enable-async-loading`**：让 HTTP 可用性先启动，再等待静态服务器完成加载。 -- **`--enable-lazy-loading`**：为服务能力启用懒加载行为。 +- **`--enable-lazy-loading`**：选择使用元工具逐步发现工具；省略该选项则直接暴露全部工具。 - **`--enable-config-reload`**：启用配置重载处理。 - **`--enable-session-persistence`**：启用 HTTP 会话持久化。 @@ -224,10 +224,10 @@ Stopped supervised background runtime in Runtime Scope /home/me/.config/1mcp (su 1mcp serve --config-dir ./config ``` -### 启用异步加载与懒加载 +### 启用懒加载 ```bash -1mcp serve --enable-async-loading --enable-lazy-loading +1mcp serve --enable-lazy-loading ``` ### 启动时筛选服务器暴露面 diff --git a/docs/zh/guide/essentials/configuration.md b/docs/zh/guide/essentials/configuration.md index ac453c6e..e3bac259 100644 --- a/docs/zh/guide/essentials/configuration.md +++ b/docs/zh/guide/essentials/configuration.md @@ -350,6 +350,27 @@ ONE_MCP_PAGINATION=true \ npx -y @1mcp/agent ``` +### 懒加载 + +懒加载是一种可选的稳定工具表面兼容模式，适用于支持逐步发现工具的客户端。默认关闭时，1MCP 会直接暴露全部后端工具；启用后，后端工具表面为 `tool_list`、`tool_schema` 和 `tool_invoke`，客户端通过这些元工具发现 Schema 并调用后端工具。由运维人员显式启用的内部管理工具仍会以 `1mcp` 命名空间直接暴露。 + +它可以减少初始工具 Schema 的负载，但不会减少后端连接或进程、让同步启动更早绑定端口，也不会修复遗留的代理进程。不要默认将它与异步加载同时启用；应根据各自的运行时行为分别选择。异步能力发布后，符合条件的后就绪服务器无需重连即可在下一次请求中可见；该请求时可见性契约见 [#392](https://github.com/1mcp-app/agent/issues/392)。 + +仅为本次启动启用： + +```bash +1mcp serve --enable-lazy-loading +``` + +也可以在 `config.toml` 中持久启用，并重启进程： + +```toml +[lazyLoading] +enabled = true +``` + +CLI 参数优先于 `config.toml`。旧的 `--lazy-mode`、`[lazyLoading].mode` 和 `--lazy-direct-expose` 输入会在一个兼容周期内继续接受，但将被忽略并产生弃用警告；`enabled` 是唯一有效的开关。 + ### 配置重载 控制配置文件热重载行为以实现无缝更新。 diff --git a/src/commands/serve/index.ts b/src/commands/serve/index.ts index 697197d9..1c0309e8 100644 --- a/src/commands/serve/index.ts +++ b/src/commands/serve/index.ts @@ -163,7 +163,7 @@ export const serverOptions = { type: 'boolean' as const, }, 'lazy-mode': { - describe: 'Lazy loading mode (metatool: 3 meta-tools only, hybrid: direct + meta-tools, full: disabled)', + describe: 'Deprecated and ignored. Use --enable-lazy-loading to expose meta-tools, or omit it for full direct exposure.', type: 'string' as const, choices: ['metatool', 'hybrid', 'full'] as const, }, @@ -192,7 +192,7 @@ export const serverOptions = { default: 'grouped', }, 'lazy-direct-expose': { - describe: 'Comma-separated glob patterns for tools to expose directly in hybrid mode', + describe: 'Deprecated and ignored. Lazy loading does not support direct tool exposure.', type: 'string' as const, default: undefined, }, @@ -338,15 +338,11 @@ INTERNAL TOOLS: • Examples: "safe" (read-only), "discovery,management" (no installation) LAZY LOADING: - Use --enable-lazy-loading to enable lazy tool loading for token optimization. - Lazy loading reduces initial token usage by ~95% by loading tools on-demand. - Modes: - • metatool: Expose only 3 meta-tools (list/describe/call) - maximum savings - • hybrid: Common tools direct + meta-tools for rest - balanced approach - • full: Disabled, load all tools normally (default, backward compatible) + Use --enable-lazy-loading to opt into stable meta-tool exposure for progressive discovery. + Omit it (the default) for full direct tool exposure. Legacy --lazy-mode and + --lazy-direct-expose inputs are accepted temporarily but ignored with a warning. Examples: - • --enable-lazy-loading --lazy-mode=metatool - • --enable-lazy-loading --lazy-mode=hybrid --lazy-preload=filesystem,search + • --enable-lazy-loading • --enable-lazy-loading --lazy-cache-max-entries=2000 For more information: https://github.com/1mcp-app/agent diff --git a/src/commands/serve/serve.stdio-filter.test.ts b/src/commands/serve/serve.stdio-filter.test.ts index 38c98d1a..31d551d4 100644 --- a/src/commands/serve/serve.stdio-filter.test.ts +++ b/src/commands/serve/serve.stdio-filter.test.ts @@ -102,7 +102,6 @@ function makeStdioOptions(overrides: Partial = {}): ServeOptions { 'async-batch-delay': 100, 'async-notify-on-ready': true, 'enable-lazy-loading': false, - 'lazy-mode': 'full', 'lazy-inline-catalog': false, 'lazy-catalog-format': 'grouped', 'lazy-cache-max-entries': 1000, diff --git a/src/commands/serve/serve.test.ts b/src/commands/serve/serve.test.ts index 62b0f03b..39c15285 100644 --- a/src/commands/serve/serve.test.ts +++ b/src/commands/serve/serve.test.ts @@ -139,7 +139,6 @@ describe('serveCommand - config-dir session isolation', () => { 'async-batch-delay': 100, 'async-notify-on-ready': true, 'enable-lazy-loading': false, - 'lazy-mode': 'full', 'lazy-inline-catalog': false, 'lazy-catalog-format': 'grouped', 'lazy-cache-max-entries': 1000, @@ -208,7 +207,6 @@ describe('serveCommand - config-dir session isolation', () => { 'async-batch-delay': 100, 'async-notify-on-ready': true, 'enable-lazy-loading': false, - 'lazy-mode': 'full', 'lazy-inline-catalog': false, 'lazy-catalog-format': 'grouped', 'lazy-cache-max-entries': 1000, @@ -272,7 +270,6 @@ describe('serveCommand - config-dir session isolation', () => { 'async-batch-delay': 100, 'async-notify-on-ready': true, 'enable-lazy-loading': false, - 'lazy-mode': 'full', 'lazy-inline-catalog': false, 'lazy-catalog-format': 'grouped', 'lazy-cache-max-entries': 1000, diff --git a/src/commands/serve/serve.transport-precedence.test.ts b/src/commands/serve/serve.transport-precedence.test.ts index cc28d461..75179012 100644 --- a/src/commands/serve/serve.transport-precedence.test.ts +++ b/src/commands/serve/serve.transport-precedence.test.ts @@ -1,5 +1,6 @@ import { AgentConfigManager } from '@src/core/server/agentConfig.js'; import { configureGlobalLogger } from '@src/logger/configureGlobalLogger.js'; +import logger from '@src/logger/logger.js'; import { displayLogo } from '@src/utils/ui/logo.js'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; @@ -110,7 +111,6 @@ describe('serveCommand - config-dir session isolation', () => { 'async-batch-delay': 100, 'async-notify-on-ready': true, 'enable-lazy-loading': false, - 'lazy-mode': 'full', 'lazy-inline-catalog': false, 'lazy-catalog-format': 'grouped', 'lazy-cache-max-entries': 1000, @@ -196,6 +196,30 @@ describe('serveCommand - config-dir session isolation', () => { ); }); + it('warns for legacy lazy selectors without changing the enabled switch', async () => { + const configManager = AgentConfigManager.getInstance(); + const updateConfigSpy = vi.mocked(configManager.updateConfig); + + try { + await serveCommand({ + ...baseOptions, + transport: 'http', + port: 3050, + host: '127.0.0.1', + 'lazy-mode': 'hybrid', + 'lazy-direct-expose': 'filesystem_*', + }); + } catch { + // Ignore errors from mocked dependencies. + } + + expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('--lazy-mode is ignored')); + expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('--lazy-direct-expose is ignored')); + expect(updateConfigSpy).toHaveBeenCalledWith( + expect.objectContaining({ lazyLoading: expect.objectContaining({ enabled: false }) }), + ); + }); + it('uses config.toml app settings when matching CLI flags are omitted', async () => { const { ConfigManager } = await import('@src/config/configManager.js'); vi.mocked(ConfigManager.getInstance).mockReturnValue({ diff --git a/src/commands/serve/serve.ts b/src/commands/serve/serve.ts index 726cb268..61696edd 100644 --- a/src/commands/serve/serve.ts +++ b/src/commands/serve/serve.ts @@ -174,6 +174,22 @@ function loadInstructionsTemplate(templatePath?: string, configDir?: string): st } } +function warnForLegacyLazyLoadingOptions(parsedArgv: ServeOptions): void { + if (parsedArgv['lazy-mode'] !== undefined) { + logger.warn( + 'DEPRECATION WARNING: --lazy-mode is ignored. Lazy loading is controlled only by --enable-lazy-loading; ' + + 'for a persistent setting, use [lazyLoading] enabled = true in config.toml. Remove --lazy-mode because it does not change runtime behavior.', + ); + } + + if (parsedArgv['lazy-direct-expose'] !== undefined) { + logger.warn( + 'DEPRECATION WARNING: --lazy-direct-expose is ignored. Lazy loading exposes only meta-tools when enabled; ' + + 'for a persistent setting, use [lazyLoading] enabled = true in config.toml. Remove --lazy-direct-expose because it does not change runtime behavior.', + ); + } +} + /** * Set up graceful shutdown handling */ @@ -436,6 +452,7 @@ export async function serveCommand(parsedArgv: ServeOptions): Promise { } const internalToolsList = parseInternalToolsList(parsedArgv['internal-tools']); + warnForLegacyLazyLoadingOptions(parsedArgv); const directExpose = parseCommaSeparatedList(parsedArgv['lazy-direct-expose']); const preloadPatterns = parseCommaSeparatedList(parsedArgv['lazy-preload']); const preloadKeywords = parseCommaSeparatedList(parsedArgv['lazy-preload-keywords']); diff --git a/src/config/configLoader.ts b/src/config/configLoader.ts index 3c3931e0..c9de28fe 100644 --- a/src/config/configLoader.ts +++ b/src/config/configLoader.ts @@ -86,6 +86,22 @@ function warnIfLegacyAppConfig(rawConfig: unknown, configFilePath: string): void ); } +function warnForLegacyLazyLoadingConfig(rawConfig: unknown, tomlPath: string): void { + if (!rawConfig || typeof rawConfig !== 'object') { + return; + } + + const lazyLoading = (rawConfig as Record).lazyLoading; + if (!lazyLoading || typeof lazyLoading !== 'object' || !('mode' in lazyLoading)) { + return; + } + + logger.warn( + `The [lazyLoading].mode setting in ${tomlPath} is deprecated and ignored. ` + + 'Lazy loading is controlled only by [lazyLoading] enabled = true. Remove mode because it does not change runtime behavior.', + ); +} + function warnForUnknownGlobalConfigKeys(rawGlobal: unknown): void { const unknownGlobalKeys = getUnknownGlobalConfigKeys(rawGlobal); if (unknownGlobalKeys.length === 0) { @@ -118,6 +134,7 @@ export function loadAppConfigFromTomlPath(tomlPath: string): ApplicationConfig { } const raw = fs.readFileSync(tomlPath, 'utf8'); const parsed = parseToml(raw); + warnForLegacyLazyLoadingConfig(parsed, tomlPath); return applicationConfigSchema.parse(parsed); } catch (error) { if (error instanceof ZodError) { diff --git a/src/config/configLoader.validation-app-config.test.ts b/src/config/configLoader.validation-app-config.test.ts index 05116679..f622bda4 100644 --- a/src/config/configLoader.validation-app-config.test.ts +++ b/src/config/configLoader.validation-app-config.test.ts @@ -4,6 +4,7 @@ import { tmpdir } from 'os'; import { join } from 'path'; import { ConfigLoader } from '@src/config/configLoader.js'; +import logger from '@src/logger/logger.js'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; @@ -63,6 +64,7 @@ describe('ConfigLoader', () => { // Ignore cleanup errors } vi.clearAllMocks(); + vi.restoreAllMocks(); resetMockAgentConfig(); }); @@ -417,6 +419,30 @@ minServers = 2 expect(result.asyncLoading?.minServers).toBe(2); }); + it('accepts the legacy lazy mode while warning that enabled controls runtime behavior', async () => { + const warnSpy = vi.spyOn(logger, 'warn').mockImplementation(() => logger); + await fsPromises.writeFile( + join(tempConfigDir, 'config.toml'), + '[lazyLoading]\nenabled = false\nmode = "hybrid"\n', + ); + + const result = loader.loadAppConfigFromToml(); + + expect(result.lazyLoading).toEqual({ enabled: false, mode: 'hybrid' }); + expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('[lazyLoading].mode')); + expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('enabled = true')); + }); + + it('does not warn when lazy loading uses only the enabled switch', async () => { + const warnSpy = vi.spyOn(logger, 'warn').mockImplementation(() => logger); + await fsPromises.writeFile(join(tempConfigDir, 'config.toml'), '[lazyLoading]\nenabled = true\n'); + + const result = loader.loadAppConfigFromToml(); + + expect(result.lazyLoading).toEqual({ enabled: true }); + expect(warnSpy).not.toHaveBeenCalled(); + }); + it('should return empty object and warn on invalid TOML', async () => { await fsPromises.writeFile(join(tempConfigDir, 'config.toml'), 'invalid = [toml'); diff --git a/src/core/capabilities/capabilityCatalog.test.ts b/src/core/capabilities/capabilityCatalog.test.ts index f3aa6f16..af6afba3 100644 --- a/src/core/capabilities/capabilityCatalog.test.ts +++ b/src/core/capabilities/capabilityCatalog.test.ts @@ -4,6 +4,7 @@ import type { TemplateHashProvider } from '@src/core/server/connectionResolver.j import { ClientStatus, type OutboundConnections } from '@src/core/types/client.js'; import { CapabilityCatalog } from './capabilityCatalog.js'; +import { capabilityVisibilityFromServerNames, createCapabilityVisibility } from './capabilityVisibility.js'; import { SchemaCache } from './schemaCache.js'; import { ToolRegistry } from './toolRegistry.js'; @@ -112,7 +113,7 @@ describe('CapabilityCatalog', () => { getAllRenderedHashesForSession: () => undefined, }).invokeVisibleTool( { server: 'template-server', toolName: 'template_tool', args: { message: 'hi' } }, - 'session-1', + createCapabilityVisibility([['template-server', 'template-server']], 'session-1'), ); expect(result.error).toBeUndefined(); @@ -130,7 +131,7 @@ describe('CapabilityCatalog', () => { getAllRenderedHashesForSession: () => undefined, }).invokeVisibleTool( { server: 'template-server', toolName: 'template_tool', args: { message: 'hi' } }, - 'missing-session', + createCapabilityVisibility([['template-server', 'template-server']], 'missing-session'), ); expect(result.error).toMatchObject({ @@ -140,8 +141,8 @@ describe('CapabilityCatalog', () => { expect(mockClient.callTool).not.toHaveBeenCalled(); }); - it('filters visibility by allowed server set', async () => { - const result = await createCatalog().listVisibleTools({}, undefined, new Set(['filesystem'])); + it('filters capability visibility by Server Candidate Set', async () => { + const result = await createCatalog().listVisibleTools({}, capabilityVisibilityFromServerNames(['filesystem'])); expect(result.tools.map((tool) => tool.server)).toEqual(['filesystem']); expect(result.routes.map((route) => route.connectionKey)).toEqual(['filesystem']); @@ -150,8 +151,7 @@ describe('CapabilityCatalog', () => { it('does not reveal disabled tool details for hidden servers', async () => { const result = await createCatalog().describeVisibleTool( { server: 'filesystem', toolName: 'write_file' }, - undefined, - new Set(['template-server']), + capabilityVisibilityFromServerNames(['template-server']), ); expect(result.error).toMatchObject({ @@ -170,9 +170,11 @@ describe('CapabilityCatalog', () => { return { changed: true, shouldNotifyListChanged: true }; }); - const result = await createCatalog(undefined, { refreshCapabilities }).listVisibleTools({}, undefined, undefined, { - refreshIntent: 'force', - }); + const result = await createCatalog(undefined, { refreshCapabilities }).listVisibleTools( + {}, + undefined, + { refreshIntent: 'force' }, + ); expect(refreshCapabilities).toHaveBeenCalledWith({ intent: 'force', reason: 'list' }); expect(result.tools.map((tool) => `${tool.server}/${tool.name}`)).toEqual(['filesystem/read_file']); diff --git a/src/core/capabilities/capabilityCatalog.ts b/src/core/capabilities/capabilityCatalog.ts index e30749c3..205225d4 100644 --- a/src/core/capabilities/capabilityCatalog.ts +++ b/src/core/capabilities/capabilityCatalog.ts @@ -5,6 +5,10 @@ import { getDisabledToolError, isToolDisabled } from '@src/core/server/disabledT import type { MCPServerParams, OutboundConnections } from '@src/core/types/index.js'; import logger from '@src/logger/logger.js'; +import { + type CapabilityVisibility, + getCapabilityVisibleServerNames, +} from './capabilityVisibility.js'; import { SchemaCache } from './schemaCache.js'; import type { ListToolsOptions, ListToolsResult as RegistryListToolsResult, ToolMetadata } from './toolRegistry.js'; import { ToolRegistry } from './toolRegistry.js'; @@ -60,7 +64,7 @@ export interface CapabilityCatalogDependencies { getServerConfigs: () => Record; loadSchema?: (server: string, toolName: string) => Promise; refreshCapabilities?: (input: CapabilityRefreshInput) => Promise; - defaultAllowedServers?: Set; + defaultVisibility?: CapabilityVisibility; templateHashProvider?: TemplateHashProvider; } @@ -97,17 +101,16 @@ export class CapabilityCatalog { public async listVisibleTools( options: ListToolsOptions = {}, - sessionId?: string, - allowedServers?: Set, + visibility?: CapabilityVisibility, queryOptions: CapabilityCatalogQueryOptions = {}, ): Promise { const refresh = await this.resolveRefreshFacts(queryOptions.refreshIntent ?? 'never', 'list'); - const registry = this.visibleToolRegistry(allowedServers); + const registry = this.visibleToolRegistry(visibility); const result = registry.listTools(options); const tools = result.tools; const servers = Array.from(new Set(tools.map((tool) => tool.server))).sort(); const routes = tools - .map((tool) => this.resolveRoute(tool.server, tool.name, sessionId)) + .map((tool) => this.resolveRoute(tool, visibility)) .filter((route): route is CapabilityRoute => route !== undefined); return { @@ -121,12 +124,11 @@ export class CapabilityCatalog { public async describeVisibleTool( args: { server?: string; toolName?: string }, - sessionId?: string, - allowedServers?: Set, + visibility?: CapabilityVisibility, queryOptions: CapabilityCatalogQueryOptions = {}, ): Promise { const refresh = await this.resolveRefreshFacts(queryOptions.refreshIntent ?? 'never', 'describe'); - const access = this.resolveVisibleToolAccess(args, sessionId, allowedServers); + const access = this.resolveVisibleToolAccess(args, visibility); if (access.error) { return { schema: {}, error: access.error, refresh }; } @@ -167,12 +169,11 @@ export class CapabilityCatalog { public async invokeVisibleTool( args: { server?: string; toolName?: string; args: unknown }, - sessionId?: string, - allowedServers?: Set, + visibility?: CapabilityVisibility, queryOptions: CapabilityCatalogQueryOptions = {}, ): Promise { const refresh = await this.resolveRefreshFacts(queryOptions.refreshIntent ?? 'never', 'invoke'); - const access = this.resolveVisibleToolAccess(args, sessionId, allowedServers); + const access = this.resolveVisibleToolAccess(args, visibility); if (access.error) { return { result: {}, @@ -261,11 +262,11 @@ export class CapabilityCatalog { }; } - private visibleToolRegistry(allowedServers?: Set): ToolRegistry { + private visibleToolRegistry(visibility?: CapabilityVisibility): ToolRegistry { let registry = this.deps.getToolRegistry(); - const effectiveAllowedServers = allowedServers ?? this.deps.defaultAllowedServers; - if (effectiveAllowedServers !== undefined) { - registry = registry.filterByServers(effectiveAllowedServers); + const effectiveVisibility = visibility ?? this.deps.defaultVisibility; + if (effectiveVisibility !== undefined) { + registry = registry.filterByServerCandidates(effectiveVisibility.serverCandidates); } if (typeof registry.getAllTools !== 'function') { @@ -284,6 +285,7 @@ export class CapabilityCatalog { inputSchema: tool.inputSchema ?? { type: 'object' }, }, server: tool.server, + connectionKey: tool.connectionKey, tags: tool.tags, })), ); @@ -291,8 +293,7 @@ export class CapabilityCatalog { private resolveVisibleToolAccess( args: { server?: string; toolName?: string }, - sessionId?: string, - allowedServers?: Set, + visibility?: CapabilityVisibility, ): | { route: CapabilityRoute; tool: ToolMetadata; error?: never } | { route?: never; tool?: never; error: CapabilityAccessError } { @@ -305,7 +306,7 @@ export class CapabilityCatalog { }; } - const visibleRegistry = this.visibleToolRegistry(allowedServers); + const visibleRegistry = this.visibleToolRegistry(visibility); if (typeof visibleRegistry.getTool !== 'function') { return { error: { @@ -318,7 +319,7 @@ export class CapabilityCatalog { const tool = visibleRegistry.getTool(args.server, args.toolName); if (!tool) { - const disabledError = this.isServerVisible(args.server, allowedServers) + const disabledError = this.isServerVisible(args.server, visibility) ? getDisabledToolError(this.deps.getServerConfigs(), args.server, args.toolName) : undefined; return { @@ -329,7 +330,7 @@ export class CapabilityCatalog { }; } - const route = this.resolveRoute(args.server, args.toolName, sessionId); + const route = this.resolveRoute(tool, visibility); if (!route) { return { error: { @@ -342,21 +343,33 @@ export class CapabilityCatalog { return { route, tool }; } - private isServerVisible(server: string, allowedServers?: Set): boolean { - const effectiveAllowedServers = allowedServers ?? this.deps.defaultAllowedServers; - return effectiveAllowedServers === undefined || effectiveAllowedServers.has(server); + private isServerVisible(server: string, visibility?: CapabilityVisibility): boolean { + const effectiveVisibility = visibility ?? this.deps.defaultVisibility; + return ( + effectiveVisibility === undefined || getCapabilityVisibleServerNames(effectiveVisibility).has(server) + ); } - private resolveRoute(server: string, toolName: string, sessionId?: string): CapabilityRoute | undefined { - const sessionResult = sessionId ? this.connectionResolver.resolveWithKey(server, sessionId) : undefined; - const result = sessionResult ?? (!sessionId ? this.connectionResolver.findByServerName(server) : undefined); + private resolveRoute(tool: ToolMetadata, visibility?: CapabilityVisibility): CapabilityRoute | undefined { + const registryConnectionKey = tool.connectionKey ?? tool.server; + if (this.deps.outboundConnections.has(registryConnectionKey)) { + return { + server: tool.server, + toolName: tool.name, + connectionKey: registryConnectionKey, + }; + } + + const sessionId = visibility?.sessionId ?? this.deps.defaultVisibility?.sessionId; + const sessionResult = sessionId ? this.connectionResolver.resolveWithKey(tool.server, sessionId) : undefined; + const result = sessionResult ?? (!sessionId ? this.connectionResolver.findByServerName(tool.server) : undefined); if (!result) { return undefined; } return { - server, - toolName, + server: tool.server, + toolName: tool.name, connectionKey: result.key, }; } diff --git a/src/core/capabilities/capabilityVisibility.ts b/src/core/capabilities/capabilityVisibility.ts new file mode 100644 index 00000000..2f8ebb25 --- /dev/null +++ b/src/core/capabilities/capabilityVisibility.ts @@ -0,0 +1,26 @@ +/** + * Request-scoped capability visibility after session, tag, and preset filtering. + * Connection keys retain template identity; server names remain the public API. + */ +export interface CapabilityVisibility { + readonly sessionId?: string; + readonly serverCandidates: ReadonlyMap; +} + +export function createCapabilityVisibility( + serverCandidates: Iterable, + sessionId?: string, +): CapabilityVisibility { + return { + sessionId, + serverCandidates: new Map(serverCandidates), + }; +} + +export function capabilityVisibilityFromServerNames(serverNames: Iterable): CapabilityVisibility { + return createCapabilityVisibility(Array.from(serverNames, (serverName) => [serverName, serverName] as const)); +} + +export function getCapabilityVisibleServerNames(visibility: CapabilityVisibility): Set { + return new Set(visibility.serverCandidates.values()); +} diff --git a/src/core/capabilities/internalCapabilitiesProvider.ts b/src/core/capabilities/internalCapabilitiesProvider.ts index c50a31d9..3f67f7e4 100644 --- a/src/core/capabilities/internalCapabilitiesProvider.ts +++ b/src/core/capabilities/internalCapabilitiesProvider.ts @@ -34,6 +34,7 @@ import { createStatusTool, } from '@src/core/capabilities/internal/managementTools.js'; import { LazyLoadingOrchestrator } from '@src/core/capabilities/lazyLoadingOrchestrator.js'; +import type { CapabilityVisibility } from '@src/core/capabilities/capabilityVisibility.js'; import { FlagManager } from '@src/core/flags/flagManager.js'; import { AgentConfigManager } from '@src/core/server/agentConfig.js'; import { @@ -259,7 +260,11 @@ export class InternalCapabilitiesProvider extends EventEmitter { /** * Execute an internal tool */ - public async executeTool(toolName: string, args: unknown, sessionId?: string): Promise { + public async executeTool( + toolName: string, + args: unknown, + visibility?: CapabilityVisibility, + ): Promise { if (!this.isInitialized) { throw new Error('Internal capabilities provider not initialized'); } @@ -274,7 +279,7 @@ export class InternalCapabilitiesProvider extends EventEmitter { if (!this.lazyLoadingOrchestrator) { throw new Error('Lazy loading not available'); } - const result = await this.lazyLoadingOrchestrator.callMetaTool(toolName, args, sessionId); + const result = await this.lazyLoadingOrchestrator.callMetaTool(toolName, args, visibility); // Check for error field in structured result if (typeof result === 'object' && result !== null && 'error' in result && result.error) { diff --git a/src/core/capabilities/lazyLoadingOrchestrator.test.ts b/src/core/capabilities/lazyLoadingOrchestrator.test.ts index 727fa118..5ea3b10d 100644 --- a/src/core/capabilities/lazyLoadingOrchestrator.test.ts +++ b/src/core/capabilities/lazyLoadingOrchestrator.test.ts @@ -4,6 +4,7 @@ import { ClientStatus, OutboundConnections } from '@src/core/types/index.js'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { capabilityVisibilityFromServerNames } from './capabilityVisibility.js'; import { LazyLoadingOrchestrator } from './lazyLoadingOrchestrator.js'; describe('LazyLoadingOrchestrator', () => { @@ -199,7 +200,7 @@ describe('LazyLoadingOrchestrator', () => { expect(orchestrator.isEnabled()).toBe(false); }); - it('should initialize successfully in hybrid mode', async () => { + it('treats legacy hybrid configuration as enabled meta-tool exposure', async () => { mockAgentConfig.get.mockImplementation((key: string) => { if (key === 'lazyLoading') { return { @@ -219,7 +220,6 @@ describe('LazyLoadingOrchestrator', () => { await orchestrator.initialize(); - // getMode() removed - hybrid mode replaced with metatool mode expect(orchestrator.isEnabled()).toBe(true); }); @@ -304,7 +304,7 @@ describe('LazyLoadingOrchestrator', () => { expect(capabilities.timestamp).toBeDefined(); }); - it('should return meta-tools + direct exposed tools in hybrid mode', async () => { + it('ignores legacy direct exposure settings when lazy loading is enabled', async () => { mockAgentConfig.get.mockImplementation((key: string) => { if (key === 'lazyLoading') { return { @@ -325,8 +325,7 @@ describe('LazyLoadingOrchestrator', () => { const capabilities = await orchestrator.getCapabilities(); - // Should have 3 meta-tools + direct exposed tools - expect(capabilities.tools.length).toBeGreaterThanOrEqual(3); + expect(capabilities.tools).toHaveLength(3); expect(capabilities.tools.map((t) => t.name)).toContain('tool_list'); expect(capabilities.tools.map((t) => t.name)).toContain('tool_schema'); expect(capabilities.tools.map((t) => t.name)).toContain('tool_invoke'); @@ -390,7 +389,7 @@ describe('LazyLoadingOrchestrator', () => { expect(orchestrator.shouldNotifyListChanged()).toBe(true); }); - it('should return false in hybrid mode when no direct tools', async () => { + it('does not notify tool-list changes for legacy hybrid configuration', async () => { mockAgentConfig.get.mockImplementation((key: string) => { if (key === 'lazyLoading') { return { @@ -462,7 +461,11 @@ describe('LazyLoadingOrchestrator', () => { }); it('should call tool_list', async () => { - const result = await orchestrator.callMetaTool('tool_list', {}); + const result = await orchestrator.callMetaTool( + 'tool_list', + {}, + capabilityVisibilityFromServerNames(['filesystem', 'database']), + ); expect(result).toBeDefined(); expect((result as any).isError).toBeFalsy(); @@ -480,7 +483,7 @@ describe('LazyLoadingOrchestrator', () => { const result = await orchestrator.callMetaTool('tool_schema', { server: 'filesystem', toolName: 'read_file', - }); + }, capabilityVisibilityFromServerNames(['filesystem'])); expect(result).toBeDefined(); }); @@ -504,7 +507,9 @@ describe('LazyLoadingOrchestrator', () => { const fullOrchestrator = new LazyLoadingOrchestrator(mockOutboundConnections, mockAgentConfig); await fullOrchestrator.initialize(); - await expect(fullOrchestrator.callMetaTool('tool_list', {})).rejects.toThrow( + await expect( + fullOrchestrator.callMetaTool('tool_list', {}, capabilityVisibilityFromServerNames([])), + ).rejects.toThrow( 'Meta-tool provider not initialized', ); }); @@ -1038,27 +1043,8 @@ describe('LazyLoadingOrchestrator', () => { await orchestrator.initialize(); }); - it('should store and retrieve session-specific allowed servers', async () => { - const sessionId = 'session-123'; - const allowedServers = new Set(['filesystem', 'database']); - - // Initially no filter - expect(orchestrator.getSessionAllowedServers(sessionId)).toBeUndefined(); - - // Store filter via getCapabilitiesForFilteredServers - await orchestrator.getCapabilitiesForFilteredServers(allowedServers, sessionId); - - // Should retrieve the same filter - const retrieved = orchestrator.getSessionAllowedServers(sessionId); - expect(retrieved).toBeDefined(); - expect(retrieved?.size).toBe(2); - expect(retrieved?.has('filesystem')).toBe(true); - expect(retrieved?.has('database')).toBe(true); - }); - - it('should filter capabilities by session when calling getCapabilitiesForFilteredServers', async () => { - const sessionId = 'session-456'; - const allowedServers = new Set(['filesystem']); + it('should filter capabilities with the set resolved for the current request', async () => { + const visibility = capabilityVisibilityFromServerNames(['filesystem']); // Mock base capabilities with multiple servers const mockCapabilities = { @@ -1077,7 +1063,7 @@ describe('LazyLoadingOrchestrator', () => { vi.spyOn(orchestrator['capabilityAggregator'], 'getCurrentCapabilities').mockReturnValue(mockCapabilities); - const filteredCaps = await orchestrator.getCapabilitiesForFilteredServers(allowedServers, sessionId); + const filteredCaps = await orchestrator.getCapabilitiesForVisibility(visibility); // Should only include filesystem resources/prompts/servers expect(filteredCaps.resources.length).toBe(1); @@ -1096,84 +1082,58 @@ describe('LazyLoadingOrchestrator', () => { expect(filteredCaps.tools.map((t) => t.name)).toContain('tool_invoke'); }); - it('should clear session filter correctly', async () => { - const sessionId = 'session-789'; - const allowedServers = new Set(['filesystem']); - - // Set filter - await orchestrator.getCapabilitiesForFilteredServers(allowedServers, sessionId); - expect(orchestrator.getSessionAllowedServers(sessionId)).toBeDefined(); - - // Clear filter - orchestrator.clearSessionFilter(sessionId); - expect(orchestrator.getSessionAllowedServers(sessionId)).toBeUndefined(); - }); - - it('should isolate filters between different sessions', async () => { - const session1 = 'session-aaa'; - const session2 = 'session-bbb'; - const allowedServers1 = new Set(['filesystem']); - const allowedServers2 = new Set(['database']); - - // Set different filters for different sessions - await orchestrator.getCapabilitiesForFilteredServers(allowedServers1, session1); - await orchestrator.getCapabilitiesForFilteredServers(allowedServers2, session2); - - // Verify isolation - const filter1 = orchestrator.getSessionAllowedServers(session1); - const filter2 = orchestrator.getSessionAllowedServers(session2); - - expect(filter1?.has('filesystem')).toBe(true); - expect(filter1?.has('database')).toBe(false); - - expect(filter2?.has('database')).toBe(true); - expect(filter2?.has('filesystem')).toBe(false); - - // Clear one session shouldn't affect the other - orchestrator.clearSessionFilter(session1); - expect(orchestrator.getSessionAllowedServers(session1)).toBeUndefined(); - expect(orchestrator.getSessionAllowedServers(session2)).toBeDefined(); - }); - - it('should apply session filter when calling callMetaTool with sessionId', async () => { - const sessionId = 'session-ccc'; - const allowedServers = new Set(['filesystem']); - - // Set session filter - await orchestrator.getCapabilitiesForFilteredServers(allowedServers, sessionId); - - // Call tool_list meta-tool with sessionId - const result = await orchestrator.callMetaTool('tool_list', {}, sessionId); - - expect(result).toBeDefined(); - expect((result as any).isError).toBeFalsy(); + it('should apply the explicitly supplied filter to every meta-tool request', async () => { + const visibility = capabilityVisibilityFromServerNames(['filesystem']); + const listed = await orchestrator.callMetaTool('tool_list', {}, visibility); + const hiddenSchema = await orchestrator.callMetaTool( + 'tool_schema', + { server: 'database', toolName: 'read_file' }, + visibility, + ); + const hiddenInvoke = await orchestrator.callMetaTool( + 'tool_invoke', + { server: 'database', toolName: 'read_file', args: {} }, + visibility, + ); - // The result should respect the session filter - // (actual filtering logic is in MetaToolProvider, we just verify it receives sessionId) - const sessionFilter = orchestrator.getSessionAllowedServers(sessionId); - expect(sessionFilter).toBeDefined(); - expect(sessionFilter?.has('filesystem')).toBe(true); + expect((listed as { servers?: string[] }).servers).toEqual(['filesystem']); + expect((hiddenSchema as { error?: { type?: string } }).error?.type).toBe('not_found'); + expect((hiddenInvoke as { error?: { type?: string } }).error?.type).toBe('not_found'); }); - it('should handle undefined sessionId gracefully', async () => { - const allowedServers = new Set(['filesystem', 'database']); - - // Set filter with undefined sessionId - await orchestrator.getCapabilitiesForFilteredServers(allowedServers, undefined); - - // Should store under undefined key - const retrieved = orchestrator.getSessionAllowedServers(undefined); - expect(retrieved).toBeDefined(); - expect(retrieved?.size).toBe(2); + it('discovers and invokes a late-ready server after its capability snapshot refresh', async () => { + mockOutboundConnections.set('late-ready', { + name: 'late-ready', + client: mockClient, + status: ClientStatus.Connected, + transport: { + tags: ['fs'], + start: async () => {}, + send: async () => undefined, + close: async () => {}, + }, + capabilities: { tools: {} }, + } as any); + + await orchestrator.refreshCapabilities(); + const visibility = capabilityVisibilityFromServerNames(['late-ready']); + const listed = await orchestrator.callMetaTool('tool_list', {}, visibility); + const invoked = await orchestrator.callMetaTool( + 'tool_invoke', + { server: 'late-ready', toolName: 'read_file', args: { path: '/tmp/example' } }, + visibility, + ); - // Clear undefined session - orchestrator.clearSessionFilter(undefined); - expect(orchestrator.getSessionAllowedServers(undefined)).toBeUndefined(); + expect((listed as { servers?: string[] }).servers).toEqual(['late-ready']); + expect((invoked as { error?: unknown }).error).toBeUndefined(); + expect(mockClient.callTool).toHaveBeenCalledWith({ + name: 'read_file', + arguments: { path: '/tmp/example' }, + }); }); it('should handle empty filter set', async () => { - const sessionId = 'session-empty'; - const emptyServers = new Set(); + const visibility = capabilityVisibilityFromServerNames([]); const mockCapabilities = { tools: [], @@ -1185,7 +1145,7 @@ describe('LazyLoadingOrchestrator', () => { vi.spyOn(orchestrator['capabilityAggregator'], 'getCurrentCapabilities').mockReturnValue(mockCapabilities); - const filteredCaps = await orchestrator.getCapabilitiesForFilteredServers(emptyServers, sessionId); + const filteredCaps = await orchestrator.getCapabilitiesForVisibility(visibility); // All resources/prompts/servers should be filtered out expect(filteredCaps.resources.length).toBe(0); @@ -1216,8 +1176,7 @@ describe('LazyLoadingOrchestrator', () => { const disabledOrchestrator = new LazyLoadingOrchestrator(mockOutboundConnections, mockAgentConfig); await disabledOrchestrator.initialize(); - const sessionId = 'session-disabled'; - const allowedServers = new Set(['filesystem']); + const visibility = capabilityVisibilityFromServerNames(['filesystem']); const mockCapabilities = { tools: [ @@ -1240,7 +1199,7 @@ describe('LazyLoadingOrchestrator', () => { mockCapabilities, ); - const caps = await disabledOrchestrator.getCapabilitiesForFilteredServers(allowedServers, sessionId); + const caps = await disabledOrchestrator.getCapabilitiesForVisibility(visibility); // When disabled, should return all capabilities without filtering expect(caps.resources.length).toBe(2); @@ -1248,8 +1207,7 @@ describe('LazyLoadingOrchestrator', () => { }); it('should filter resources with complex server names correctly', async () => { - const sessionId = 'session-complex'; - const allowedServers = new Set(['server-with-dashes', 'server_with_underscores']); + const visibility = capabilityVisibilityFromServerNames(['server-with-dashes', 'server_with_underscores']); const mockCapabilities = { tools: [], @@ -1265,7 +1223,7 @@ describe('LazyLoadingOrchestrator', () => { vi.spyOn(orchestrator['capabilityAggregator'], 'getCurrentCapabilities').mockReturnValue(mockCapabilities); - const filteredCaps = await orchestrator.getCapabilitiesForFilteredServers(allowedServers, sessionId); + const filteredCaps = await orchestrator.getCapabilitiesForVisibility(visibility); expect(filteredCaps.resources.length).toBe(2); expect(filteredCaps.resources.map((r) => r.name)).toContain('server-with-dashes_1mcp_resource1'); diff --git a/src/core/capabilities/lazyLoadingOrchestrator.ts b/src/core/capabilities/lazyLoadingOrchestrator.ts index 887de8f9..13124e54 100644 --- a/src/core/capabilities/lazyLoadingOrchestrator.ts +++ b/src/core/capabilities/lazyLoadingOrchestrator.ts @@ -12,6 +12,7 @@ import logger, { debugIf, errorIf } from '@src/logger/logger.js'; import { AsyncLoadingOrchestrator } from './asyncLoadingOrchestrator.js'; import { AsyncLoadingOrchestratorEvent } from './asyncLoadingOrchestratorEvent.js'; +import { type CapabilityVisibility, getCapabilityVisibleServerNames } from './capabilityVisibility.js'; import { AggregatedCapabilities, CapabilityAggregator } from './capabilityAggregator.js'; import { MetaToolProvider } from './metaToolProvider.js'; import { SchemaCache, SchemaCacheConfig } from './schemaCache.js'; @@ -62,8 +63,6 @@ export class LazyLoadingOrchestrator extends EventEmitter { private capabilityAggregator: CapabilityAggregator; private isInitialized: boolean = false; private asyncOrchestrator?: AsyncLoadingOrchestrator; - // Session-specific filters for allowed servers (keyed by sessionId) - private sessionAllowedServers: Map> = new Map(); private connectionResolver: ConnectionResolver; constructor( @@ -161,8 +160,7 @@ export class LazyLoadingOrchestrator extends EventEmitter { */ private async buildToolRegistry(): Promise { // Build tools map for registry by fetching tools directly from each connection - const toolsMap = new Map(); - const serverTags = new Map(); + const registryTools: Array<{ tool: Tool; server: string; connectionKey: string; tags: string[] }> = []; const failedServers: Array<{ server: string; error: string }> = []; const serverConfigs = McpConfigManager.getInstance().getTransportConfig(); @@ -182,11 +180,15 @@ export class LazyLoadingOrchestrator extends EventEmitter { // Map keys for template servers include hash suffix (e.g., "template-server:abc123") // but connection.name is the clean name (e.g., "template-server") // This ensures tool registry uses consistent server names - toolsMap.set(effectiveServerName, serverTools); - - // Get tags from transport const tags = connection.transport.tags || []; - serverTags.set(effectiveServerName, tags); + registryTools.push( + ...serverTools.map((tool) => ({ + tool, + server: effectiveServerName, + connectionKey: serverName, + tags, + })), + ); } } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); @@ -209,7 +211,7 @@ export class LazyLoadingOrchestrator extends EventEmitter { } } - this.toolRegistry = ToolRegistry.fromToolsMap(toolsMap, serverTags); + this.toolRegistry = ToolRegistry.fromToolsWithServer(registryTools); } /** @@ -340,22 +342,10 @@ export class LazyLoadingOrchestrator extends EventEmitter { /** * Get aggregated capabilities for a filtered set of servers * - * This method stores session-specific server filters that will be used by: - * - callMetaTool() to filter tool listings via MetaToolProvider - * - * The filter persists for the session and can be cleared by calling clearSessionFilter(). - * - * @param filteredServerNames - Set of server names to include in capabilities - * @param sessionId - Optional session ID to associate the filter with + * @param visibility - Request-scoped capability visibility * @returns Aggregated capabilities filtered to only include specified servers */ - public async getCapabilitiesForFilteredServers( - filteredServerNames: Set, - sessionId?: string, - ): Promise { - // Store the session-specific filter - this.sessionAllowedServers.set(sessionId, filteredServerNames); - + public async getCapabilitiesForVisibility(visibility: CapabilityVisibility): Promise { // Get the base capabilities const lazyConfig = this.config.get('lazyLoading'); const baseCapabilities = this.capabilityAggregator.getCurrentCapabilities(); @@ -368,6 +358,7 @@ export class LazyLoadingOrchestrator extends EventEmitter { // Meta-tools are always included (they're gateway tools) // The filter will be applied when tools are listed via meta-tools const metaTools = this.metaToolProvider?.getMetaTools() || []; + const visibleServerNames = getCapabilityVisibleServerNames(visibility); // Filter resources to only include those from filtered servers const filteredResources = baseCapabilities.resources.filter((resource) => { @@ -376,7 +367,7 @@ export class LazyLoadingOrchestrator extends EventEmitter { // Extract server name from resource URI const parts = resourceName.split(MCP_URI_SEPARATOR); const serverName = parts[0]; - return filteredServerNames.has(serverName); + return visibleServerNames.has(serverName); }); // Filter prompts to only include those from filtered servers @@ -385,12 +376,12 @@ export class LazyLoadingOrchestrator extends EventEmitter { // Prompts are namespaced with server prefix (e.g., "server_1mcp_prompt") const parts = promptName.split(MCP_URI_SEPARATOR); const serverName = parts[0]; - return filteredServerNames.has(serverName); + return visibleServerNames.has(serverName); }); // Filter ready servers const filteredReadyServers = baseCapabilities.readyServers.filter((serverName) => - filteredServerNames.has(serverName), + visibleServerNames.has(serverName), ); return { @@ -402,23 +393,6 @@ export class LazyLoadingOrchestrator extends EventEmitter { }; } - /** - * Clear the session-specific filter - * @param sessionId - Optional session ID whose filter should be cleared - */ - public clearSessionFilter(sessionId?: string): void { - this.sessionAllowedServers.delete(sessionId); - } - - /** - * Get the allowed servers for a specific session - * @param sessionId - Optional session ID - * @returns Set of allowed server names, or undefined if no filter is set - */ - public getSessionAllowedServers(sessionId?: string): Set | undefined { - return this.sessionAllowedServers.get(sessionId); - } - /** * Handle listChanged notifications based on lazy loading state */ @@ -522,24 +496,18 @@ export class LazyLoadingOrchestrator extends EventEmitter { * Call a meta-tool if in meta-tool mode * @param name - Meta-tool name * @param args - Meta-tool arguments - * @param sessionId - Optional session ID to retrieve filter for - * @param allowedServers - Optional explicit allowed servers set (takes precedence over session filter) + * @param visibility - Request-scoped Filter Selection and Server Candidate Set */ public async callMetaTool( name: string, args: unknown, - sessionId?: string, - allowedServers?: Set, + visibility?: CapabilityVisibility, ): Promise { if (!this.metaToolProvider) { throw new Error('Meta-tool provider not initialized'); } - // Use explicitly provided allowedServers, or fall back to session-specific filter - const effectiveAllowedServers = - allowedServers !== undefined ? allowedServers : this.sessionAllowedServers.get(sessionId); - - return this.metaToolProvider.callMetaTool(name, args, sessionId, effectiveAllowedServers); + return this.metaToolProvider.callMetaTool(name, args, visibility); } /** diff --git a/src/core/capabilities/lazyLoadingPerformance.test.ts b/src/core/capabilities/lazyLoadingPerformance.test.ts index 84e6081f..453102e5 100644 --- a/src/core/capabilities/lazyLoadingPerformance.test.ts +++ b/src/core/capabilities/lazyLoadingPerformance.test.ts @@ -330,20 +330,17 @@ describe('Lazy Loading Performance Tests', () => { }); }); - describe('Mode Comparison', () => { - it('should compare metatool vs hybrid vs full modes', async () => { - const modes = ['metatool', 'hybrid', 'full'] as const; + describe('Boolean lazy-loading contract', () => { + it('uses only enabled to select meta-tool or direct exposure', async () => { + const states = [true, false] as const; + const results: boolean[] = []; - const results: Array<{ mode: string; enabled: boolean }> = []; - - for (const mode of modes) { + for (const enabled of states) { mockAgentConfig.get.mockImplementation((key: string) => { if (key === 'lazyLoading') { return { - enabled: mode !== 'full', - mode, + enabled, metaTools: { enabled: true, inlineCatalog: false, catalogFormat: 'grouped' }, - directExpose: mode === 'hybrid' ? ['filesystem_*'] : [], cache: { maxEntries: 1000, strategy: 'lru' }, preload: { patterns: [], keywords: [] }, fallback: { onError: 'skip', timeoutMs: 5000 }, @@ -357,22 +354,10 @@ describe('Lazy Loading Performance Tests', () => { const stats = orchestrator.getStatistics(); - results.push({ - mode, - enabled: stats.enabled, - }); + results.push(stats.enabled); } - // Metatool and hybrid should be enabled - const metatoolResult = results.find((r) => r.mode === 'metatool'); - expect(metatoolResult?.enabled).toBe(true); - - const hybridResult = results.find((r) => r.mode === 'hybrid'); - expect(hybridResult?.enabled).toBe(true); - - // Full mode should be disabled - const fullResult = results.find((r) => r.mode === 'full'); - expect(fullResult?.enabled).toBe(false); + expect(results).toEqual([true, false]); }); }); }); diff --git a/src/core/capabilities/metaToolProvider-template-servers.test.ts b/src/core/capabilities/metaToolProvider-template-servers.test.ts index d7824cd5..bdc20b37 100644 --- a/src/core/capabilities/metaToolProvider-template-servers.test.ts +++ b/src/core/capabilities/metaToolProvider-template-servers.test.ts @@ -12,6 +12,7 @@ import { ClientStatus, OutboundConnections } from '@src/core/types/client.js'; import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { createCapabilityVisibility } from './capabilityVisibility.js'; import { MetaToolProvider } from './metaToolProvider.js'; import { SchemaCache } from './schemaCache.js'; import { ToolRegistry } from './toolRegistry.js'; @@ -162,7 +163,27 @@ describe('MetaToolProvider - Template Server Support', () => { it('should resolve list routes through the caller session for shareable templates', async () => { const sessionAwareProvider = new MetaToolProvider( - () => toolRegistry, + () => + ToolRegistry.fromToolsWithServer([ + { + tool: { + name: 'template_tool', + description: 'A tool from the session server', + inputSchema: { type: 'object', properties: {} }, + }, + server: 'template-server', + connectionKey: 'template-server:session-123', + }, + { + tool: { + name: 'template_tool', + description: 'A tool from the shareable server', + inputSchema: { type: 'object', properties: {} }, + }, + server: 'template-server', + connectionKey: 'template-server:abc123', + }, + ]), schemaCache, outboundConnections, undefined, @@ -174,10 +195,103 @@ describe('MetaToolProvider - Template Server Support', () => { ); const listVisibleTools = vi.spyOn((sessionAwareProvider as any).capabilityCatalog, 'listVisibleTools'); - const result = await sessionAwareProvider.callMetaTool('tool_list', {}, 'session-123'); + const visibility = createCapabilityVisibility( + [['template-server:abc123', 'template-server']], + 'session-123', + ); + const result = await sessionAwareProvider.callMetaTool('tool_list', {}, visibility); expect((result as any).error).toBeUndefined(); - expect(listVisibleTools).toHaveBeenCalledWith({}, 'session-123', undefined); + expect(listVisibleTools).toHaveBeenCalledWith({}, visibility); + }); + + it('keeps differing per-session tool surfaces, schemas, and invocation routes isolated', async () => { + const sessionAClient = { + listTools: vi.fn().mockResolvedValue({ + tools: [{ name: 'tool_a', description: 'session A', inputSchema: { type: 'object', title: 'A' } }], + }), + callTool: vi.fn().mockResolvedValue({ content: [{ type: 'text', text: 'A' }] }), + } as unknown as Client; + const sessionBClient = { + listTools: vi.fn().mockResolvedValue({ + tools: [{ name: 'tool_b', description: 'session B', inputSchema: { type: 'object', title: 'B' } }], + }), + callTool: vi.fn().mockResolvedValue({ content: [{ type: 'text', text: 'B' }] }), + } as unknown as Client; + const connections: OutboundConnections = new Map([ + [ + 'template-server:session-a', + { + name: 'template-server', + client: sessionAClient, + status: ClientStatus.Connected, + capabilities: { tools: {} }, + transport: 'stdio' as any, + }, + ], + [ + 'template-server:session-b', + { + name: 'template-server', + client: sessionBClient, + status: ClientStatus.Connected, + capabilities: { tools: {} }, + transport: 'stdio' as any, + }, + ], + ]); + const registry = ToolRegistry.fromToolsWithServer([ + { + tool: { name: 'tool_a', description: 'session A', inputSchema: { type: 'object', title: 'A' } }, + server: 'template-server', + connectionKey: 'template-server:session-a', + }, + { + tool: { name: 'tool_b', description: 'session B', inputSchema: { type: 'object', title: 'B' } }, + server: 'template-server', + connectionKey: 'template-server:session-b', + }, + ]); + const provider = new MetaToolProvider( + () => registry, + schemaCache, + connections, + async (connectionKey, toolName) => { + const result = await connections.get(connectionKey)!.client.listTools(); + return result.tools.find((tool) => tool.name === toolName)!; + }, + ); + const sessionAVisibility = createCapabilityVisibility( + [['template-server:session-a', 'template-server']], + 'session-a', + ); + const sessionBVisibility = createCapabilityVisibility( + [['template-server:session-b', 'template-server']], + 'session-b', + ); + + const listedA = await provider.callMetaTool('tool_list', {}, sessionAVisibility); + const listedB = await provider.callMetaTool('tool_list', {}, sessionBVisibility); + expect((listedA as any).tools.map((tool: Tool) => tool.name)).toEqual(['tool_a']); + expect((listedB as any).tools.map((tool: Tool) => tool.name)).toEqual(['tool_b']); + expect((listedA as any).tools[0].server).toBe('template-server'); + + const schemaA = await provider.callMetaTool( + 'tool_schema', + { server: 'template-server', toolName: 'tool_a' }, + sessionAVisibility, + ); + const invokeB = await provider.callMetaTool( + 'tool_invoke', + { server: 'template-server', toolName: 'tool_b', args: {} }, + sessionBVisibility, + ); + + expect((schemaA as any).schema.inputSchema.title).toBe('A'); + expect(schemaCache.getIfCached('template-server:session-a', 'tool_a')).not.toBeNull(); + expect((invokeB as any).error).toBeUndefined(); + expect(sessionBClient.callTool).toHaveBeenCalledOnce(); + expect(sessionAClient.callTool).not.toHaveBeenCalled(); }); }); @@ -316,7 +430,27 @@ describe('MetaToolProvider - Template Server Support', () => { ]); const sessionAwareProvider = new MetaToolProvider( - () => toolRegistry, + () => + ToolRegistry.fromToolsWithServer([ + { + tool: { + name: 'template_tool', + description: 'A tool from the session server', + inputSchema: { type: 'object', properties: {} }, + }, + server: 'template-server', + connectionKey: 'template-server:session-123', + }, + { + tool: { + name: 'template_tool', + description: 'A tool from the shareable server', + inputSchema: { type: 'object', properties: {} }, + }, + server: 'template-server', + connectionKey: 'template-server:abc123', + }, + ]), schemaCache, multiConnections, undefined, @@ -334,7 +468,7 @@ describe('MetaToolProvider - Template Server Support', () => { toolName: 'template_tool', args: { message: 'test' }, }, - 'session-123', + createCapabilityVisibility([['template-server:session-123', 'template-server']], 'session-123'), ); expect((result as any).error).toBeUndefined(); diff --git a/src/core/capabilities/metaToolProvider.test.ts b/src/core/capabilities/metaToolProvider.test.ts index bf8ea1e2..e252286d 100644 --- a/src/core/capabilities/metaToolProvider.test.ts +++ b/src/core/capabilities/metaToolProvider.test.ts @@ -4,11 +4,13 @@ import { ClientStatus, OutboundConnections } from '@src/core/types/index.js'; import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { capabilityVisibilityFromServerNames, createCapabilityVisibility } from './capabilityVisibility.js'; import { MetaToolProvider } from './metaToolProvider.js'; import { SchemaCache } from './schemaCache.js'; import { ToolRegistry } from './toolRegistry.js'; const mockGetTransportConfig = vi.fn().mockReturnValue({}); +const visibility = (...serverNames: string[]) => capabilityVisibilityFromServerNames(serverNames); vi.mock('@src/config/mcpConfigManager.js', () => ({ McpConfigManager: { @@ -675,8 +677,20 @@ describe('MetaToolProvider', () => { ]) as OutboundConnections; const sessionSchemaCache = new SchemaCache({ maxEntries: 100 }); + const sessionRegistry = ToolRegistry.fromToolsWithServer([ + { + tool: { name: 'read_file', description: 'session-a', inputSchema: { type: 'object' } }, + server: 'filesystem', + connectionKey: 'filesystem:session-a', + }, + { + tool: { name: 'read_file', description: 'session-b', inputSchema: { type: 'object' } }, + server: 'filesystem', + connectionKey: 'filesystem:session-b', + }, + ]); const sessionProvider = new MetaToolProvider( - () => toolRegistry, + () => sessionRegistry, sessionSchemaCache, sessionScopedConnections, mockSchemaLoader, @@ -692,7 +706,7 @@ describe('MetaToolProvider', () => { server: 'filesystem', toolName: 'read_file', }, - 'session-a', + createCapabilityVisibility([['filesystem:session-a', 'filesystem']], 'session-a'), ); await sessionProvider.callMetaTool( 'tool_schema', @@ -700,7 +714,7 @@ describe('MetaToolProvider', () => { server: 'filesystem', toolName: 'read_file', }, - 'session-b', + createCapabilityVisibility([['filesystem:session-b', 'filesystem']], 'session-b'), ); expect(mockSchemaLoader).toHaveBeenNthCalledWith(1, 'filesystem:session-a', 'read_file'); @@ -922,9 +936,9 @@ describe('MetaToolProvider', () => { multiServerProvider = new MetaToolProvider(() => multiServerRegistry, schemaCache, multiServerConnections); }); - it('should filter tool_list results when allowedServers is set', async () => { - // Set allowed servers to only filesystem and search - multiServerProvider.setAllowedServers(new Set(['filesystem', 'search'])); + it('should filter tool_list results when Capability Visibility is set', async () => { + // Set Server Candidate Set to only filesystem and search + multiServerProvider.setCapabilityVisibility(visibility('filesystem', 'search')); const result = await multiServerProvider.callMetaTool('tool_list', {}); @@ -954,9 +968,9 @@ describe('MetaToolProvider', () => { } }); - it('should return all tools when allowedServers is undefined', async () => { + it('should return all tools when Capability Visibility is undefined', async () => { // Clear any filtering - multiServerProvider.setAllowedServers(undefined); + multiServerProvider.setCapabilityVisibility(undefined); const result = await multiServerProvider.callMetaTool('tool_list', {}); @@ -980,9 +994,9 @@ describe('MetaToolProvider', () => { } }); - it('should return empty list when allowedServers is empty set', async () => { - // Set allowed servers to empty set - multiServerProvider.setAllowedServers(new Set([])); + it('should return empty list when Capability Visibility is empty set', async () => { + // Set Server Candidate Set to empty set + multiServerProvider.setCapabilityVisibility(visibility()); const result = await multiServerProvider.callMetaTool('tool_list', {}); @@ -1001,9 +1015,9 @@ describe('MetaToolProvider', () => { } }); - it('should filter tool_schema access to allowed servers only', async () => { - // Set allowed servers to only filesystem - multiServerProvider.setAllowedServers(new Set(['filesystem'])); + it('should filter tool_schema access to Server Candidate Set only', async () => { + // Set Server Candidate Set to only filesystem + multiServerProvider.setCapabilityVisibility(visibility('filesystem')); // Cache a tool from filesystem server const filesystemTool: Tool = { @@ -1050,9 +1064,9 @@ describe('MetaToolProvider', () => { } }); - it('should filter tool_invoke access to allowed servers only', async () => { - // Set allowed servers to only search - multiServerProvider.setAllowedServers(new Set(['search'])); + it('should filter tool_invoke access to Server Candidate Set only', async () => { + // Set Server Candidate Set to only search + multiServerProvider.setCapabilityVisibility(visibility('search')); const searchClient = multiServerConnections.get('search')?.client; const filesystemClient = multiServerConnections.get('filesystem')?.client; @@ -1105,8 +1119,8 @@ describe('MetaToolProvider', () => { }); it('should return not_found error for filtered servers', async () => { - // Set allowed servers to only database - multiServerProvider.setAllowedServers(new Set(['database'])); + // Set Server Candidate Set to only database + multiServerProvider.setCapabilityVisibility(visibility('database')); // Try to access a tool from filtered-out filesystem server const result = await multiServerProvider.callMetaTool('tool_invoke', { @@ -1131,9 +1145,9 @@ describe('MetaToolProvider', () => { } }); - it('should dynamically update filtering when setAllowedServers is called multiple times', async () => { + it('should dynamically update filtering when setCapabilityVisibility is called multiple times', async () => { // Start with all servers allowed - multiServerProvider.setAllowedServers(undefined); + multiServerProvider.setCapabilityVisibility(undefined); let result = await multiServerProvider.callMetaTool('tool_list', {}); if ('tools' in result && 'totalCount' in result) { @@ -1141,7 +1155,7 @@ describe('MetaToolProvider', () => { } // Filter to only filesystem - multiServerProvider.setAllowedServers(new Set(['filesystem'])); + multiServerProvider.setCapabilityVisibility(visibility('filesystem')); result = await multiServerProvider.callMetaTool('tool_list', {}); if ('tools' in result && 'totalCount' in result) { @@ -1149,7 +1163,7 @@ describe('MetaToolProvider', () => { } // Filter to filesystem and database - multiServerProvider.setAllowedServers(new Set(['filesystem', 'database'])); + multiServerProvider.setCapabilityVisibility(visibility('filesystem', 'database')); result = await multiServerProvider.callMetaTool('tool_list', {}); if ('tools' in result && 'totalCount' in result) { @@ -1157,7 +1171,7 @@ describe('MetaToolProvider', () => { } // Clear filter again - multiServerProvider.setAllowedServers(undefined); + multiServerProvider.setCapabilityVisibility(undefined); result = await multiServerProvider.callMetaTool('tool_list', {}); if ('tools' in result && 'totalCount' in result) { @@ -1165,10 +1179,10 @@ describe('MetaToolProvider', () => { } }); - it('should prefer per-call allowedServers over shared provider state', async () => { - multiServerProvider.setAllowedServers(new Set(['database'])); + it('should prefer per-call Capability Visibility over shared provider state', async () => { + multiServerProvider.setCapabilityVisibility(visibility('database')); - const result = await multiServerProvider.callMetaTool('tool_list', {}, undefined, new Set(['filesystem'])); + const result = await multiServerProvider.callMetaTool('tool_list', {}, visibility('filesystem')); if ('error' in result && result.error) { throw new Error(result.error.message); @@ -1190,7 +1204,7 @@ describe('MetaToolProvider', () => { throw new Error('Registry corrupted'); }), getServers: vi.fn(() => []), - filterByServers: vi.fn(function (this: any) { + filterByConnectionKeys: vi.fn(function (this: any) { return this; }), hasTool: vi.fn(), @@ -1210,7 +1224,7 @@ describe('MetaToolProvider', () => { const failingRegistry: any = { listTools: vi.fn(() => ({ tools: [] })), getServers: vi.fn(() => []), - filterByServers: vi.fn(function (this: any) { + filterByConnectionKeys: vi.fn(function (this: any) { return this; }), hasTool: vi.fn(() => true), // Tool exists @@ -1250,7 +1264,7 @@ describe('MetaToolProvider', () => { throw new Error('Internal failure'); }), getServers: vi.fn(() => []), - filterByServers: vi.fn(function (this: any) { + filterByConnectionKeys: vi.fn(function (this: any) { return this; }), hasTool: vi.fn(), diff --git a/src/core/capabilities/metaToolProvider.ts b/src/core/capabilities/metaToolProvider.ts index 40bcf9e5..a4a764bf 100644 --- a/src/core/capabilities/metaToolProvider.ts +++ b/src/core/capabilities/metaToolProvider.ts @@ -1,13 +1,13 @@ import { Tool } from '@modelcontextprotocol/sdk/types.js'; import { McpConfigManager } from '@src/config/mcpConfigManager.js'; -import { ConnectionResolver, TemplateHashProvider } from '@src/core/server/connectionResolver.js'; -import { getDisabledToolError } from '@src/core/server/disabledTools.js'; +import { TemplateHashProvider } from '@src/core/server/connectionResolver.js'; import { OutboundConnections } from '@src/core/types/index.js'; import logger, { errorIf } from '@src/logger/logger.js'; import { zodToInputSchema, zodToOutputSchema } from '@src/utils/schemaUtils.js'; import { CapabilityCatalog } from './capabilityCatalog.js'; +import type { CapabilityVisibility } from './capabilityVisibility.js'; import { SchemaCache } from './schemaCache.js'; import { ToolInvokeInputSchema, @@ -20,7 +20,7 @@ import { ToolSchemaOutput, ToolSchemaOutputSchema, } from './schemas/metaToolSchemas.js'; -import { ToolMetadata, ToolRegistry } from './toolRegistry.js'; +import { type ToolMetadata, ToolRegistry } from './toolRegistry.js'; /** * Result types for meta-tools @@ -86,8 +86,7 @@ export class MetaToolProvider { private schemaCache: SchemaCache; private outboundConnections: OutboundConnections; private loadSchema?: SchemaLoader; - private allowedServers?: Set; - private connectionResolver: ConnectionResolver; + private defaultVisibility?: CapabilityVisibility; private capabilityCatalog: CapabilityCatalog; private templateHashProvider?: TemplateHashProvider; @@ -96,101 +95,42 @@ export class MetaToolProvider { schemaCache: SchemaCache, outboundConnections: OutboundConnections, loadSchema?: SchemaLoader, - allowedServers?: Set, + defaultVisibility?: CapabilityVisibility, templateHashProvider?: TemplateHashProvider, ) { this.getToolRegistry = getToolRegistry; this.schemaCache = schemaCache; this.outboundConnections = outboundConnections; this.loadSchema = loadSchema; - this.allowedServers = allowedServers; + this.defaultVisibility = defaultVisibility; this.templateHashProvider = templateHashProvider; - this.connectionResolver = new ConnectionResolver(outboundConnections, templateHashProvider); this.capabilityCatalog = new CapabilityCatalog({ getToolRegistry, schemaCache, outboundConnections, loadSchema, - defaultAllowedServers: allowedServers, + defaultVisibility, templateHashProvider, getServerConfigs: () => McpConfigManager.getInstance().getTransportConfig(), }); } /** - * Resolve a clean server name to the actual connection key. - * Delegates to ConnectionResolver.findByServerName for unified resolution logic. - * - * @param cleanServerName - The clean server name (without hash suffix) - * @returns The actual connection key, or the original name if not found + * Set the default capability visibility for callers without request context. */ - private resolveConnectionKey(cleanServerName: string, sessionId?: string): string { - const result = - this.connectionResolver.resolveWithKey(cleanServerName, sessionId) ?? - this.connectionResolver.findByServerName(cleanServerName); - return result?.key ?? cleanServerName; - } - - /** - * Set the allowed servers filter - * @param serverNames - Set of server names to allow, or undefined to allow all - */ - public setAllowedServers(serverNames?: Set): void { - this.allowedServers = serverNames; + public setCapabilityVisibility(visibility?: CapabilityVisibility): void { + this.defaultVisibility = visibility; this.capabilityCatalog = new CapabilityCatalog({ getToolRegistry: this.getToolRegistry, schemaCache: this.schemaCache, outboundConnections: this.outboundConnections, loadSchema: this.loadSchema, - defaultAllowedServers: serverNames, + defaultVisibility: visibility, templateHashProvider: this.templateHashProvider, getServerConfigs: () => McpConfigManager.getInstance().getTransportConfig(), }); } - private getDisabledError(logicalServerName: string, toolName: string) { - return getDisabledToolError(McpConfigManager.getInstance().getTransportConfig(), logicalServerName, toolName); - } - - private validateResolvedToolAccess( - args: DescribeToolArgs | CallToolArgs, - allowedServers?: Set, - ): { error?: DescribeToolResult['error'] | CallToolResult['error'] } { - if (!args.server || !args.toolName) { - return { - error: { - type: 'validation', - message: 'Validation Error: "server" and "toolName" are required parameters', - }, - }; - } - - if (!this.toolRegistry(allowedServers).hasTool(args.server, args.toolName)) { - return { - error: { - type: 'not_found', - message: `Tool not found: ${args.server}:${args.toolName}. Call tool_list to see available tools.`, - }, - }; - } - - return { - error: this.getDisabledError(args.server, args.toolName), - }; - } - - /** - * Get the current tool registry, optionally filtered by allowed servers - */ - private toolRegistry(allowedServers?: Set): ToolRegistry { - const registry = this.getToolRegistry(); - const effectiveAllowedServers = allowedServers ?? this.allowedServers; - if (effectiveAllowedServers !== undefined) { - return registry.filterByServers(effectiveAllowedServers); - } - return registry; - } - /** * Get all available meta-tools (3 discovery tools) */ @@ -204,8 +144,7 @@ export class MetaToolProvider { public async callMetaTool( name: string, args: unknown, - sessionId?: string, - allowedServers?: Set, + visibility?: CapabilityVisibility, ): Promise { switch (name) { case 'tool_list': { @@ -222,7 +161,7 @@ export class MetaToolProvider { }, } as ListToolsResult; } - return this.listAvailableTools(parsed.data, sessionId, allowedServers); + return this.listAvailableTools(parsed.data, visibility); } case 'tool_schema': { const parsed = ToolSchemaInputSchema.safeParse(args); @@ -235,7 +174,7 @@ export class MetaToolProvider { }, } as DescribeToolResult; } - return this.describeTool(parsed.data, sessionId, allowedServers); + return this.describeTool(parsed.data, visibility); } case 'tool_invoke': { const parsed = ToolInvokeInputSchema.safeParse(args); @@ -250,7 +189,7 @@ export class MetaToolProvider { }, } as CallToolResult; } - return this.callTool(parsed.data, sessionId, allowedServers); + return this.callTool(parsed.data, visibility); } default: return { @@ -283,11 +222,10 @@ export class MetaToolProvider { */ private async listAvailableTools( args: ListAvailableToolsArgs, - sessionId?: string, - allowedServers?: Set, + visibility?: CapabilityVisibility, ): Promise { try { - const result = await this.capabilityCatalog.listVisibleTools(args, sessionId, allowedServers); + const result = await this.capabilityCatalog.listVisibleTools(args, visibility); // Format tools for response const tools = result.tools.map((tool: ToolMetadata) => ({ @@ -347,11 +285,10 @@ export class MetaToolProvider { */ private async describeTool( args: DescribeToolArgs, - sessionId?: string, - allowedServers?: Set, + visibility?: CapabilityVisibility, ): Promise { try { - const result = await this.capabilityCatalog.describeVisibleTool(args, sessionId, allowedServers); + const result = await this.capabilityCatalog.describeVisibleTool(args, visibility); if (result.error) { return { schema: {}, @@ -397,11 +334,10 @@ export class MetaToolProvider { */ private async callTool( args: CallToolArgs, - sessionId?: string, - allowedServers?: Set, + visibility?: CapabilityVisibility, ): Promise { try { - const result = await this.capabilityCatalog.invokeVisibleTool(args, sessionId, allowedServers); + const result = await this.capabilityCatalog.invokeVisibleTool(args, visibility); if (result.error) { return { result: {}, diff --git a/src/core/capabilities/toolRegistry.filter-servers.test.ts b/src/core/capabilities/toolRegistry.filter-servers.test.ts index 8c24dfe5..acf2df73 100644 --- a/src/core/capabilities/toolRegistry.filter-servers.test.ts +++ b/src/core/capabilities/toolRegistry.filter-servers.test.ts @@ -23,10 +23,10 @@ describe('ToolRegistry', () => { } }); - describe('filterByServers', () => { + describe('filterByConnectionKeys', () => { it('should return new registry with only tools from specified servers', () => { - const serverNames = new Set(['filesystem', 'database']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['filesystem', 'database']); + const filtered = registry.filterByConnectionKeys(connectionKeys); expect(filtered.size()).toBe(4); expect(filtered.hasTool('filesystem', 'read_file')).toBe(true); @@ -38,24 +38,24 @@ describe('ToolRegistry', () => { }); it('should return empty registry when no servers match', () => { - const serverNames = new Set(['nonexistent', 'another_nonexistent']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['nonexistent', 'another_nonexistent']); + const filtered = registry.filterByConnectionKeys(connectionKeys); expect(filtered.size()).toBe(0); expect(filtered.getAllTools()).toHaveLength(0); }); it('should handle empty server set', () => { - const serverNames = new Set(); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(); + const filtered = registry.filterByConnectionKeys(connectionKeys); expect(filtered.size()).toBe(0); expect(filtered.getAllTools()).toHaveLength(0); }); it('should preserve tool metadata in filtered registry', () => { - const serverNames = new Set(['filesystem']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['filesystem']); + const filtered = registry.filterByConnectionKeys(connectionKeys); const readFileTool = filtered.getTool('filesystem', 'read_file'); expect(readFileTool).toEqual({ @@ -75,8 +75,8 @@ describe('ToolRegistry', () => { }); it('should work correctly with getAllTools() after filtering', () => { - const serverNames = new Set(['database', 'git']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['database', 'git']); + const filtered = registry.filterByConnectionKeys(connectionKeys); const allTools = filtered.getAllTools(); expect(allTools).toHaveLength(3); @@ -87,8 +87,8 @@ describe('ToolRegistry', () => { }); it('should work correctly with listTools() after filtering', () => { - const serverNames = new Set(['filesystem']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['filesystem']); + const filtered = registry.filterByConnectionKeys(connectionKeys); const result = filtered.listTools({}); expect(result.tools).toHaveLength(2); @@ -98,8 +98,8 @@ describe('ToolRegistry', () => { }); it('should work correctly with getServers() after filtering', () => { - const serverNames = new Set(['filesystem', 'database']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['filesystem', 'database']); + const filtered = registry.filterByConnectionKeys(connectionKeys); const servers = filtered.getServers(); expect(servers).toEqual(['database', 'filesystem']); @@ -108,8 +108,8 @@ describe('ToolRegistry', () => { }); it('should work correctly with getTags() after filtering', () => { - const serverNames = new Set(['filesystem']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['filesystem']); + const filtered = registry.filterByConnectionKeys(connectionKeys); const tags = filtered.getTags(); expect(tags).toContain('fs'); @@ -119,8 +119,8 @@ describe('ToolRegistry', () => { }); it('should work correctly with getToolCountByServer() after filtering', () => { - const serverNames = new Set(['filesystem', 'search']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['filesystem', 'search']); + const filtered = registry.filterByConnectionKeys(connectionKeys); const counts = filtered.getToolCountByServer(); expect(counts).toEqual({ @@ -132,12 +132,12 @@ describe('ToolRegistry', () => { }); it('should support chaining multiple filters', () => { - const serverNames1 = new Set(['filesystem', 'database', 'git']); - const filtered1 = registry.filterByServers(serverNames1); + const connectionKeys1 = new Set(['filesystem', 'database', 'git']); + const filtered1 = registry.filterByConnectionKeys(connectionKeys1); expect(filtered1.size()).toBe(5); - const serverNames2 = new Set(['filesystem', 'git']); - const filtered2 = filtered1.filterByServers(serverNames2); + const connectionKeys2 = new Set(['filesystem', 'git']); + const filtered2 = filtered1.filterByConnectionKeys(connectionKeys2); expect(filtered2.size()).toBe(3); expect(filtered2.hasTool('filesystem', 'read_file')).toBe(true); expect(filtered2.hasTool('git', 'git_status')).toBe(true); @@ -145,8 +145,8 @@ describe('ToolRegistry', () => { }); it('should create independent registry instance', () => { - const serverNames = new Set(['filesystem']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['filesystem']); + const filtered = registry.filterByConnectionKeys(connectionKeys); // Original registry should remain unchanged expect(registry.size()).toBe(6); @@ -158,8 +158,8 @@ describe('ToolRegistry', () => { }); it('should filter single server correctly', () => { - const serverNames = new Set(['search']); - const filtered = registry.filterByServers(serverNames); + const connectionKeys = new Set(['search']); + const filtered = registry.filterByConnectionKeys(connectionKeys); expect(filtered.size()).toBe(1); expect(filtered.hasTool('search', 'search')).toBe(true); diff --git a/src/core/capabilities/toolRegistry.ts b/src/core/capabilities/toolRegistry.ts index e05b3472..022ae639 100644 --- a/src/core/capabilities/toolRegistry.ts +++ b/src/core/capabilities/toolRegistry.ts @@ -8,6 +8,7 @@ import logger, { debugIf, errorIf } from '@src/logger/logger.js'; export interface ToolMetadata { name: string; server: string; + connectionKey?: string; description: string; inputSchema?: Tool['inputSchema']; tags?: string[]; @@ -112,11 +113,12 @@ export class ToolRegistry { * @returns A new ToolRegistry instance */ public static fromToolsWithServer( - toolsWithServer: Array<{ tool: Tool; server: string; tags?: string[] }>, + toolsWithServer: Array<{ tool: Tool; server: string; connectionKey?: string; tags?: string[] }>, ): ToolRegistry { - const tools: ToolMetadata[] = toolsWithServer.map(({ tool, server, tags }) => ({ + const tools: ToolMetadata[] = toolsWithServer.map(({ tool, server, connectionKey, tags }) => ({ name: tool.name, server, + connectionKey: connectionKey ?? server, description: tool.description || '', inputSchema: tool.inputSchema, tags: tags || [], @@ -305,8 +307,18 @@ export class ToolRegistry { * @param serverNames - Set of server names to include * @returns A new ToolRegistry instance with filtered tools */ - public filterByServers(serverNames: Set): ToolRegistry { - const filteredTools = this.tools.filter((tool) => serverNames.has(tool.server)); + public filterByConnectionKeys(connectionKeys: ReadonlySet): ToolRegistry { + const filteredTools = this.tools.filter((tool) => connectionKeys.has(tool.connectionKey ?? tool.server)); + return new ToolRegistry(filteredTools); + } + + /** Filter by exact connection identity, with clean-name fallback for legacy metadata only. */ + public filterByServerCandidates(serverCandidates: ReadonlyMap): ToolRegistry { + const connectionKeys = new Set(serverCandidates.keys()); + const publicServerNames = new Set(serverCandidates.values()); + const filteredTools = this.tools.filter((tool) => + tool.connectionKey ? connectionKeys.has(tool.connectionKey) : publicServerNames.has(tool.server), + ); return new ToolRegistry(filteredTools); } diff --git a/src/core/instructions/templateTypes.ts b/src/core/instructions/templateTypes.ts index acbf81e7..2f8f7eaa 100644 --- a/src/core/instructions/templateTypes.ts +++ b/src/core/instructions/templateTypes.ts @@ -88,7 +88,7 @@ export interface LazyLoadingState { enabled: boolean; /** Lazy loading mode */ - mode: 'metatool' | 'hybrid' | 'full'; + mode: 'metatool' | 'full'; /** Total number of tools available across all servers */ availableToolsCount: number; @@ -96,7 +96,7 @@ export interface LazyLoadingState { /** Number of tools exposed via tools/list (varies by mode) */ exposedToolsCount: number; - /** Number of tools in directExpose list (hybrid mode only) */ + /** Number of directly exposed tools (always zero when lazy loading is enabled) */ directExposeCount: number; /** Number of tools currently cached with schemas */ diff --git a/src/core/protocol/requestHandlerUtils.test.ts b/src/core/protocol/requestHandlerUtils.test.ts new file mode 100644 index 00000000..1612540e --- /dev/null +++ b/src/core/protocol/requestHandlerUtils.test.ts @@ -0,0 +1,51 @@ +import { ClientStatus, type InboundConnection, type OutboundConnections } from '@src/core/types/index.js'; + +import { describe, expect, it, vi } from 'vitest'; + +vi.mock('@src/core/server/serverManager.js', () => ({ + ServerManager: { + get current() { + return { getTemplateServerManager: () => undefined }; + }, + }, +})); + +import { resolveLazyCapabilityVisibility } from './requestHandlerUtils.js'; + +describe('resolveLazyCapabilityVisibility', () => { + it('re-evaluates template scope, tags, connection state, and tool capability for each request', () => { + const connections = new Map([ + ['visible', { name: 'visible', status: ClientStatus.Connected, capabilities: { tools: {} }, transport: { tags: ['safe'] } }], + ['hidden', { name: 'hidden', status: ClientStatus.Connected, capabilities: { tools: {} }, transport: { tags: ['private'] } }], + ['late', { name: 'late', status: ClientStatus.Restarting, capabilities: { tools: {} }, transport: { tags: ['safe'] } }], + ['not-a-tool', { name: 'not-a-tool', status: ClientStatus.Connected, capabilities: {}, transport: { tags: ['safe'] } }], + ['template:other-session', { name: 'template', status: ClientStatus.Connected, capabilities: { tools: {} }, transport: { tags: ['safe'] } }], + ['template:session-1', { name: 'template', status: ClientStatus.Connected, capabilities: { tools: {} }, transport: { tags: ['safe'] } }], + ]) as OutboundConnections; + const inbound = { tags: ['safe'], tagFilterMode: 'simple-or' } as InboundConnection; + + const initialVisibility = resolveLazyCapabilityVisibility(connections, inbound, 'session-1'); + expect(Array.from(initialVisibility.serverCandidates.entries()).sort()).toEqual([ + ['template:session-1', 'template'], + ['visible', 'visible'], + ]); + expect(initialVisibility.sessionId).toBe('session-1'); + + connections.get('late')!.status = ClientStatus.Connected; + + expect( + Array.from(resolveLazyCapabilityVisibility(connections, inbound, 'session-1').serverCandidates.entries()).sort(), + ).toEqual([ + ['late', 'late'], + ['template:session-1', 'template'], + ['visible', 'visible'], + ]); + + connections.get('late')!.status = ClientStatus.Disconnected; + connections.delete('visible'); + + expect( + Array.from(resolveLazyCapabilityVisibility(connections, inbound, 'session-1').serverCandidates.entries()), + ).toEqual([['template:session-1', 'template']]); + }); +}); diff --git a/src/core/protocol/requestHandlerUtils.ts b/src/core/protocol/requestHandlerUtils.ts index ac9e032e..d07cb3c8 100644 --- a/src/core/protocol/requestHandlerUtils.ts +++ b/src/core/protocol/requestHandlerUtils.ts @@ -1,7 +1,9 @@ import { McpConfigManager } from '@src/config/mcpConfigManager.js'; import { CapabilityCatalog } from '@src/core/capabilities/capabilityCatalog.js'; +import { createCapabilityVisibility, type CapabilityVisibility } from '@src/core/capabilities/capabilityVisibility.js'; import { SchemaCache } from '@src/core/capabilities/schemaCache.js'; import { ToolRegistry } from '@src/core/capabilities/toolRegistry.js'; +import { byCapabilities } from '@src/core/filtering/clientFiltering.js'; import { FilteringService } from '@src/core/filtering/filteringService.js'; import { createConnectionResolver } from '@src/core/server/connectionResolver.js'; import { ServerManager } from '@src/core/server/serverManager.js'; @@ -69,3 +71,20 @@ export function filterConnectionsForSession( const resolver = createConnectionResolver(outboundConns, templateServerManager); return resolver.filterForSession(sessionId); } + +/** Resolve request-time Filter Selection into a Server Candidate Set. */ +export function resolveLazyCapabilityVisibility( + outboundConns: OutboundConnections, + inboundConn: InboundConnection, + sessionId: string | undefined, +): CapabilityVisibility { + // Scope template instances before applying client filters and availability. + const sessionScoped = filterConnectionsForSession(outboundConns, sessionId); + const tagAndPresetScoped = FilteringService.getFilteredConnections(sessionScoped, inboundConn); + const toolCapable = byCapabilities({ tools: {} })(tagAndPresetScoped); + + return createCapabilityVisibility( + Array.from(toolCapable.entries(), ([connectionKey, connection]) => [connectionKey, connection.name] as const), + sessionId, + ); +} diff --git a/src/core/protocol/toolRequestHandlers.ts b/src/core/protocol/toolRequestHandlers.ts index 517ea1ed..0ea97c7f 100644 --- a/src/core/protocol/toolRequestHandlers.ts +++ b/src/core/protocol/toolRequestHandlers.ts @@ -23,6 +23,7 @@ import { createCapabilityCatalogFromConnections, filterConnectionsForSession, getRequestSession, + resolveLazyCapabilityVisibility, resolveOutboundConnection, } from './requestHandlerUtils.js'; @@ -39,20 +40,15 @@ export function registerToolHandlers( ListToolsRequestSchema, withErrorHandling(async (request: ListToolsRequest) => { if (lazyLoadingEnabled && lazyLoadingOrchestrator) { - const sessionFilteredConns = filterConnectionsForSession(outboundConns, sessionId); - const capabilityFilteredClients = byCapabilities({ tools: {} })(sessionFilteredConns); - const filteredClients = FilteringService.getFilteredConnections(capabilityFilteredClients, inboundConn); - - const filteredServerNames = new Set(Array.from(filteredClients.values()).map((conn) => conn.name)); + const visibility = resolveLazyCapabilityVisibility(outboundConns, inboundConn, sessionId); + const visibleServerNames = Array.from(new Set(visibility.serverCandidates.values())); infoIf(() => ({ message: 'Lazy loading: filtered servers', meta: { totalOutbound: outboundConns.size, - sessionFiltered: sessionFilteredConns.size, - capabilityFiltered: capabilityFilteredClients.size, - finalFiltered: filteredClients.size, - filteredServerNames: Array.from(filteredServerNames), + finalFiltered: visibility.serverCandidates.size, + visibleServerNames, inboundConfig: { tagFilterMode: inboundConn.tagFilterMode, tags: inboundConn.tags, @@ -61,10 +57,7 @@ export function registerToolHandlers( }, })); - const capabilities = await lazyLoadingOrchestrator.getCapabilitiesForFilteredServers( - filteredServerNames, - sessionId, - ); + const capabilities = await lazyLoadingOrchestrator.getCapabilitiesForVisibility(visibility); const internalProvider = InternalCapabilitiesProvider.getInstance(); await internalProvider.initialize(); @@ -87,7 +80,7 @@ export function registerToolHandlers( const filteredClients = FilteringService.getFilteredConnections(capabilityFilteredClients, inboundConn); const catalog = await createCapabilityCatalogFromConnections(filteredClients, getServerConfigs); - const result = await catalog.listVisibleTools(request.params || {}, sessionId); + const result = await catalog.listVisibleTools(request.params || {}); const internalProvider = InternalCapabilitiesProvider.getInstance(); await internalProvider.initialize(); @@ -121,7 +114,8 @@ export function registerToolHandlers( if (isUnprefixedMetaTool && lazyLoadingOrchestrator) { let result; try { - result = await lazyLoadingOrchestrator.callMetaTool(toolName, request.params.arguments, sessionId); + const visibility = resolveLazyCapabilityVisibility(outboundConns, inboundConn, sessionId); + result = await lazyLoadingOrchestrator.callMetaTool(toolName, request.params.arguments, visibility); } catch (metaToolError) { logger.error(`Meta-tool ${toolName} execution failed: ${metaToolError}`); throw new Error( @@ -146,7 +140,14 @@ export function registerToolHandlers( if (clientName === '1mcp') { const internalProvider = InternalCapabilitiesProvider.getInstance(); await internalProvider.initialize(); - const result = await internalProvider.executeTool(extractedToolName, request.params.arguments, sessionId); + const visibility = lazyLoadingEnabled + ? resolveLazyCapabilityVisibility(outboundConns, inboundConn, sessionId) + : undefined; + const result = await internalProvider.executeTool( + extractedToolName, + request.params.arguments, + visibility, + ); return structuredToolResult(result); } diff --git a/src/core/protocol/toolRequestHandlers.visibility.test.ts b/src/core/protocol/toolRequestHandlers.visibility.test.ts new file mode 100644 index 00000000..a6116ca4 --- /dev/null +++ b/src/core/protocol/toolRequestHandlers.visibility.test.ts @@ -0,0 +1,79 @@ +import { ClientStatus, type InboundConnection, type OutboundConnections } from '@src/core/types/index.js'; +import type { LazyLoadingOrchestrator } from '@src/core/capabilities/lazyLoadingOrchestrator.js'; + +import { describe, expect, it, vi } from 'vitest'; + +vi.mock('@src/core/server/serverManager.js', () => ({ + ServerManager: { + get current() { + return { getTemplateServerManager: () => undefined }; + }, + }, +})); + +import { registerToolHandlers } from './toolRequestHandlers.js'; + +describe('registerToolHandlers capability visibility', () => { + it('re-resolves the Server Candidate Set for each meta-tool request', async () => { + type CapturedHandler = (request: { params: { name: string; arguments: unknown } }) => Promise; + const handlers: CapturedHandler[] = []; + const inbound = { + context: { sessionId: 'session-1' }, + tags: ['safe'], + tagFilterMode: 'simple-or', + server: { + setRequestHandler: vi.fn((_schema, handler) => handlers.push(handler)), + }, + } as unknown as InboundConnection; + const connections = new Map([ + [ + 'ready', + { + name: 'ready', + status: ClientStatus.Connected, + capabilities: { tools: {} }, + transport: { tags: ['safe'] }, + }, + ], + [ + 'late', + { + name: 'late', + status: ClientStatus.Restarting, + capabilities: { tools: {} }, + transport: { tags: ['safe'] }, + }, + ], + [ + 'excluded', + { + name: 'excluded', + status: ClientStatus.Connected, + capabilities: { tools: {} }, + transport: { tags: ['private'] }, + }, + ], + ]) as OutboundConnections; + const callMetaTool = vi.fn().mockResolvedValue({ tools: [] }); + const orchestrator = { + isEnabled: () => true, + isMetaTool: () => true, + callMetaTool, + } as unknown as LazyLoadingOrchestrator; + + registerToolHandlers(connections, inbound, orchestrator); + const callHandler = handlers[1]; + + await callHandler({ params: { name: 'tool_list', arguments: {} } }); + expect(Array.from(callMetaTool.mock.calls[0][2].serverCandidates.entries())).toEqual([['ready', 'ready']]); + + connections.get('late')!.status = ClientStatus.Connected; + await callHandler({ params: { name: 'tool_list', arguments: {} } }); + + expect(Array.from(callMetaTool.mock.calls[1][2].serverCandidates.entries())).toEqual([ + ['ready', 'ready'], + ['late', 'late'], + ]); + expect(callMetaTool.mock.calls[1][2].sessionId).toBe('session-1'); + }); +}); diff --git a/src/core/server/connectionManager.test.ts b/src/core/server/connectionManager.test.ts index 25e1dae1..f36147ad 100644 --- a/src/core/server/connectionManager.test.ts +++ b/src/core/server/connectionManager.test.ts @@ -418,165 +418,19 @@ describe('ConnectionManager', () => { }); }); - describe('initializeSessionFilter - template servers', () => { - it('should use clean names from connection.name for template servers with hash-suffixed keys', async () => { - const sessionId = 'test-session-template'; - - // Mock lazy loading orchestrator - const mockOrchestrator = { - isEnabled: vi.fn().mockReturnValue(true), - getCapabilitiesForFilteredServers: vi.fn().mockResolvedValue(undefined), - }; - connectionManager.setLazyLoadingOrchestrator(mockOrchestrator as any); - - // Add template servers with hash-suffixed keys to outboundConnections - mockOutboundConns.set('serena:abc123', { - name: 'serena', // Clean name - client: {} as any, - capabilities: { tools: {} }, - tags: ['memory'], - } as any); - mockOutboundConns.set('context7:def456', { - name: 'context7', // Clean name - client: {} as any, - capabilities: { tools: {} }, - tags: ['docs'], - } as any); - - const opts = { - tags: ['memory', 'docs'], - enablePagination: false, - }; - - await connectionManager.connectTransport(mockTransport, sessionId, opts); - - // Verify that getCapabilitiesForFilteredServers was called with clean names - expect(mockOrchestrator.getCapabilitiesForFilteredServers).toHaveBeenCalled(); - const [filteredServerNames, calledSessionId] = mockOrchestrator.getCapabilitiesForFilteredServers.mock.calls[0]; - - expect(calledSessionId).toBe(sessionId); - expect(filteredServerNames).toBeInstanceOf(Set); - // Should contain clean names, not hash-suffixed keys - expect(Array.from(filteredServerNames)).toEqual(expect.arrayContaining(['serena', 'context7'])); - expect(Array.from(filteredServerNames)).not.toContain('serena:abc123'); - expect(Array.from(filteredServerNames)).not.toContain('context7:def456'); - }); - - it('should handle mix of template servers and regular servers', async () => { - const sessionId = 'test-session-mixed'; - - const mockOrchestrator = { - isEnabled: vi.fn().mockReturnValue(true), - getCapabilitiesForFilteredServers: vi.fn().mockResolvedValue(undefined), - }; - connectionManager.setLazyLoadingOrchestrator(mockOrchestrator as any); - - // Add mix of template servers (hash-suffixed) and regular servers - mockOutboundConns.set('serena:abc123', { - name: 'serena', // Template server with clean name - client: {} as any, - capabilities: { tools: {} }, - tags: ['memory'], - } as any); - mockOutboundConns.set('filesystem', { - name: 'filesystem', // Regular server - client: {} as any, - capabilities: { tools: {} }, - tags: ['files'], - } as any); - mockOutboundConns.set('postgres:session123', { - name: 'postgres', // Template server with session suffix - client: {} as any, - capabilities: { tools: {} }, - tags: ['database'], - } as any); - - const opts = { - tags: ['memory', 'files', 'database'], - enablePagination: false, - }; - - await connectionManager.connectTransport(mockTransport, sessionId, opts); - - const [filteredServerNames] = mockOrchestrator.getCapabilitiesForFilteredServers.mock.calls[0]; - - // Should contain clean names for all servers - expect(Array.from(filteredServerNames)).toEqual(expect.arrayContaining(['serena', 'filesystem', 'postgres'])); - // Should not contain any suffixed keys - expect(Array.from(filteredServerNames)).not.toContain('serena:abc123'); - expect(Array.from(filteredServerNames)).not.toContain('postgres:session123'); - }); - - it('should filter template servers by tags correctly', async () => { - const sessionId = 'test-session-tag-filter'; - - const mockOrchestrator = { - isEnabled: vi.fn().mockReturnValue(true), - getCapabilitiesForFilteredServers: vi.fn().mockResolvedValue(undefined), - }; - connectionManager.setLazyLoadingOrchestrator(mockOrchestrator as any); - - // Add template servers with different tags - mockOutboundConns.set('serena:abc123', { - name: 'serena', - client: {} as any, - capabilities: { tools: {} }, - tags: ['memory', 'search'], - } as any); - mockOutboundConns.set('context7:def456', { - name: 'context7', - client: {} as any, - capabilities: { tools: {} }, - tags: ['docs', 'reference'], - } as any); - mockOutboundConns.set('filesystem', { - name: 'filesystem', - client: {} as any, - capabilities: { tools: {} }, - tags: ['files', 'storage'], - } as any); - - const opts = { - tags: ['memory'], // Only filter for memory tag - enablePagination: false, - }; - - await connectionManager.connectTransport(mockTransport, sessionId, opts); - - const [filteredServerNames] = mockOrchestrator.getCapabilitiesForFilteredServers.mock.calls[0]; + it('does not materialize lazy allowed-server state when a transport connects', async () => { + const mockOrchestrator = { + isEnabled: vi.fn().mockReturnValue(true), + getCapabilitiesForVisibility: vi.fn(), + }; + connectionManager.setLazyLoadingOrchestrator(mockOrchestrator as any); - // Should only contain servers with memory tag - expect(Array.from(filteredServerNames)).toContain('serena'); - // Should use clean name, not hash-suffixed key - expect(Array.from(filteredServerNames)).not.toContain('serena:abc123'); + await connectionManager.connectTransport(mockTransport, 'session-1', { + tags: ['memory'], + enablePagination: false, }); - it('should not call lazy loading orchestrator when disabled', async () => { - const sessionId = 'test-session-disabled'; - - const mockOrchestrator = { - isEnabled: vi.fn().mockReturnValue(false), // Disabled - getCapabilitiesForFilteredServers: vi.fn().mockResolvedValue(undefined), - }; - connectionManager.setLazyLoadingOrchestrator(mockOrchestrator as any); - - mockOutboundConns.set('serena:abc123', { - name: 'serena', - client: {} as any, - capabilities: { tools: {} }, - tags: ['memory'], - } as any); - - const opts = { - tags: ['memory'], - enablePagination: false, - }; - - await connectionManager.connectTransport(mockTransport, sessionId, opts); - - // Should not call getCapabilitiesForFilteredServers when lazy loading is disabled - expect(mockOrchestrator.getCapabilitiesForFilteredServers).not.toHaveBeenCalled(); - }); + expect(mockOrchestrator.getCapabilitiesForVisibility).not.toHaveBeenCalled(); }); describe('getInboundConnections', () => { diff --git a/src/core/server/connectionManager.ts b/src/core/server/connectionManager.ts index 344ccefc..f2231ae3 100644 --- a/src/core/server/connectionManager.ts +++ b/src/core/server/connectionManager.ts @@ -3,8 +3,6 @@ import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; import { setupCapabilities } from '@src/core/capabilities/capabilityManager.js'; import { LazyLoadingOrchestrator } from '@src/core/capabilities/lazyLoadingOrchestrator.js'; -import { byCapabilities } from '@src/core/filtering/clientFiltering.js'; -import { FilteringService } from '@src/core/filtering/filteringService.js'; import type { OutboundConnections } from '@src/core/types/client.js'; import { InboundConnection, InboundConnectionConfig, OperationOptions, ServerStatus } from '@src/core/types/index.js'; import { @@ -266,11 +264,6 @@ export class ConnectionManager { serverInfo.status = ServerStatus.Connected; serverInfo.lastConnected = new Date(); - // Initialize session filter for lazy loading if enabled - if (this.lazyLoadingOrchestrator?.isEnabled()) { - await this.initializeSessionFilter(sessionId, serverInfo); - } - // Register client with preset notification service if preset is used if (opts.presetName) { await this.registerClientForPresets(sessionId, opts.presetName, serverInfo); @@ -279,44 +272,6 @@ export class ConnectionManager { logger.info(`Connected transport for session ${sessionId}`); } - /** - * Initialize session filter for lazy loading - * This applies tag/preset filtering to the session so that tool_list meta-tool - * returns only tools from filtered servers - */ - private async initializeSessionFilter(sessionId: string, serverInfo: InboundConnection): Promise { - if (!this.lazyLoadingOrchestrator) { - return; - } - - // Apply the same filtering chain as used in tools/list handler - const sessionFilteredConns = this.outboundConns; - const capabilityFilteredConns = byCapabilities({ tools: {} })(sessionFilteredConns); - const filteredConns = FilteringService.getFilteredConnections(capabilityFilteredConns, serverInfo); - - // Get the set of filtered server names using clean names from connection.name - // This ensures consistency with how the tool registry is built - // Template servers use hash-suffixed keys (e.g., "serena:abc123") but - // connection.name contains the clean name (e.g., "serena") - const filteredServerNames = new Set(Array.from(filteredConns.values()).map((conn) => conn.name)); - - debugIf(() => ({ - message: 'Initializing session filter for lazy loading', - meta: { - sessionId, - totalServers: this.outboundConns.size, - filteredServers: filteredServerNames.size, - serverNames: Array.from(filteredServerNames), - mapKeys: Array.from(filteredConns.keys()), // For debugging key mismatch - tagFilterMode: serverInfo.tagFilterMode, - tags: serverInfo.tags, - }, - })); - - // Store the session filter so tool_list returns filtered results - await this.lazyLoadingOrchestrator.getCapabilitiesForFilteredServers(filteredServerNames, sessionId); - } - /** * Register client with preset notification service */ diff --git a/src/core/server/connectionResolver.ts b/src/core/server/connectionResolver.ts index b4de5400..079ac4ec 100644 --- a/src/core/server/connectionResolver.ts +++ b/src/core/server/connectionResolver.ts @@ -136,6 +136,14 @@ export class ConnectionResolver { sessionHashes && sessionHashes.has(identity.templateName) && sessionHashes.get(identity.templateName) === suffix + ) { + filtered.set(key, conn); + continue; + } + + if ( + sessionId && + this.templateHashProvider?.getRenderedHashForSession(sessionId, identity.templateName) === suffix ) { filtered.set(key, conn); } @@ -153,7 +161,7 @@ export class ConnectionResolver { if (!sessionId || !this.templateHashProvider) { return undefined; } - return this.templateHashProvider.getAllRenderedHashesForSession(sessionId); + return this.templateHashProvider.getAllRenderedHashesForSession?.(sessionId); } /** diff --git a/src/core/types/transport.ts b/src/core/types/transport.ts index bf78b590..be34f86c 100644 --- a/src/core/types/transport.ts +++ b/src/core/types/transport.ts @@ -340,7 +340,7 @@ export const applicationConfigSchema = z.object({ lazyLoading: z .object({ enabled: z.boolean().optional().describe('Enable lazy loading for tools'), - mode: z.enum(['full', 'metatool', 'hybrid']).optional().describe('Lazy loading mode'), + mode: z.enum(['full', 'metatool', 'hybrid']).optional().describe('Deprecated and ignored lazy loading mode'), cacheMaxEntries: z.number().int().min(0).optional().describe('Maximum tool schemas to cache'), inlineCatalog: z.boolean().optional().describe('Include full tool catalog in initialize template'), }) diff --git a/src/transport/http/routes/apiRoutes.tool-invocations.test.ts b/src/transport/http/routes/apiRoutes.tool-invocations.test.ts index c62e8823..0ee323c0 100644 --- a/src/transport/http/routes/apiRoutes.tool-invocations.test.ts +++ b/src/transport/http/routes/apiRoutes.tool-invocations.test.ts @@ -431,7 +431,6 @@ describe('apiRoutes /api/tool-invocations', () => { 'tool_invoke', { server: 'alpha', toolName: 'mytool', args: { x: 1 } }, undefined, - undefined, ); expect(res.body).toEqual(mockResult); }); @@ -468,8 +467,7 @@ describe('apiRoutes /api/tool-invocations', () => { expect(callMetaTool).toHaveBeenCalledWith( 'tool_invoke', { server: 'alpha', toolName: 'mytool', args: { x: 1 } }, - 'rest-session-123', - undefined, + expect.objectContaining({ sessionId: 'rest-session-123', serverCandidates: expect.any(Map) }), ); }); }); diff --git a/src/transport/http/routes/apiRoutes.tools.test.ts b/src/transport/http/routes/apiRoutes.tools.test.ts index 99fd551d..03e0be40 100644 --- a/src/transport/http/routes/apiRoutes.tools.test.ts +++ b/src/transport/http/routes/apiRoutes.tools.test.ts @@ -173,6 +173,60 @@ describe('apiRoutes /api/tools', () => { expect(secondBody.nextCursor).toBeUndefined(); }); + it('keeps same-name fallback template instances isolated by the selected connection key', async () => { + const filteredScopeMiddleware: RequestHandler = (_req, res, next) => { + res.locals.validatedTags = ['session-one']; + res.locals.tagFilterMode = 'simple-or'; + next(); + }; + const serverManager = { + getLazyLoadingOrchestrator: vi.fn(() => undefined), + getClients: vi.fn( + () => + new Map([ + [ + 'instance-one', + { + name: 'template', + status: 'connected', + transport: { tags: ['session-one'] }, + client: { + listTools: vi.fn().mockResolvedValue({ + tools: [{ name: 'first_tool', description: 'First instance', inputSchema: {} }], + }), + }, + }, + ], + [ + 'instance-two', + { + name: 'template', + status: 'connected', + transport: { tags: ['session-two'] }, + client: { + listTools: vi.fn().mockResolvedValue({ + tools: [{ name: 'second_tool', description: 'Second instance', inputSchema: {} }], + }), + }, + }, + ], + ]), + ), + }; + const handler = createToolsHandler(serverManager as never); + const res = createMockResponse(); + + await invokeInspectRoute(filteredScopeMiddleware, { query: {} }, res); + await invokeInspectRoute(handler, { query: {} }, res); + + expect(res.statusCode).toBe(200); + expect(res.body).toMatchObject({ + totalCount: 1, + servers: ['template'], + tools: [{ name: 'first_tool', server: 'template', description: 'First instance' }], + }); + }); + it('reports only servers with matching tools in fallback mode', async () => { const serverManager = { getLazyLoadingOrchestrator: vi.fn(() => undefined), @@ -285,7 +339,6 @@ describe('apiRoutes /api/tools', () => { cursor: 'abc', }, undefined, - undefined, ); expect(res.body).toEqual(mockResult); }); @@ -380,7 +433,11 @@ describe('apiRoutes /api/tools', () => { {}, 'ephemeral', ); - expect(callMetaTool).toHaveBeenCalledWith('tool_list', expect.any(Object), 'header-session', undefined); + expect(callMetaTool).toHaveBeenCalledWith( + 'tool_list', + expect.any(Object), + expect.objectContaining({ sessionId: 'header-session', serverCandidates: expect.any(Map) }), + ); expect(res.setHeader).toHaveBeenCalledWith('mcp-session-id', 'header-session'); expect(context.sessionId).toBe('context-session'); }); @@ -440,7 +497,11 @@ describe('apiRoutes /api/tools', () => { ); expect(registerTemplate).toHaveBeenCalledWith('serena', templateConfig); expect(refreshCapabilities).toHaveBeenCalledOnce(); - expect(callMetaTool).toHaveBeenCalledWith('tool_list', expect.any(Object), 'derived-session-id', undefined); + expect(callMetaTool).toHaveBeenCalledWith( + 'tool_list', + expect.any(Object), + expect.objectContaining({ sessionId: 'derived-session-id', serverCandidates: expect.any(Map) }), + ); expect(res.setHeader).toHaveBeenCalledWith('mcp-session-id', 'derived-session-id'); }); diff --git a/src/transport/http/routes/toolRoutes.ts b/src/transport/http/routes/toolRoutes.ts index e86d1888..da26c61b 100644 --- a/src/transport/http/routes/toolRoutes.ts +++ b/src/transport/http/routes/toolRoutes.ts @@ -2,11 +2,16 @@ import type { Tool } from '@modelcontextprotocol/sdk/types.js'; import { McpConfigManager } from '@src/config/mcpConfigManager.js'; import { CapabilityCatalog } from '@src/core/capabilities/capabilityCatalog.js'; +import { + createCapabilityVisibility, + getCapabilityVisibleServerNames, + type CapabilityVisibility, +} from '@src/core/capabilities/capabilityVisibility.js'; import { ToolInvokeOutput, ToolListOutput } from '@src/core/capabilities/schemas/metaToolSchemas.js'; import { ToolRegistry } from '@src/core/capabilities/toolRegistry.js'; import { FilteringService } from '@src/core/filtering/filteringService.js'; import { type ServerAdapter, ServerType } from '@src/core/server/adapters/types.js'; -import type { TemplateHashProvider } from '@src/core/server/connectionResolver.js'; +import { createConnectionResolver, type TemplateHashProvider } from '@src/core/server/connectionResolver.js'; import { getDisabledToolError } from '@src/core/server/disabledTools.js'; import { ServerManager } from '@src/core/server/serverManager.js'; import { ClientStatus } from '@src/core/types/client.js'; @@ -26,20 +31,34 @@ function getServerConfigs() { return McpConfigManager.getInstance().getTransportConfig(); } -function getAllowedServersFromRequest(serverManager: ServerManager, res: Response): Set | undefined { +function getCapabilityVisibilityFromRequest( + serverManager: ServerManager, + res: Response, + sessionId?: string, +): CapabilityVisibility | undefined { const filterConfig = buildFilterConfig(res); - if (filterConfig.tagFilterMode === 'none' && (!filterConfig.tags || filterConfig.tags.length === 0)) { + const hasFilterSelection = + filterConfig.tagFilterMode !== 'none' || (filterConfig.tags !== undefined && filterConfig.tags.length > 0); + if (!hasFilterSelection && !sessionId) { return undefined; } - const allConnections = serverManager.getClients(); - const filteredConnections = FilteringService.getFilteredConnections(allConnections, filterConfig); - return new Set( - Array.from(filteredConnections.entries()).flatMap(([key, connection]) => { - const logicalName = key.includes(':') ? key.split(':')[0] : key; - return connection.name && connection.name !== logicalName - ? [key, logicalName, connection.name] - : [key, logicalName]; + const getClients = (serverManager as { getClients?: () => ReturnType }).getClients; + if (typeof getClients !== 'function') { + return sessionId ? createCapabilityVisibility([], sessionId) : undefined; + } + const allConnections = getClients.call(serverManager); + const sessionScoped = createConnectionResolver(allConnections, getTemplateHashProvider(serverManager)).filterForSession( + sessionId, + ); + const filteredConnections = hasFilterSelection + ? FilteringService.getFilteredConnections(sessionScoped, filterConfig) + : sessionScoped; + return createCapabilityVisibility( + Array.from(filteredConnections.entries(), ([connectionKey, connection]) => { + const publicServerName = connection.name || connectionKey.split(':')[0]; + return [connectionKey, publicServerName] as const; }), + sessionId, ); } @@ -70,8 +89,7 @@ async function createFallbackCapabilityCatalog( serverManager: ServerManager, ): Promise<{ catalog: CapabilityCatalog; degradedServers: string[] }> { const clients = serverManager.getClients(); - const toolsByServer = new Map(); - const serverTags = new Map(); + const registryTools: Array<{ tool: Tool; server: string; connectionKey: string; tags: string[] }> = []; const degradedServers: string[] = []; for (const [connectionKey, conn] of clients) { @@ -80,13 +98,15 @@ async function createFallbackCapabilityCatalog( const logicalServerName = conn.name || (connectionKey.includes(':') ? connectionKey.split(':')[0] : connectionKey); const result = await conn.client.listTools(); - toolsByServer.set(logicalServerName, result.tools ?? []); - const tags = Array.isArray((conn.transport as { tags?: unknown }).tags) - ? ((conn.transport as { tags?: unknown }).tags as unknown[]).filter( + const transportTags = (conn.transport as { tags?: unknown } | undefined)?.tags; + const tags = Array.isArray(transportTags) + ? transportTags.filter( (tag): tag is string => typeof tag === 'string', ) : []; - serverTags.set(logicalServerName, tags); + registryTools.push( + ...(result.tools ?? []).map((tool) => ({ tool, server: logicalServerName, connectionKey, tags })), + ); } catch (err) { logger.error(`Failed to list tools for ${connectionKey}:`, err); degradedServers.push(connectionKey); @@ -94,7 +114,7 @@ async function createFallbackCapabilityCatalog( } const catalog = new CapabilityCatalog({ - getToolRegistry: () => ToolRegistry.fromToolsMap(toolsByServer, serverTags), + getToolRegistry: () => ToolRegistry.fromToolsWithServer(registryTools), schemaCache: { getIfCached: () => null, getOrLoad: async (_server: string, _toolName: string) => { @@ -129,7 +149,7 @@ export function createToolsHandler(serverManager: ServerManager): RequestHandler const limit = limitParam !== undefined && Number.isFinite(limitParam) && limitParam > 0 ? limitParam : undefined; const requestSessionId = await initializeRequestContextForApi(serverManager, req, res); - const allowedServers = getAllowedServersFromRequest(serverManager, res); + const visibility = getCapabilityVisibilityFromRequest(serverManager, res, requestSessionId); const lazyOrchestrator = serverManager.getLazyLoadingOrchestrator(); if (!lazyOrchestrator) { @@ -142,8 +162,7 @@ export function createToolsHandler(serverManager: ServerManager): RequestHandler limit, cursor, }, - requestSessionId, - allowedServers, + visibility, ); res.json({ @@ -173,8 +192,7 @@ export function createToolsHandler(serverManager: ServerManager): RequestHandler limit, cursor, }, - requestSessionId, - allowedServers, + visibility, ); if (catalogResult.tools.length > 0 || catalogResult.totalCount > 0) { res.json({ @@ -196,8 +214,7 @@ export function createToolsHandler(serverManager: ServerManager): RequestHandler limit, cursor, }, - requestSessionId, - allowedServers, + visibility, )) as ToolListOutput; if (result.error) { @@ -240,12 +257,16 @@ export function createToolInvocationsHandler(serverManager: ServerManager): Requ return; } - const allowedServers = getAllowedServersFromRequest(serverManager, res); + const visibility = getCapabilityVisibilityFromRequest(serverManager, res, requestSessionId); + const visibleServerNames = visibility ? getCapabilityVisibleServerNames(visibility) : undefined; + const filterConfig = buildFilterConfig(res); + const hasFilterSelection = + filterConfig.tagFilterMode !== 'none' || (filterConfig.tags !== undefined && filterConfig.tags.length > 0); const lazyOrchestrator = serverManager.getLazyLoadingOrchestrator(); if (!lazyOrchestrator) { - if (allowedServers && !allowedServers.has(target.serverName)) { + if (hasFilterSelection && visibleServerNames && !visibleServerNames.has(target.serverName)) { res.status(404).json({ error: `Server not found: ${target.serverName}` }); return; } @@ -255,10 +276,7 @@ export function createToolInvocationsHandler(serverManager: ServerManager): Requ return; } const allConnections = serverManager.getClients(); - const filteredConnections = - allowedServers === undefined - ? allConnections - : FilteringService.getFilteredConnections(allConnections, buildFilterConfig(res)); + const filteredConnections = FilteringService.getFilteredConnections(allConnections, buildFilterConfig(res)); const serverRegistry = getServerRegistry(serverManager); const sessionConnection = requestSessionId ? serverRegistry?.resolveConnection?.(target.serverName, { sessionId: requestSessionId }) @@ -287,7 +305,7 @@ export function createToolInvocationsHandler(serverManager: ServerManager): Requ return; } - if (!allowedServers || allowedServers.has(target.serverName)) { + if (!hasFilterSelection || !visibleServerNames || visibleServerNames.has(target.serverName)) { const disabledToolError = getDisabledToolInvocationError(target.serverName, target.toolName); if (disabledToolError) { res.status(404).json({ error: disabledToolError }); @@ -305,8 +323,7 @@ export function createToolInvocationsHandler(serverManager: ServerManager): Requ }); const catalogResult = await catalog.invokeVisibleTool( { server: target.serverName, toolName: target.toolName, args: toolArgs }, - requestSessionId, - allowedServers, + visibility, ); if (!catalogResult.error) { res.json({ result: catalogResult.result, server: catalogResult.server, tool: catalogResult.tool }); @@ -325,8 +342,7 @@ export function createToolInvocationsHandler(serverManager: ServerManager): Requ toolName: target.toolName, args: toolArgs, }, - requestSessionId, - allowedServers, + visibility, )) as ToolInvokeOutput; if (result.error) {