feat: mcp supports

.gitignore

@@ -38,4 +38,5 @@ coverage.txt
 *.o
 build/
 .env
-example/10-ai/*.yaml
+zipper.yml
+*.md

ai/register.go

@@ -0,0 +1,61 @@
+package ai
+
+import (
+	"sync"
+
+	"github.com/sashabaranov/go-openai"
+	"github.com/yomorun/yomo/core/metadata"
+)
+
+var (
+	mu              sync.Mutex
+	defaultRegister Register
+)
+
+// SetRegister sets the default register
+func SetRegister(r Register) {
+	mu.Lock()
+	defer mu.Unlock()
+	defaultRegister = r
+}
+
+// GetRegister gets the default register
+func GetRegister() Register {
+	mu.Lock()
+	defer mu.Unlock()
+	return defaultRegister
+}
+
+// ListToolCalls returns the list of tool calls
+func ListToolCalls(md metadata.M) ([]openai.Tool, error) {
+	if defaultRegister == nil {
+		return nil, nil
+	}
+	return defaultRegister.ListToolCalls(md)
+}
+
+// RegisterFunction registers a function calling function
+func RegisterFunction(functionDefinition *openai.FunctionDefinition, connID uint64, md metadata.M) error {
+	if defaultRegister == nil {
+		return nil
+	}
+	return defaultRegister.RegisterFunction(functionDefinition, connID, md)
+}
+
+// UnregisterFunction unregisters a function calling function
+func UnregisterFunction(connID uint64, md metadata.M) {
+	if defaultRegister == nil {
+		return
+	}
+	defaultRegister.UnregisterFunction(connID, md)
+}
+
+// Register provides an stateful register for registering and unregistering functions
+type Register interface {
+	// ListToolCalls returns the list of tool calls
+	ListToolCalls(md metadata.M) ([]openai.Tool, error)
+	// RegisterFunction registers a function calling function
+	RegisterFunction(fd *openai.FunctionDefinition, connID uint64, md metadata.M) error
+	// UnregisterFunction unregisters a function calling function
+	UnregisterFunction(connID uint64, md metadata.M)
+}

cli/serve.go

@@ -44,6 +44,8 @@
 	"github.com/yomorun/yomo/pkg/bridge/ai/provider/vertexai"
 	"github.com/yomorun/yomo/pkg/bridge/ai/provider/vllm"
 	"github.com/yomorun/yomo/pkg/bridge/ai/provider/xai"
+	"github.com/yomorun/yomo/pkg/bridge/llm"
+	"github.com/yomorun/yomo/pkg/bridge/mcp"
 )
 
 // serveCmd represents the serve command
@@ -99,6 +101,16 @@
 			// add AI connection middleware
 			options = append(options, yomo.WithFrameListener(listener))
 		}
+		// check and parse the mcp server config
+		mcpConfig, err := mcp.ParseConfig(bridgeConf)
+		if err != nil {
+			if err == mcp.ErrConfigNotFound {
+				ylog.Warn("mcp server is disabled")
+			} else {
+				log.FailureStatusEvent(os.Stdout, "%s", err.Error())
+				return
+			}
+		}
 		// new zipper
 		zipper, err := yomo.NewZipper(
 			conf.Name,
@@ -110,8 +122,8 @@
 		}
 		zipper.Logger().Info("using config file", "file_path", config)
 
-		// AI Server
 		if aiConfig != nil {
+			// AI Server
 			// register the llm provider
 			registerAIProvider(aiConfig)
 			// start the llm api server
@@ -122,12 +134,23 @@
 				conn2, _ := listener.Dial()
 				reducer := ai.NewReducer(conn2, auth.NewCredential(fmt.Sprintf("token:%s", tokenString)))
 
-				err := ai.Serve(aiConfig, ylog.Default(), source, reducer)
+				err := llm.Serve(aiConfig, ylog.Default(), source, reducer)
 				if err != nil {
 					log.FailureStatusEvent(os.Stdout, "%s", err.Error())
 					return
 				}
 			}()
+			// MCP Server
+			if mcpConfig != nil {
+				defer mcp.Stop()
+				go func() {
+					err = mcp.Start(mcpConfig, aiConfig, listenAddr, ylog.Default())
+					if err != nil {
+						log.FailureStatusEvent(os.Stdout, "%s", err.Error())
+						return
+					}
+				}()
+			}
 		}
 
 		// start the zipper

core/server.go

@@ -19,7 +19,7 @@
 
 	// authentication implements, Currently, only token authentication is implemented
 	_ "github.com/yomorun/yomo/pkg/auth"
-	"github.com/yomorun/yomo/pkg/bridge/ai/register"
+
 	"github.com/yomorun/yomo/pkg/frame-codec/y3codec"
 	yquic "github.com/yomorun/yomo/pkg/listener/quic"
 	pkgtls "github.com/yomorun/yomo/pkg/tls"
@@ -187,7 +187,7 @@
 
 	if conn.ClientType() == ClientTypeStreamFunction {
 		s.router.Remove(conn.ID())
-		register.UnregisterFunction(conn.ID(), conn.Metadata())
+		ai.UnregisterFunction(conn.ID(), conn.Metadata())
 	}
 	_ = s.connector.Remove(conn.ID())
 }
@@ -284,7 +284,7 @@
 	if err := json.Unmarshal([]byte(definition), &fd); err != nil {
 		return fmt.Errorf("unmarshal function definition error: %s", err.Error())
 	}
-	if err := register.RegisterFunction(&fd, conn.ID(), md); err != nil {
+	if err := ai.RegisterFunction(&fd, conn.ID(), md); err != nil {
 		return err
 	}
 	s.logger.Info("register ai function success", "function_name", fd.Name, "definition", string(definition))

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/invopop/jsonschema v0.13.0
 	github.com/joho/godotenv v1.5.1
 	github.com/lmittmann/tint v1.0.7
+	github.com/mark3labs/mcp-go v0.20.1
 	github.com/matoous/go-nanoid/v2 v2.1.0
 	github.com/quic-go/quic-go v0.50.1
 	github.com/robfig/cron/v3 v3.0.1
@@ -25,8 +26,6 @@ require (
 	github.com/spf13/viper v1.20.0
 	github.com/stretchr/testify v1.10.0
 	github.com/tetratelabs/wazero v1.9.0
-	github.com/tidwall/gjson v1.18.0
-	github.com/tidwall/sjson v1.2.5
 	github.com/vmihailenco/msgpack/v5 v5.4.1
 	github.com/yomorun/y3 v1.0.5
 	go.opentelemetry.io/otel v1.35.0
@@ -83,12 +82,15 @@ require (
 	github.com/spf13/afero v1.14.0 // indirect
 	github.com/spf13/cast v1.7.1 // indirect
 	github.com/subosito/gotenv v1.6.0 // indirect
+	github.com/tidwall/gjson v1.18.0 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.1 // indirect
+	github.com/tidwall/sjson v1.2.5 // indirect
 	github.com/tklauser/go-sysconf v0.3.15 // indirect
 	github.com/tklauser/numcpus v0.10.0 // indirect
 	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
 	github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
+	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
 	github.com/yusufpapurcu/wmi v1.2.4 // indirect
 	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0 // indirect