From 4680ddf3c066b81b944411154ecdfec0781a854c Mon Sep 17 00:00:00 2001 From: alban bertolini Date: Fri, 16 Jan 2026 21:01:02 +0100 Subject: [PATCH 1/2] refactor(ai-proxy): make RemoteTool abstract and add McpServerRemoteTool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Make RemoteTool an abstract base class - Add McpServerRemoteTool for MCP server tools (sourceType: 'mcp-server') - Update mcp-client to use McpServerRemoteTool - Update tests accordingly 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- packages/ai-proxy/src/mcp-client.ts | 13 ++++--------- packages/ai-proxy/src/mcp-server-remote-tool.ts | 11 +++++++++++ packages/ai-proxy/src/remote-tool.ts | 14 ++++---------- packages/ai-proxy/test/mcp-client.test.ts | 8 +++----- 4 files changed, 22 insertions(+), 24 deletions(-) create mode 100644 packages/ai-proxy/src/mcp-server-remote-tool.ts diff --git a/packages/ai-proxy/src/mcp-client.ts b/packages/ai-proxy/src/mcp-client.ts index 760df5de87..e71250cc05 100644 --- a/packages/ai-proxy/src/mcp-client.ts +++ b/packages/ai-proxy/src/mcp-client.ts @@ -3,7 +3,7 @@ import type { Logger } from '@forestadmin/datasource-toolkit'; import { MultiServerMCPClient } from '@langchain/mcp-adapters'; import { McpConnectionError } from './errors'; -import RemoteTool from './remote-tool'; +import McpServerRemoteTool from './mcp-server-remote-tool'; export type McpConfiguration = { configs: MultiServerMCPClient['config']['mcpServers']; @@ -13,7 +13,7 @@ export default class McpClient { private readonly mcpClients: Record = {}; private readonly logger?: Logger; - readonly tools: RemoteTool[] = []; + readonly tools: McpServerRemoteTool[] = []; constructor(config: McpConfiguration, logger?: Logger) { this.logger = logger; @@ -27,7 +27,7 @@ export default class McpClient { }); } - async loadTools(): Promise { + async loadTools(): Promise { const errors: Array<{ server: string; error: Error }> = []; await Promise.all( @@ -35,12 +35,7 @@ export default class McpClient { try { const tools = (await client.getTools()) ?? []; const extendedTools = tools.map( - tool => - new RemoteTool({ - tool, - sourceId: name, - sourceType: 'mcp-server', - }), + tool => new McpServerRemoteTool({ tool, sourceId: name }), ); this.tools.push(...extendedTools); } catch (error) { diff --git a/packages/ai-proxy/src/mcp-server-remote-tool.ts b/packages/ai-proxy/src/mcp-server-remote-tool.ts new file mode 100644 index 0000000000..ef19a29bbe --- /dev/null +++ b/packages/ai-proxy/src/mcp-server-remote-tool.ts @@ -0,0 +1,11 @@ +import type { StructuredToolInterface } from '@langchain/core/tools'; + +import RemoteTool from './remote-tool'; + +export default class McpServerRemoteTool extends RemoteTool { + constructor(options: { tool: StructuredToolInterface; sourceId?: string }) { + super({ ...options, sourceType: 'mcp-server' }); + this.base = options.tool; + this.sourceId = options.sourceId; + } +} diff --git a/packages/ai-proxy/src/remote-tool.ts b/packages/ai-proxy/src/remote-tool.ts index 5e10a748eb..9c31f3f1d1 100644 --- a/packages/ai-proxy/src/remote-tool.ts +++ b/packages/ai-proxy/src/remote-tool.ts @@ -1,16 +1,14 @@ import type { StructuredToolInterface } from '@langchain/core/tools'; -export type SourceType = 'server' | 'mcp-server'; - -export default class RemoteTool { +export default abstract class RemoteTool { base: StructuredToolInterface; sourceId: string; - sourceType: SourceType; + sourceType: string; constructor(options: { tool: StructuredToolInterface; sourceId?: string; - sourceType?: SourceType; + sourceType?: string; }) { this.base = options.tool; this.sourceId = options.sourceId; @@ -18,12 +16,8 @@ export default class RemoteTool { } get sanitizedName() { - return this.sanitizeName(this.base.name); - } - - private sanitizeName(name: string): string { // OpenAI function names must be alphanumeric and can contain underscores // This function replaces non-alphanumeric characters with underscores - return name.replace(/[^a-zA-Z0-9_-]/g, '_'); + return this.base.name.replace(/[^a-zA-Z0-9_-]/g, '_'); } } diff --git a/packages/ai-proxy/test/mcp-client.test.ts b/packages/ai-proxy/test/mcp-client.test.ts index 7a350f8ed9..6e52883511 100644 --- a/packages/ai-proxy/test/mcp-client.test.ts +++ b/packages/ai-proxy/test/mcp-client.test.ts @@ -4,7 +4,7 @@ import { tool } from '@langchain/core/tools'; import { McpConnectionError } from '../src'; import McpClient from '../src/mcp-client'; -import RemoteTool from '../src/remote-tool'; +import McpServerRemoteTool from '../src/mcp-server-remote-tool'; const getToolsMock = jest.fn(); const closeMock = jest.fn(); @@ -56,15 +56,13 @@ describe('McpClient', () => { await mcpClient.loadTools(); expect(mcpClient.tools).toEqual([ - new RemoteTool({ + new McpServerRemoteTool({ tool: tool1, sourceId: 'slack', - sourceType: 'mcp-server', }), - new RemoteTool({ + new McpServerRemoteTool({ tool: tool2, sourceId: 'slack', - sourceType: 'mcp-server', }), ]); }); From 11fba7118792538974ab2b2a2c4fed9ef4c9bb21 Mon Sep 17 00:00:00 2001 From: alban bertolini Date: Fri, 16 Jan 2026 21:05:05 +0100 Subject: [PATCH 2/2] refactor(ai-proxy): add ServerRemoteTool and ServerRemoteToolBuilder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add ServerRemoteTool for server-side tools (sourceType: 'server') - Extract tool building logic into ServerRemoteToolBuilder - Update RemoteTools to use ServerRemoteToolBuilder - Update router types - Update tests accordingly 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- packages/ai-proxy/src/remote-tools.ts | 25 ++++------------- packages/ai-proxy/src/router.ts | 5 ++-- .../src/server-remote-tool-builder.ts | 28 +++++++++++++++++++ packages/ai-proxy/src/server-remote-tool.ts | 11 ++++++++ packages/ai-proxy/test/remote-tools.test.ts | 7 +++-- 5 files changed, 51 insertions(+), 25 deletions(-) create mode 100644 packages/ai-proxy/src/server-remote-tool-builder.ts create mode 100644 packages/ai-proxy/src/server-remote-tool.ts diff --git a/packages/ai-proxy/src/remote-tools.ts b/packages/ai-proxy/src/remote-tools.ts index b16dacb06e..c89c7b38a2 100644 --- a/packages/ai-proxy/src/remote-tools.ts +++ b/packages/ai-proxy/src/remote-tools.ts @@ -1,35 +1,20 @@ +import type RemoteTool from './remote-tool'; +import type { ServerRemoteToolsApiKeys } from './server-remote-tool-builder'; import type { ResponseFormat } from '@langchain/core/tools'; import type { ChatCompletionCreateParamsNonStreaming } from 'openai/resources/chat/completions'; -import { BraveSearch } from '@langchain/community/tools/brave_search'; import { toJsonSchema } from '@langchain/core/utils/json_schema'; import { AIToolNotFoundError, AIToolUnprocessableError } from './errors'; -import RemoteTool from './remote-tool'; +import ServerRemoteToolBuilder from './server-remote-tool-builder'; export type Messages = ChatCompletionCreateParamsNonStreaming['messages']; -export type RemoteToolsApiKeys = - | { ['AI_REMOTE_TOOL_BRAVE_SEARCH_API_KEY']: string } - | Record; // To avoid to cast the object because env is not always well typed from the caller - export class RemoteTools { - private readonly apiKeys?: RemoteToolsApiKeys; readonly tools: RemoteTool[] = []; - constructor(apiKeys: RemoteToolsApiKeys, tools?: RemoteTool[]) { - this.apiKeys = apiKeys; - this.tools.push(...(tools ?? [])); - - if (this.apiKeys?.AI_REMOTE_TOOL_BRAVE_SEARCH_API_KEY) { - this.tools.push( - new RemoteTool({ - sourceId: 'brave_search', - sourceType: 'server', - tool: new BraveSearch({ apiKey: this.apiKeys.AI_REMOTE_TOOL_BRAVE_SEARCH_API_KEY }), - }), - ); - } + constructor(apiKeys: ServerRemoteToolsApiKeys, tools?: RemoteTool[]) { + this.tools.push(...ServerRemoteToolBuilder.buildTools(apiKeys), ...(tools ?? [])); } get toolDefinitionsForFrontend() { diff --git a/packages/ai-proxy/src/router.ts b/packages/ai-proxy/src/router.ts index ced1e9f2b2..25ec97407e 100644 --- a/packages/ai-proxy/src/router.ts +++ b/packages/ai-proxy/src/router.ts @@ -1,6 +1,7 @@ import type { McpConfiguration } from './mcp-client'; import type { AiConfiguration, DispatchBody } from './provider-dispatcher'; -import type { Messages, RemoteToolsApiKeys } from './remote-tools'; +import type { Messages } from './remote-tools'; +import type { ServerRemoteToolsApiKeys } from './server-remote-tool-builder'; import type { Logger } from '@forestadmin/datasource-toolkit'; import { AIUnprocessableError, ProviderDispatcher } from './index'; @@ -13,7 +14,7 @@ export type Route = 'ai-query' | 'remote-tools' | 'invoke-remote-tool'; export type Query = { 'tool-name'?: string; }; -export type ApiKeys = RemoteToolsApiKeys; +export type ApiKeys = ServerRemoteToolsApiKeys; export class Router { private readonly localToolsApiKeys?: ApiKeys; diff --git a/packages/ai-proxy/src/server-remote-tool-builder.ts b/packages/ai-proxy/src/server-remote-tool-builder.ts new file mode 100644 index 0000000000..3d3d9b2b5d --- /dev/null +++ b/packages/ai-proxy/src/server-remote-tool-builder.ts @@ -0,0 +1,28 @@ +import { BraveSearch } from '@langchain/community/tools/brave_search'; + +import ServerRemoteTool from './server-remote-tool'; + +export type BraveApiKey = { ['AI_REMOTE_TOOL_BRAVE_SEARCH_API_KEY']: string }; + +export type ServerRemoteToolsApiKeys = BraveApiKey | Record; // To avoid to cast the object because env is not always well typed from the caller + +export default class ServerRemoteToolBuilder { + static buildTools(apiKeys?: ServerRemoteToolsApiKeys) { + const tools: ServerRemoteTool[] = []; + this.addBrave(apiKeys, tools); + // TODO: slack, infogreffe etc. tools will be built here. + + return tools; + } + + private static addBrave(apiKeys: ServerRemoteToolsApiKeys, tools: ServerRemoteTool[]) { + if (apiKeys?.AI_REMOTE_TOOL_BRAVE_SEARCH_API_KEY) { + tools.push( + new ServerRemoteTool({ + sourceId: 'brave_search', + tool: new BraveSearch({ apiKey: apiKeys.AI_REMOTE_TOOL_BRAVE_SEARCH_API_KEY }), + }), + ); + } + } +} diff --git a/packages/ai-proxy/src/server-remote-tool.ts b/packages/ai-proxy/src/server-remote-tool.ts new file mode 100644 index 0000000000..2f39e757cc --- /dev/null +++ b/packages/ai-proxy/src/server-remote-tool.ts @@ -0,0 +1,11 @@ +import type { StructuredToolInterface } from '@langchain/core/tools'; + +import RemoteTool from './remote-tool'; + +export default class ServerRemoteTool extends RemoteTool { + constructor(options: { tool: StructuredToolInterface; sourceId?: string }) { + super({ ...options, sourceType: 'server' }); + this.base = options.tool; + this.sourceId = options.sourceId; + } +} diff --git a/packages/ai-proxy/test/remote-tools.test.ts b/packages/ai-proxy/test/remote-tools.test.ts index 3a4d9a98d7..899238f32e 100644 --- a/packages/ai-proxy/test/remote-tools.test.ts +++ b/packages/ai-proxy/test/remote-tools.test.ts @@ -4,7 +4,7 @@ import type { JSONSchema } from '@langchain/core/utils/json_schema'; import { toJsonSchema } from '@langchain/core/utils/json_schema'; import { RemoteTools } from '../src'; -import RemoteTool from '../src/remote-tool'; +import McpServerRemoteTool from '../src/mcp-server-remote-tool'; describe('RemoteTools', () => { const apiKeys = { AI_REMOTE_TOOL_BRAVE_SEARCH_API_KEY: 'api-key' }; @@ -31,7 +31,7 @@ describe('RemoteTools', () => { describe('when tools are passed in the constructor', () => { it('should return the tools', () => { const tools = [ - new RemoteTool({ + new McpServerRemoteTool({ tool: { name: 'tool1', description: 'description1', @@ -42,7 +42,8 @@ describe('RemoteTools', () => { ]; const remoteTools = new RemoteTools(apiKeys, tools); expect(remoteTools.tools.length).toEqual(2); - expect(remoteTools.tools[0].base.name).toEqual('tool1'); + // Server tools (brave-search) are added first, then passed tools + expect(remoteTools.tools[1].base.name).toEqual('tool1'); }); }); });