add proxy info

Author: idoubi

.env.example.toml

@@ -4,10 +4,10 @@ port = 8025
 [api_server]
 port = 8027
 
-[mcp_server_commands]
-puppeteer = "npx -y @modelcontextprotocol/server-puppeteer"
-fetch = "uvx mcp-server-fetch"
-time = "docker run -i --rm mcp/time"
+[mcp_servers]
+puppeteer = { command="npx -y @modelcontextprotocol/server-puppeteer", share_process=true }
+fetch = { command="uvx mcp-server-fetch", share_process=true }
+time = { command="docker run -i --rm mcp/time", share_process=true }
 
 [remote_apis]
-get_server_command = "https://mcp.so/api/get-server-command"
+get_server_config = "http://127.0.0.1:3000/api/get-server-config"

handler/proxy/messages.go

@@ -1,8 +1,10 @@
 package proxy
 
 import (
+	"encoding/json"
 	"fmt"
 	"net/http"
+	"time"
 
 	"github.com/chatmcp/mcprouter/service/jsonrpc"
 	"github.com/chatmcp/mcprouter/service/mcpclient"
@@ -32,12 +34,37 @@ func Messages(c echo.Context) error {
 		return ctx.JSONRPCError(jsonrpc.ErrorParseError, nil)
 	}
 
+	proxyInfo := session.ProxyInfo()
 	sseKey := session.Key()
 
+	proxyInfo.JSONRPCVersion = request.JSONRPC
+	proxyInfo.RequestMethod = request.Method
+	proxyInfo.RequestTime = time.Now()
+	proxyInfo.RequestParams = request.Params
+
+	if request.ID != nil {
+		proxyInfo.RequestID = request.ID
+	}
+
+	if request.Method == "initialize" {
+		paramsB, _ := json.Marshal(request.Params)
+		params := &jsonrpc.InitializeParams{}
+		if err := json.Unmarshal(paramsB, params); err != nil {
+			return ctx.JSONRPCError(jsonrpc.ErrorParseError, nil)
+		}
+
+		proxyInfo.ClientName = params.ClientInfo.Name
+		proxyInfo.ClientVersion = params.ClientInfo.Version
+		proxyInfo.ProtocolVersion = params.ProtocolVersion
+
+		session.SetProxyInfo(proxyInfo)
+		ctx.StoreSession(sessionID, session)
+	}
+
 	client := ctx.GetClient(sseKey)
 
 	if client == nil {
-		command := session.Command()
+		command := proxyInfo.ServerCommand
 		_client, err := mcpclient.NewStdioClient(command)
 		if err != nil {
 			fmt.Printf("connect to mcp server failed: %v\n", err)
@@ -50,7 +77,6 @@ func Messages(c echo.Context) error {
 		}
 
 		ctx.StoreClient(sseKey, _client)
-		ctx.StoreSession(sessionID, session)
 
 		client = _client
 
@@ -67,13 +93,36 @@ func Messages(c echo.Context) error {
 	response, err := client.ForwardMessage(request)
 	if err != nil {
 		fmt.Printf("forward message failed: %v\n", err)
+		session.Close()
+		ctx.DeleteClient(sseKey)
 		return ctx.JSONRPCError(jsonrpc.ErrorProxyError, request.ID)
 	}
 
 	if response != nil {
+		if request.Method == "initialize" && response.Result != nil {
+			resultB, _ := json.Marshal(response.Result)
+			result := &jsonrpc.InitializeResult{}
+			if err := json.Unmarshal(resultB, result); err != nil {
+				fmt.Printf("unmarshal initialize result failed: %v\n", err)
+				return ctx.JSONRPCError(jsonrpc.ErrorParseError, request.ID)
+			}
+
+			proxyInfo.ServerName = result.ServerInfo.Name
+			proxyInfo.ServerVersion = result.ServerInfo.Version
+
+			session.SetProxyInfo(proxyInfo)
+			ctx.StoreSession(sessionID, session)
+		}
+
+		// not notification message, send sse message
 		session.SendMessage(response.String())
 	}
 
-	// notification message
+	proxyInfo.ResponseTime = time.Now()
+	proxyInfo.ResponseDuration = time.Since(proxyInfo.RequestTime)
+	proxyInfoB, _ := json.Marshal(proxyInfo)
+
+	fmt.Printf("proxyInfo: %s, cost: %f\n", string(proxyInfoB), proxyInfo.ResponseDuration.Seconds())
+
 	return ctx.JSONRPCResponse(response)
 }

handler/proxy/sse.go

@@ -7,7 +7,7 @@ import (
 
 	"github.com/chatmcp/mcprouter/service/mcpserver"
 	"github.com/chatmcp/mcprouter/service/proxy"
-	"github.com/google/uuid"
+	"github.com/chatmcp/mcprouter/util"
 	"github.com/labstack/echo/v4"
 )
 
@@ -25,19 +25,32 @@ func SSE(c echo.Context) error {
 		return c.String(http.StatusBadRequest, "Key is required")
 	}
 
-	command := mcpserver.GetCommand(key)
-	if command == "" {
-		return c.String(http.StatusBadRequest, "Server command not found")
+	serverConfig := mcpserver.GetServerConfig(key)
+	if serverConfig == nil {
+		return c.String(http.StatusBadRequest, "Invalid server config")
 	}
 
 	writer, err := proxy.NewSSEWriter(c)
 	if err != nil {
 		return c.String(http.StatusInternalServerError, err.Error())
 	}
 
+	// sessionID := uuid.New().String()
+	sessionID := util.MD5(key)
+
+	proxyInfo := &proxy.ProxyInfo{
+		ServerKey:          key,
+		SSERequestTime:     time.Now(),
+		SessionID:          sessionID,
+		ServerUUID:         serverConfig.ServerUUID,
+		ServerConfigName:   serverConfig.ServerName,
+		ServerShareProcess: serverConfig.ShareProcess,
+		ServerCommand:      serverConfig.Command,
+		ServerCommandHash:  serverConfig.CommandHash,
+	}
+
 	// store session
-	sessionID := uuid.New().String()
-	session := proxy.NewSSESession(writer, key, command)
+	session := proxy.NewSSESession(writer, proxyInfo)
 	ctx.StoreSession(sessionID, session)
 	defer ctx.DeleteSession(sessionID)
 
@@ -68,10 +81,10 @@ func SSE(c echo.Context) error {
 				return
 			case <-heartbeatTicker.C:
 				// Send heartbeat comment
-				if err := writer.SendHeartbeat(); err != nil {
-					session.Close()
-					return
-				}
+				// if err := writer.SendHeartbeat(); err != nil {
+				// 	session.Close()
+				// 	return
+				// }
 			case <-idleTimer.C:
 				// Close connection due to inactivity
 				session.Close()

service/api/context.go

@@ -29,7 +29,7 @@ type APIResponse struct {
 // APIContext is the context for the API request
 type APIContext struct {
 	echo.Context
-	command string
+	serverConfig *mcpserver.ServerConfig
 }
 
 func createAPIMiddleware() echo.MiddlewareFunc {
@@ -49,12 +49,12 @@ func createAPIMiddleware() echo.MiddlewareFunc {
 				return ctx.RespNoAuthMsg("no apikey")
 			}
 
-			command := mcpserver.GetCommand(apikey)
-			if command == "" {
+			serverConfig := mcpserver.GetServerConfig(apikey)
+			if serverConfig == nil {
 				return ctx.RespNoAuthMsg("invalid apikey")
 			}
 
-			ctx.command = command
+			ctx.serverConfig = serverConfig
 
 			return next(ctx)
 		}
@@ -86,12 +86,19 @@ func (c *APIContext) Valid(req interface{}) error {
 	return nil
 }
 
-func (c *APIContext) Command() string {
-	return c.command
+// ServerConfig returns the server config
+func (c *APIContext) ServerConfig() *mcpserver.ServerConfig {
+	return c.serverConfig
 }
 
+// ServerCommand returns the server command
+func (c *APIContext) ServerCommand() string {
+	return c.ServerConfig().Command
+}
+
+// Connect connects to the mcp server
 func (c *APIContext) Connect() (*mcpclient.StdioClient, error) {
-	command := c.Command()
+	command := c.ServerCommand()
 	if command == "" {
 		return nil, fmt.Errorf("invalid command")
 	}

service/jsonrpc/error.go

@@ -40,5 +40,5 @@ var (
 	ErrorInternalError = NewError(-32603, "Internal error", nil)
 
 	// ErrorProxyError is the error returned when the proxy error occurs.
-	ErrorProxyError = NewError(-32000, "Proxy error", nil)
+	ErrorProxyError = NewError(-32000, "Proxy error, Please restart client", nil)
 )

service/mcpclient/stdio.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"os/exec"
 	"sync"
+	"time"
 
 	"github.com/chatmcp/mcprouter/service/jsonrpc"
 	"github.com/tidwall/gjson"
@@ -106,7 +107,7 @@ func (c *StdioClient) listen() {
 
 		default:
 			message, err := c.stdout.ReadBytes('\n')
-			fmt.Printf("stdout read message: %s, %v\n", message, err)
+			// fmt.Printf("stdout read message: %s, %v\n", message, err)
 
 			if err != nil {
 				if err != io.EOF {
@@ -197,11 +198,17 @@ func (c *StdioClient) SendMessage(message []byte) ([]byte, error) {
 		return nil, fmt.Errorf("failed to write request message: %w", err)
 	}
 
-	fmt.Printf("stdin write request message: %s\n", message)
+	// fmt.Printf("stdin write request message: %s\n", message)
+
+	timeout := time.After(30 * time.Second)
 
 	// wait for response
 	for {
 		select {
+		case <-timeout:
+			fmt.Println("timeout waiting for response after 30 seconds")
+			c.Close()
+			return nil, fmt.Errorf("timeout waiting for response after 30 seconds")
 		case <-c.done:
 			fmt.Println("client closed with no response")
 			return nil, fmt.Errorf("client closed with no response")

service/mcpserver/command.go

@@ -2,6 +2,7 @@ package mcpserver
 
 import (
 	"bytes"
+	"crypto/md5"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -11,43 +12,73 @@ import (
 	"github.com/tidwall/gjson"
 )
 
-// GetCommand returns the command for the given key
-func GetCommand(key string) string {
-	command := viper.GetString(fmt.Sprintf("mcp_server_commands.%s", key))
-	if command == "" {
-		return getRemoteCommand(key)
+// GetServerCommand returns the command for the given key
+func GetServerCommand(key string) string {
+	config := GetServerConfig(key)
+	if config == nil {
+		return ""
+	}
+
+	return config.Command
+}
+
+// GetServerConfig returns the config for the given key
+func GetServerConfig(key string) *ServerConfig {
+	config := &ServerConfig{}
+	err := viper.UnmarshalKey(fmt.Sprintf("mcp_servers.%s", key), config)
+
+	if config.Command == "" {
+		fmt.Printf("get local config failed: %v, try to get remote config\n", err)
+
+		config, err = getRemoteServerConfig(key)
+		if err != nil {
+			fmt.Printf("get remote config failed: %v\n", err)
+			return nil
+		}
 	}
 
-	return command
+	return config
 }
 
-// getRemoteCommand returns the command for the given key from the remote API
-func getRemoteCommand(key string) string {
-	apiUrl := viper.GetString("remote_apis.get_server_command")
+// getRemoteServerConfig returns the config for the given key from the remote API
+func getRemoteServerConfig(key string) (*ServerConfig, error) {
+	apiUrl := viper.GetString("remote_apis.get_server_config")
 
-	fmt.Printf("get remote command from %s, with key: %s\n", apiUrl, key)
+	fmt.Printf("get remote config from %s, with key: %s\n", apiUrl, key)
 
 	params := map[string]string{
 		"server_key": key,
 	}
 
 	jsonData, err := json.Marshal(params)
 	if err != nil {
-		return ""
+		return nil, err
 	}
 
 	response, err := http.Post(apiUrl, "application/json", bytes.NewBuffer(jsonData))
 	if err != nil {
-		return ""
+		return nil, err
 	}
 
 	body, err := io.ReadAll(response.Body)
 	if err != nil {
-		return ""
+		return nil, err
 	}
 
 	data := gjson.ParseBytes(body)
-	command := data.Get("data.server_command").String()
 
-	return command
+	config := &ServerConfig{
+		ServerUUID:   data.Get("data.server_uuid").String(),
+		ServerName:   data.Get("data.server_name").String(),
+		ServerKey:    data.Get("data.server_key").String(),
+		Command:      data.Get("data.command").String(),
+		CommandHash:  data.Get("data.command_hash").String(),
+		ShareProcess: data.Get("data.share_process").Bool(),
+	}
+
+	if config.CommandHash == "" {
+		config.CommandHash = fmt.Sprintf("%x", md5.Sum([]byte(config.Command)))
+	}
+
+	return config, nil
 }

service/mcpserver/config.go

@@ -0,0 +1,11 @@
+package mcpserver
+
+// ServerConfig is the config for the remote mcp server
+type ServerConfig struct {
+	ServerUUID   string `json:"server_uuid,omitempty"`
+	ServerName   string `json:"server_name,omitempty"`
+	ServerKey    string `json:"server_key,omitempty"`
+	Command      string `json:"command"`
+	CommandHash  string `json:"command_hash,omitempty"`
+	ShareProcess bool   `json:"share_process,omitempty"`
+}