Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions packages/ai-proxy/src/forest-integration-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ export type CustomConfig = ZendeskConfig | KolarConfig | SnowflakeConfig;
export type ForestIntegrationName = 'Zendesk' | 'Kolar' | 'Snowflake';

export interface ForestIntegrationConfig {
id?: string;
integrationName: ForestIntegrationName;
config: CustomConfig;
isForestConnector: true;
Expand All @@ -39,16 +40,16 @@ export default class ForestIntegrationClient implements ToolProvider {
async loadTools(): Promise<RemoteTool[]> {
const tools: RemoteTool[] = [];

this.configs.forEach(({ integrationName, config }) => {
this.configs.forEach(({ id: mcpServerId, integrationName, config }) => {
switch (integrationName) {
case 'Zendesk':
tools.push(...getZendeskTools(config as ZendeskConfig));
tools.push(...getZendeskTools(config as ZendeskConfig, mcpServerId));
break;
case 'Kolar':
tools.push(...getKolarTools(config as KolarConfig));
tools.push(...getKolarTools(config as KolarConfig, mcpServerId));
break;
case 'Snowflake':
tools.push(...getSnowflakeTools(config as SnowflakeConfig));
tools.push(...getSnowflakeTools(config as SnowflakeConfig, mcpServerId));
break;
default:
this.logger?.('Warn', `Unsupported integration: ${integrationName}`);
Expand Down
3 changes: 2 additions & 1 deletion packages/ai-proxy/src/integrations/kolar/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ export interface KolarConfig {
apiKey: string;
}

export default function getKolarTools(config: KolarConfig): RemoteTool[] {
export default function getKolarTools(config: KolarConfig, mcpServerId?: string): RemoteTool[] {
const { baseUrl, headers } = getKolarConfig(config);

return [
Expand All @@ -23,6 +23,7 @@ export default function getKolarTools(config: KolarConfig): RemoteTool[] {
tool =>
new ServerRemoteTool({
sourceId: 'kolar',
mcpServerId,
tool,
}),
);
Expand Down
6 changes: 5 additions & 1 deletion packages/ai-proxy/src/integrations/snowflake/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,10 @@ export interface SnowflakeConfig {
defaultRole?: string;
}

export default function getSnowflakeTools(config: SnowflakeConfig): RemoteTool[] {
export default function getSnowflakeTools(
config: SnowflakeConfig,
mcpServerId?: string,
): RemoteTool[] {
const headers = getSnowflakeAuthHeaders(config);
const baseUrl = getSnowflakeBaseUrl(config.accountIdentifier);

Expand All @@ -27,6 +30,7 @@ export default function getSnowflakeTools(config: SnowflakeConfig): RemoteTool[]
tool =>
new ServerRemoteTool({
sourceId: 'snowflake',
mcpServerId,
tool,
}),
);
Expand Down
3 changes: 2 additions & 1 deletion packages/ai-proxy/src/integrations/zendesk/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ export interface ZendeskConfig {
apiToken: string;
}

export default function getZendeskTools(config: ZendeskConfig): RemoteTool[] {
export default function getZendeskTools(config: ZendeskConfig, mcpServerId?: string): RemoteTool[] {
const { baseUrl, headers } = getZendeskConfig(config);

return [
Expand All @@ -29,6 +29,7 @@ export default function getZendeskTools(config: ZendeskConfig): RemoteTool[] {
tool =>
new ServerRemoteTool({
sourceId: 'zendesk',
mcpServerId,
tool,
}),
);
Expand Down
19 changes: 15 additions & 4 deletions packages/ai-proxy/src/mcp-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,23 +8,29 @@ import McpServerRemoteTool from './mcp-server-remote-tool';

export type McpServers = MultiServerMCPClient['config']['mcpServers'];

export type McpServerConfig = MultiServerMCPClient['config']['mcpServers'][string];
export type McpServerConfig = MultiServerMCPClient['config']['mcpServers'][string] & {
id?: string;
};

export type McpConfiguration = {
configs: McpServers;
configs: Record<string, McpServerConfig>;
} & Omit<MultiServerMCPClient['config'], 'mcpServers'>;

export default class McpClient implements ToolProvider {
private readonly mcpClients: Record<string, MultiServerMCPClient> = {};
private readonly mcpServerIdsByName: Record<string, string | undefined> = {};
private readonly logger?: Logger;

constructor(config: McpConfiguration, logger?: Logger) {
this.logger = logger;
// split the config into several clients to be more resilient
// if a mcp server is down, the others will still work
Object.entries(config.configs).forEach(([name, serverConfig]) => {
const { id: mcpServerId, ...rest } = serverConfig as McpServerConfig &
Record<string, unknown>;
this.mcpServerIdsByName[name] = mcpServerId;
this.mcpClients[name] = new MultiServerMCPClient({
mcpServers: { [name]: serverConfig },
mcpServers: { [name]: rest as McpServerConfig },
...config,
});
});
Expand All @@ -39,7 +45,12 @@ export default class McpClient implements ToolProvider {
try {
const loadedTools = (await client.getTools()) ?? [];
const extendedTools = loadedTools.map(
tool => new McpServerRemoteTool({ tool, sourceId: name }),
tool =>
new McpServerRemoteTool({
tool,
sourceId: name,
mcpServerId: this.mcpServerIdsByName[name],
}),
);
tools.push(...extendedTools);
} catch (error) {
Expand Down
6 changes: 5 additions & 1 deletion packages/ai-proxy/src/mcp-server-remote-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,11 @@ import type { StructuredToolInterface } from '@langchain/core/tools';
import RemoteTool from './remote-tool';

export default class McpServerRemoteTool<ToolType = unknown> extends RemoteTool {
constructor(options: { tool: StructuredToolInterface<ToolType>; sourceId?: string }) {
constructor(options: {
tool: StructuredToolInterface<ToolType>;
sourceId?: string;
mcpServerId?: string;
}) {
super({ ...options, sourceType: 'mcp-server' });
}
}
5 changes: 5 additions & 0 deletions packages/ai-proxy/src/remote-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,20 @@ export default abstract class RemoteTool<ToolType = unknown> {
base: StructuredToolInterface<ToolType>;
sourceId: string;
sourceType: string;
// Shared across McpServerRemoteTool and ServerRemoteTool because the orchestrator
// stores Forest connectors alongside user MCP servers in the same ai_mcp_configs table.
mcpServerId?: string;

constructor(options: {
tool: StructuredToolInterface<ToolType>;
sourceId?: string;
sourceType?: string;
mcpServerId?: string;
}) {
this.base = options.tool;
this.sourceId = options.sourceId;
this.sourceType = options.sourceType;
this.mcpServerId = options.mcpServerId;
}

get sanitizedName() {
Expand Down
6 changes: 5 additions & 1 deletion packages/ai-proxy/src/server-remote-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,11 @@ import type { StructuredToolInterface } from '@langchain/core/tools';
import RemoteTool from './remote-tool';

export default class ServerRemoteTool<ToolType = unknown> extends RemoteTool {
constructor(options: { tool: StructuredToolInterface<ToolType>; sourceId?: string }) {
constructor(options: {
tool: StructuredToolInterface<ToolType>;
sourceId?: string;
mcpServerId?: string;
}) {
super({ ...options, sourceType: 'server' });
}
}
96 changes: 88 additions & 8 deletions packages/ai-proxy/test/forest-integration-client.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
import ForestIntegrationClient from '../src/forest-integration-client';
import getKolarTools from '../src/integrations/kolar/tools';
import { validateKolarConfig } from '../src/integrations/kolar/utils';
import getSnowflakeTools from '../src/integrations/snowflake/tools';
import { validateSnowflakeConfig } from '../src/integrations/snowflake/utils';
import getZendeskTools from '../src/integrations/zendesk/tools';
import { validateZendeskConfig } from '../src/integrations/zendesk/utils';

const mockZendeskTools = [{ name: 'zendesk_get_tickets' }, { name: 'zendesk_get_ticket' }];
Expand Down Expand Up @@ -37,6 +40,7 @@ describe('ForestIntegrationClient', () => {
it('should load zendesk tools when integration is zendesk', async () => {
const client = new ForestIntegrationClient([
{
id: '1',
integrationName: 'Zendesk',
config: { subdomain: 'test', email: '[email protected]', apiToken: 'tok' },
isForestConnector: true,
Expand All @@ -52,7 +56,7 @@ describe('ForestIntegrationClient', () => {
const logger = jest.fn();
const client = new ForestIntegrationClient(
// @ts-expect-error Testing unsupported integration
[{ integrationName: 'unknown', config: {} as any, isForestConnector: true }],
[{ id: '1', integrationName: 'unknown', config: {} as any, isForestConnector: true }],
logger,
);

Expand All @@ -64,6 +68,7 @@ describe('ForestIntegrationClient', () => {
it('should load kolar tools when integration is Kolar', async () => {
const client = new ForestIntegrationClient([
{
id: '1',
integrationName: 'Kolar',
config: { apiKey: 'key' },
isForestConnector: true,
Expand All @@ -78,6 +83,7 @@ describe('ForestIntegrationClient', () => {
it('should load snowflake tools when integration is Snowflake', async () => {
const client = new ForestIntegrationClient([
{
id: '1',
integrationName: 'Snowflake',
config: {
accountIdentifier: 'a',
Expand All @@ -101,11 +107,13 @@ describe('ForestIntegrationClient', () => {
it('should load tools from multiple configs', async () => {
const client = new ForestIntegrationClient([
{
id: '1',
integrationName: 'Zendesk',
config: { subdomain: 'a', email: '[email protected]', apiToken: 'tok' },
isForestConnector: true,
},
{
id: '2',
integrationName: 'Zendesk',
config: { subdomain: 'b', email: '[email protected]', apiToken: 'tok2' },
isForestConnector: true,
Expand All @@ -122,7 +130,7 @@ describe('ForestIntegrationClient', () => {
it('should call validateZendeskConfig for Zendesk integration', async () => {
const zendeskConfig = { subdomain: 'test', email: '[email protected]', apiToken: 'tok' };
const client = new ForestIntegrationClient([
{ integrationName: 'Zendesk', config: zendeskConfig, isForestConnector: true },
{ id: '1', integrationName: 'Zendesk', config: zendeskConfig, isForestConnector: true },
]);

await client.checkConnection();
Expand All @@ -133,6 +141,7 @@ describe('ForestIntegrationClient', () => {
it('should return true on success', async () => {
const client = new ForestIntegrationClient([
{
id: '1',
integrationName: 'Zendesk',
config: { subdomain: 'test', email: '[email protected]', apiToken: 'tok' },
isForestConnector: true,
Expand All @@ -147,7 +156,7 @@ describe('ForestIntegrationClient', () => {
it('should call validateKolarConfig for Kolar integration', async () => {
const kolarConfig = { apiKey: 'key' };
const client = new ForestIntegrationClient([
{ integrationName: 'Kolar', config: kolarConfig, isForestConnector: true },
{ id: '1', integrationName: 'Kolar', config: kolarConfig, isForestConnector: true },
]);

await client.checkConnection();
Expand All @@ -161,7 +170,7 @@ describe('ForestIntegrationClient', () => {
programmaticAccessToken: 'tok',
};
const client = new ForestIntegrationClient([
{ integrationName: 'Snowflake', config: snowflakeConfig, isForestConnector: true },
{ id: '1', integrationName: 'Snowflake', config: snowflakeConfig, isForestConnector: true },
]);

await client.checkConnection();
Expand All @@ -172,12 +181,10 @@ describe('ForestIntegrationClient', () => {
it('should throw for unsupported integration', async () => {
const client = new ForestIntegrationClient([
// @ts-expect-error Testing unsupported integration
{ integrationName: 'Unknown', config: {}, isForestConnector: true },
{ id: '1', integrationName: 'Unknown', config: {}, isForestConnector: true },
]);

await expect(client.checkConnection()).rejects.toThrow(
'Unsupported integration: Unknown',
);
await expect(client.checkConnection()).rejects.toThrow('Unsupported integration: Unknown');
});
});

Expand All @@ -188,4 +195,77 @@ describe('ForestIntegrationClient', () => {
await expect(client.dispose()).resolves.toBeUndefined();
});
});

describe('ForestIntegrationConfig.id is threaded as mcpServerId into integration tool factories', () => {
it('passes the config id to getZendeskTools so produced tools can be matched by step.mcpServerId', async () => {
const client = new ForestIntegrationClient([
{
id: 'forest-zendesk-42',
integrationName: 'Zendesk',
config: { subdomain: 'test', email: '[email protected]', apiToken: 'tok' },
isForestConnector: true,
},
]);

await client.loadTools();

expect(getZendeskTools).toHaveBeenCalledWith(
expect.objectContaining({ subdomain: 'test' }),
'forest-zendesk-42',
);
});

it('passes the config id to getKolarTools', async () => {
const client = new ForestIntegrationClient([
{
id: 'forest-kolar-7',
integrationName: 'Kolar',
config: { apiKey: 'key' },
isForestConnector: true,
},
]);

await client.loadTools();

expect(getKolarTools).toHaveBeenCalledWith(
expect.objectContaining({ apiKey: 'key' }),
'forest-kolar-7',
);
});

it('passes the config id to getSnowflakeTools', async () => {
const client = new ForestIntegrationClient([
{
id: 'forest-snowflake-99',
integrationName: 'Snowflake',
config: { accountIdentifier: 'a', programmaticAccessToken: 'tok' },
isForestConnector: true,
},
]);

await client.loadTools();

expect(getSnowflakeTools).toHaveBeenCalledWith(
expect.objectContaining({ accountIdentifier: 'a' }),
'forest-snowflake-99',
);
});

it('passes undefined to the factory when the config entry has no id', async () => {
const client = new ForestIntegrationClient([
{
integrationName: 'Zendesk',
config: { subdomain: 'test', email: '[email protected]', apiToken: 'tok' },
isForestConnector: true,
},
]);

await client.loadTools();

expect(getZendeskTools).toHaveBeenCalledWith(
expect.objectContaining({ subdomain: 'test' }),
undefined,
);
});
});
});
19 changes: 19 additions & 0 deletions packages/ai-proxy/test/integrations/kolar/tools.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,4 +26,23 @@ describe('getKolarTools', () => {
'kolar_get_screening_result',
]);
});

describe('mcpServerId propagation', () => {
it('sets RemoteTool.mcpServerId on every produced tool when mcpServerId is provided', () => {
const tools = getKolarTools(config, 'forest-kolar-7');

expect(tools).not.toHaveLength(0);
tools.forEach(tool => {
expect(tool.mcpServerId).toBe('forest-kolar-7');
});
});

it('leaves RemoteTool.mcpServerId undefined when no mcpServerId is provided', () => {
const tools = getKolarTools(config);

tools.forEach(tool => {
expect(tool.mcpServerId).toBeUndefined();
});
});
});
});
Loading
Loading