From fe67f79050c7f4b8971f9b9aabc22c5fd23bac07 Mon Sep 17 00:00:00 2001 From: Kadxy <2230318258@qq.com> Date: Sun, 29 Dec 2024 09:24:52 +0800 Subject: [PATCH] feat: MCP message type --- app/mcp/actions.ts | 9 +++++-- app/mcp/client.ts | 6 ++++- app/mcp/types.ts | 61 ++++++++++++++++++++++++++++++++++++++++++++++ app/store/chat.ts | 48 ++++++++++++++++++++++-------------- 4 files changed, 103 insertions(+), 21 deletions(-) create mode 100644 app/mcp/types.ts diff --git a/app/mcp/actions.ts b/app/mcp/actions.ts index af8683440..5fe611b3a 100644 --- a/app/mcp/actions.ts +++ b/app/mcp/actions.ts @@ -3,8 +3,9 @@ import { createClient, executeRequest } from "./client"; import { MCPClientLogger } from "./logger"; import conf from "./mcp_config.json"; +import { McpRequestMessage } from "./types"; -const logger = new MCPClientLogger("MCP Server"); +const logger = new MCPClientLogger("MCP Actions"); // Use Map to store all clients const clientsMap = new Map(); @@ -51,7 +52,10 @@ export async function initializeMcpClients() { } // Execute MCP request -export async function executeMcpAction(clientId: string, request: any) { +export async function executeMcpAction( + clientId: string, + request: McpRequestMessage, +) { try { // Find the corresponding client const client = clientsMap.get(clientId); @@ -61,6 +65,7 @@ export async function executeMcpAction(clientId: string, request: any) { } logger.info(`Executing MCP request for ${clientId}`); + // Execute request and return result return await executeRequest(client, request); } catch (error) { diff --git a/app/mcp/client.ts b/app/mcp/client.ts index 7eb55fb82..0600f00be 100644 --- a/app/mcp/client.ts +++ b/app/mcp/client.ts @@ -1,6 +1,7 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js"; import { MCPClientLogger } from "./logger"; +import { McpRequestMessage } from "./types"; import { z } from "zod"; export interface ServerConfig { @@ -79,6 +80,9 @@ export async function listPrimitives(client: Client) { } /** Execute a request */ -export async function executeRequest(client: Client, request: any) { +export async function executeRequest( + client: Client, + request: McpRequestMessage, +) { return client.request(request, z.any()); } diff --git a/app/mcp/types.ts b/app/mcp/types.ts new file mode 100644 index 000000000..763121bad --- /dev/null +++ b/app/mcp/types.ts @@ -0,0 +1,61 @@ +// ref: https://spec.modelcontextprotocol.io/specification/basic/messages/ + +import { z } from "zod"; + +export interface McpRequestMessage { + jsonrpc?: "2.0"; + id?: string | number; + method: "tools/call" | string; + params?: { + [key: string]: unknown; + }; +} + +export const McpRequestMessageSchema: z.ZodType = z.object({ + jsonrpc: z.literal("2.0").optional(), + id: z.union([z.string(), z.number()]).optional(), + method: z.string(), + params: z.record(z.unknown()).optional(), +}); + +export interface McpResponseMessage { + jsonrpc?: "2.0"; + id?: string | number; + result?: { + [key: string]: unknown; + }; + error?: { + code: number; + message: string; + data?: unknown; + }; +} + +export const McpResponseMessageSchema: z.ZodType = z.object( + { + jsonrpc: z.literal("2.0").optional(), + id: z.union([z.string(), z.number()]).optional(), + result: z.record(z.unknown()).optional(), + error: z + .object({ + code: z.number(), + message: z.string(), + data: z.unknown().optional(), + }) + .optional(), + }, +); + +export interface McpNotifications { + jsonrpc?: "2.0"; + method: string; + params?: { + [key: string]: unknown; + }; +} + +export const McpNotificationsSchema: z.ZodType = z.object({ + jsonrpc: z.literal("2.0").optional(), + method: z.string(), + params: z.record(z.unknown()).optional(), +}); diff --git a/app/store/chat.ts b/app/store/chat.ts index d30fa1fea..e0ee95621 100644 --- a/app/store/chat.ts +++ b/app/store/chat.ts @@ -1,4 +1,9 @@ -import { getMessageTextContent, trimTopic } from "../utils"; +import { + getMessageTextContent, + isDalle3, + safeLocalStorage, + trimTopic, +} from "../utils"; import { indexedDBStorage } from "@/app/utils/indexedDB-storage"; import { nanoid } from "nanoid"; @@ -14,14 +19,13 @@ import { DEFAULT_INPUT_TEMPLATE, DEFAULT_MODELS, DEFAULT_SYSTEM_TEMPLATE, + GEMINI_SUMMARIZE_MODEL, KnowledgeCutOffDate, + ServiceProvider, StoreKey, SUMMARIZE_MODEL, - GEMINI_SUMMARIZE_MODEL, - ServiceProvider, } from "../constant"; import Locale, { getLang } from "../locales"; -import { isDalle3, safeLocalStorage } from "../utils"; import { prettyObject } from "../utils/format"; import { createPersistStore } from "../utils/store"; import { estimateTokenLength } from "../utils/token"; @@ -55,6 +59,7 @@ export type ChatMessage = RequestMessage & { model?: ModelType; tools?: ChatMessageTool[]; audio_url?: string; + isMcpResponse?: boolean; }; export function createMessage(override: Partial): ChatMessage { @@ -368,20 +373,22 @@ export const useChatStore = createPersistStore( get().summarizeSession(false, targetSession); }, - async onUserInput(content: string, attachImages?: string[]) { + async onUserInput( + content: string, + attachImages?: string[], + isMcpResponse?: boolean, + ) { const session = get().currentSession(); const modelConfig = session.mask.modelConfig; - const userContent = fillTemplateWith(content, modelConfig); - console.log("[User Input] after template: ", userContent); + // MCP Response no need to fill template + let mContent: string | MultimodalContent[] = isMcpResponse + ? content + : fillTemplateWith(content, modelConfig); - let mContent: string | MultimodalContent[] = userContent; - - if (attachImages && attachImages.length > 0) { + if (!isMcpResponse && attachImages && attachImages.length > 0) { mContent = [ - ...(userContent - ? [{ type: "text" as const, text: userContent }] - : []), + ...(content ? [{ type: "text" as const, text: content }] : []), ...attachImages.map((url) => ({ type: "image_url" as const, image_url: { url }, @@ -392,6 +399,7 @@ export const useChatStore = createPersistStore( let userMessage: ChatMessage = createMessage({ role: "user", content: mContent, + isMcpResponse, }); const botMessage: ChatMessage = createMessage({ @@ -770,9 +778,10 @@ export const useChatStore = createPersistStore( lastInput, }); }, + + /** check if the message contains MCP JSON and execute the MCP action */ checkMcpJson(message: ChatMessage) { - const content = - typeof message.content === "string" ? message.content : ""; + const content = getMessageTextContent(message); if (isMcpJson(content)) { try { const mcpRequest = extractMcpJson(content); @@ -782,11 +791,14 @@ export const useChatStore = createPersistStore( executeMcpAction(mcpRequest.clientId, mcpRequest.mcp) .then((result) => { console.log("[MCP Response]", result); - // 直接使用onUserInput发送结果 - get().onUserInput( + const mcpResponse = typeof result === "object" ? JSON.stringify(result) - : String(result), + : String(result); + get().onUserInput( + `\`\`\`json:mcp:${mcpRequest.clientId}\n${mcpResponse}\n\`\`\``, + [], + true, ); }) .catch((error) => showToast(String(error)));