From 9dcec69fd52fb16827e8eed562031e3563581004 Mon Sep 17 00:00:00 2001 From: Acbox Date: Fri, 6 Feb 2026 21:29:28 +0800 Subject: [PATCH] feat: add mcp support --- agent/src/agent.ts | 39 ++++++++++++++++++++++++------ agent/src/models.ts | 27 ++++++++++++++++++++- agent/src/modules/chat.ts | 38 +++++++++++++++-------------- agent/src/tools/mcp.ts | 51 +++++++++++++++++++++++++++++++++++++++ agent/src/types/agent.ts | 2 ++ agent/src/types/index.ts | 3 ++- agent/src/types/mcp.ts | 29 ++++++++++++++++++++++ 7 files changed, 162 insertions(+), 27 deletions(-) create mode 100644 agent/src/tools/mcp.ts create mode 100644 agent/src/types/mcp.ts diff --git a/agent/src/agent.ts b/agent/src/agent.ts index 21d94877..aa98e6fb 100644 --- a/agent/src/agent.ts +++ b/agent/src/agent.ts @@ -12,6 +12,7 @@ import { AttachmentsStreamExtractor, } from './utils/attachments' import type { ContainerFileAttachment } from './types/attachment' +import { getMCPTools } from './tools/mcp' export const createAgent = ({ model: modelConfig, @@ -21,6 +22,7 @@ export const createAgent = ({ allowedActions = allActions, identity, channels = [], + mcpConnections = [], currentChannel = 'Unknown Channel', }: AgentParams, fetch: AuthFetcher) => { const model = createModel(modelConfig) @@ -36,12 +38,20 @@ export const createAgent = ({ }) } - const tools = getTools(allowedActions, { - fetch, - model: modelConfig, - brave, - identity, - }) + const getAgentTools = async () => { + const tools = getTools(allowedActions, { + fetch, + model: modelConfig, + brave, + identity, + }) + const { tools: mcpTools, close: closeMCP } = await getMCPTools(mcpConnections) + Object.assign(tools, mcpTools) + return { + tools, + close: closeMCP, + } + } const generateUserPrompt = (input: AgentInput) => { const images = input.attachments.filter(attachment => attachment.type === 'image') @@ -67,6 +77,7 @@ export const createAgent = ({ const userPrompt = generateUserPrompt(input) const messages = [...input.messages, userPrompt] const systemPrompt = generateSystemPrompt() + const { tools, close } = await getAgentTools() const { response, reasoning, text, usage } = await generateText({ model, messages, @@ -77,6 +88,9 @@ export const createAgent = ({ system: systemPrompt, } }, + onFinish: async () => { + await close() + }, tools, }) const { cleanedText, attachments: textAttachments } = extractAttachmentsFromText(text) @@ -111,6 +125,7 @@ export const createAgent = ({ }) } const messages = [...params.messages, userPrompt] + const { tools, close } = await getAgentTools() const { response, reasoning, text, usage } = await generateText({ model, messages, @@ -121,6 +136,9 @@ export const createAgent = ({ system: generateSubagentSystemPrompt(), } }, + onFinish: async () => { + await close() + }, tools, }) return { @@ -142,11 +160,16 @@ export const createAgent = ({ ] } const messages = [...params.messages, scheduleMessage] + const { tools, close } = await getAgentTools() const { response, reasoning, text, usage } = await generateText({ model, messages, system: generateSystemPrompt(), stopWhen: stepCountIs(Infinity), + onFinish: async () => { + await close() + }, + tools, }) return { messages: [scheduleMessage, ...response.messages], @@ -170,6 +193,7 @@ export const createAgent = ({ reasoning: [], usage: null } + const { tools, close } = await getAgentTools() const { fullStream } = streamText({ model, messages, @@ -181,7 +205,8 @@ export const createAgent = ({ } }, tools, - onFinish: ({ usage, reasoning, response }) => { + onFinish: async ({ usage, reasoning, response }) => { + await close() result.usage = usage as never result.reasoning = reasoning.map(part => part.text) result.messages = response.messages diff --git a/agent/src/models.ts b/agent/src/models.ts index e4df5b6d..f4aca341 100644 --- a/agent/src/models.ts +++ b/agent/src/models.ts @@ -54,4 +54,29 @@ export const FileAttachmentModel = z.object({ metadata: z.record(z.string(), z.any()).optional(), }) -export const AttachmentModel = z.union([ImageAttachmentModel, FileAttachmentModel]) \ No newline at end of file +export const AttachmentModel = z.union([ImageAttachmentModel, FileAttachmentModel]) + +export const HTTPMCPConnectionModel = z.object({ + name: z.string().min(1, 'Name is required'), + type: z.literal('http'), + url: z.string().min(1, 'URL is required'), + headers: z.record(z.string(), z.string()).optional(), +}) + +export const SSEMCPConnectionModel = z.object({ + name: z.string().min(1, 'Name is required'), + type: z.literal('sse'), + url: z.string().min(1, 'URL is required'), + headers: z.record(z.string(), z.string()).optional(), +}) + +export const StdioMCPConnectionModel = z.object({ + name: z.string().min(1, 'Name is required'), + type: z.literal('stdio'), + command: z.string().min(1, 'Command is required'), + args: z.array(z.string()), + env: z.record(z.string(), z.string()).optional(), + cwd: z.string().optional(), +}) + +export const MCPConnectionModel = z.union([HTTPMCPConnectionModel, SSEMCPConnectionModel, StdioMCPConnectionModel]) \ No newline at end of file diff --git a/agent/src/modules/chat.ts b/agent/src/modules/chat.ts index 5805468f..00b7a21a 100644 --- a/agent/src/modules/chat.ts +++ b/agent/src/modules/chat.ts @@ -4,7 +4,7 @@ import { createAgent } from '../agent' import { createAuthFetcher } from '../index' import { ModelConfig } from '../types' import { bearerMiddleware } from '../middlewares/bearer' -import { AllowedActionModel, AttachmentModel, IdentityContextModel, ModelConfigModel } from '../models' +import { AllowedActionModel, AttachmentModel, IdentityContextModel, MCPConnectionModel, ModelConfigModel } from '../models' import { allActions } from '../types' const AgentModel = z.object({ @@ -18,6 +18,7 @@ const AgentModel = z.object({ query: z.string(), identity: IdentityContextModel, attachments: z.array(AttachmentModel).optional().default([]), + mcpConnections: z.array(MCPConnectionModel).optional().default([]), }) export const chatModule = new Elysia({ prefix: '/chat' }) @@ -31,6 +32,7 @@ export const chatModule = new Elysia({ prefix: '/chat' }) currentChannel: body.currentChannel, allowedActions: body.allowedActions, identity: body.identity, + mcpConnections: body.mcpConnections, }, authFetcher) return ask({ query: body.query, @@ -44,22 +46,23 @@ export const chatModule = new Elysia({ prefix: '/chat' }) .post('/stream', async function* ({ body, bearer }) { try { const authFetcher = createAuthFetcher(bearer) - const { stream } = createAgent({ - model: body.model as ModelConfig, - activeContextTime: body.activeContextTime, - channels: body.channels, - currentChannel: body.currentChannel, - allowedActions: body.allowedActions, - identity: body.identity, - }, authFetcher) - for await (const action of stream({ - query: body.query, - messages: body.messages, - skills: body.skills, - attachments: body.attachments, - })) { - yield sse(JSON.stringify(action)) - } + const { stream } = createAgent({ + model: body.model as ModelConfig, + activeContextTime: body.activeContextTime, + channels: body.channels, + currentChannel: body.currentChannel, + allowedActions: body.allowedActions, + identity: body.identity, + mcpConnections: body.mcpConnections, + }, authFetcher) + for await (const action of stream({ + query: body.query, + messages: body.messages, + skills: body.skills, + attachments: body.attachments, + })) { + yield sse(JSON.stringify(action)) + } } catch (error) { console.error(error) yield sse(JSON.stringify({ @@ -70,4 +73,3 @@ export const chatModule = new Elysia({ prefix: '/chat' }) }, { body: AgentModel, }) - \ No newline at end of file diff --git a/agent/src/tools/mcp.ts b/agent/src/tools/mcp.ts new file mode 100644 index 00000000..583cad02 --- /dev/null +++ b/agent/src/tools/mcp.ts @@ -0,0 +1,51 @@ +import { HTTPMCPConnection, MCPConnection, SSEMCPConnection, StdioMCPConnection } from '../types' +import { createMCPClient } from '@ai-sdk/mcp' + +export const getMCPTools = async (connections: MCPConnection[]) => { + const closeCallbacks: Array<() => Promise> = [] + + const getHTTPTools = async (connection: HTTPMCPConnection) => { + const client = await createMCPClient({ + transport: { + type: 'http', + url: connection.url, + headers: connection.headers, + } + }) + closeCallbacks.push(client.close) + return await client.tools() + } + + const getSSETools = async (connection: SSEMCPConnection) => { + const client = await createMCPClient({ + transport: { + type: 'sse', + url: connection.url, + headers: connection.headers, + } + }) + closeCallbacks.push(client.close) + return await client.tools() + } + + const getStdioTools = async (connection: StdioMCPConnection) => { + // TODO: Implement stdio tools + return [] + } + + return { + tools: await Promise.all(connections.map(connection => { + switch (connection.type) { + case 'http': + return getHTTPTools(connection) + case 'sse': + return getSSETools(connection) + case 'stdio': + return getStdioTools(connection) + } + })), + close: async () => { + await Promise.all(closeCallbacks.map(callback => callback())) + } + } +} \ No newline at end of file diff --git a/agent/src/types/agent.ts b/agent/src/types/agent.ts index 08cac1bf..bd505fe0 100644 --- a/agent/src/types/agent.ts +++ b/agent/src/types/agent.ts @@ -1,6 +1,7 @@ import { ModelMessage } from 'ai' import { ModelConfig } from './model' import { AgentAttachment } from './attachment' +import { MCPConnection } from './mcp' export interface IdentityContext { botId: string @@ -43,6 +44,7 @@ export interface AgentParams { identity: IdentityContext channels?: string[] currentChannel?: string + mcpConnections?: MCPConnection[] } export interface AgentInput { diff --git a/agent/src/types/index.ts b/agent/src/types/index.ts index 9428f833..dd5abe62 100644 --- a/agent/src/types/index.ts +++ b/agent/src/types/index.ts @@ -1,4 +1,5 @@ export * from './agent' export * from './model' export * from './schedule' -export * from './attachment' \ No newline at end of file +export * from './attachment' +export * from './mcp' \ No newline at end of file diff --git a/agent/src/types/mcp.ts b/agent/src/types/mcp.ts new file mode 100644 index 00000000..ab5dc2f5 --- /dev/null +++ b/agent/src/types/mcp.ts @@ -0,0 +1,29 @@ +export interface BaseMCPConnection { + type: string + name: string +} + +export interface StdioMCPConnection extends BaseMCPConnection { + type: 'stdio' + command: string + args: string[] + env?: Record + cwd?: string +} + +export interface HTTPMCPConnection extends BaseMCPConnection { + type: 'http' + url: string + headers?: Record +} + +export interface SSEMCPConnection extends BaseMCPConnection { + type: 'sse' + url: string + headers?: Record +} + +export type MCPConnection = + | StdioMCPConnection + | HTTPMCPConnection + | SSEMCPConnection \ No newline at end of file