Skip to content

Commit 3fb2728

Browse files
zijiren233zijiren233
committed
feat: support embed mcp server
1 parent 8225567 commit 3fb2728

4 files changed

Lines changed: 220 additions & 1 deletion

File tree

‎core/controller/publicmcp-server.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
"github.com/labring/aiproxy/core/common"
1818
"github.com/labring/aiproxy/core/common/mcpproxy"
1919
statelessmcp "github.com/labring/aiproxy/core/common/stateless-mcp"
20+
"github.com/labring/aiproxy/core/embedmcp"
2021
"github.com/labring/aiproxy/core/middleware"
2122
"github.com/labring/aiproxy/core/model"
2223
"github.com/labring/aiproxy/openapi-mcp/convert"
@@ -145,6 +146,17 @@ func PublicMCPSseServer(c *gin.Context) {
145146
return
146147
}
147148
handleSSEMCPServer(c, server, model.PublicMCPTypeOpenAPI)
149+
case model.PublicMCPTypeEmbed:
150+
server, err := embedmcp.GetServer(publicMcp.ID, publicMcp.EmbeddingConfig.Init, nil)
151+
if err != nil {
152+
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
153+
mcp.NewRequestId(nil),
154+
mcp.INVALID_REQUEST,
155+
err.Error(),
156+
))
157+
return
158+
}
159+
handleSSEMCPServer(c, server, model.PublicMCPTypeEmbed)
148160
default:
149161
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
150162
mcp.NewRequestId(nil),
@@ -381,6 +393,8 @@ func PublicMCPMessage(c *gin.Context) {
381393
)
382394
case model.PublicMCPTypeOpenAPI:
383395
sendMCPSSEMessage(c, mcpTypeStr, sessionID)
396+
case model.PublicMCPTypeEmbed:
397+
sendMCPSSEMessage(c, mcpTypeStr, sessionID)
384398
default:
385399
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
386400
mcp.NewRequestId(nil),
@@ -456,6 +470,17 @@ func PublicMCPStreamable(c *gin.Context) {
456470
return
457471
}
458472
handleStreamableMCPServer(c, server)
473+
case model.PublicMCPTypeEmbed:
474+
server, err := embedmcp.GetServer(publicMcp.ID, publicMcp.EmbeddingConfig.Init, nil)
475+
if err != nil {
476+
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
477+
mcp.NewRequestId(nil),
478+
mcp.INVALID_REQUEST,
479+
err.Error(),
480+
))
481+
return
482+
}
483+
handleStreamableMCPServer(c, server)
459484
default:
460485
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
461486
mcp.NewRequestId(nil),

‎core/embedmcp/embedmcp.go‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
package embedmcp
2+
3+
import (
4+
"fmt"
5+
6+
"github.com/labring/aiproxy/core/model"
7+
"github.com/mark3labs/mcp-go/server"
8+
)
9+
10+
type ConfigValueValidator func(value string) error
11+
12+
type ConfigRequiredType int
13+
14+
const (
15+
ConfigRequiredTypeInitOptional ConfigRequiredType = iota
16+
ConfigRequiredTypeReusingOptional
17+
ConfigRequiredTypeInitOnly
18+
ConfigRequiredTypeReusingOnly
19+
ConfigRequiredTypeInitOrReusingOnly
20+
)
21+
22+
func (c ConfigRequiredType) Validate(config any, reusingConfig any) error {
23+
switch c {
24+
case ConfigRequiredTypeInitOnly:
25+
if config == nil {
26+
return fmt.Errorf("config is required")
27+
}
28+
case ConfigRequiredTypeReusingOnly:
29+
if reusingConfig == nil {
30+
return fmt.Errorf("reusing config is required")
31+
}
32+
case ConfigRequiredTypeInitOrReusingOnly:
33+
if config == nil && reusingConfig == nil {
34+
return fmt.Errorf("config or reusing config is required")
35+
}
36+
if config != nil && reusingConfig != nil {
37+
return fmt.Errorf("config and reusing config are both provided, but only one is allowed")
38+
}
39+
}
40+
return nil
41+
}
42+
43+
type ConfigTemplate struct {
44+
Required ConfigRequiredType `json:"required"`
45+
Example string `json:"example,omitempty"`
46+
Help string `json:"help,omitempty"`
47+
Validator ConfigValueValidator `json:"-"`
48+
}
49+
50+
type ConfigTemplates = map[string]ConfigTemplate
51+
52+
func ValidateConfigTemplates(ct ConfigTemplates, config map[string]string, reusingConfig map[string]string) error {
53+
if len(ct) == 0 {
54+
return nil
55+
}
56+
57+
for key, template := range ct {
58+
c := config[key]
59+
rc := reusingConfig[key]
60+
if err := template.Required.Validate(c, rc); err != nil {
61+
return fmt.Errorf("config required %s is invalid: %w", key, err)
62+
}
63+
if template.Validator != nil {
64+
if c != "" {
65+
if err := template.Validator(c); err != nil {
66+
return fmt.Errorf("config %s is invalid: %w", key, err)
67+
}
68+
} else if rc != "" {
69+
if err := template.Validator(rc); err != nil {
70+
return fmt.Errorf("reusing config %s is invalid: %w", key, err)
71+
}
72+
}
73+
}
74+
}
75+
76+
return nil
77+
}
78+
79+
func CheckConfigTemplatesExample(ct ConfigTemplates) error {
80+
for key, value := range ct {
81+
if value.Example == "" || value.Validator == nil {
82+
continue
83+
}
84+
if err := value.Validator(value.Example); err != nil {
85+
return fmt.Errorf("config %s example is invalid: %w", key, err)
86+
}
87+
}
88+
return nil
89+
}
90+
91+
type NewServerFunc func(config map[string]string, reusingConfig map[string]string) (*server.MCPServer, error)
92+
93+
type EmbeddingMcp struct {
94+
ID string
95+
Name string
96+
Readme string
97+
Tags []string
98+
ConfigTemplates ConfigTemplates
99+
NewServer NewServerFunc
100+
}
101+
102+
func (e *EmbeddingMcp) ToPublicMCP() *model.PublicMCP {
103+
return &model.PublicMCP{
104+
ID: e.ID,
105+
Name: e.Name,
106+
Readme: e.Readme,
107+
Tags: e.Tags,
108+
}
109+
}

‎core/embedmcp/register.go‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package embedmcp
2+
3+
import (
4+
"fmt"
5+
"sync"
6+
7+
"github.com/mark3labs/mcp-go/server"
8+
)
9+
10+
var (
11+
servers = make(map[string]EmbeddingMcp)
12+
mcpServerCache = make(map[string]*server.MCPServer)
13+
mcpServerCacheLock = sync.RWMutex{}
14+
)
15+
16+
func Register(mcp EmbeddingMcp) {
17+
if mcp.ID == "" {
18+
panic("mcp id is required")
19+
}
20+
if mcp.Name == "" {
21+
panic("mcp name is required")
22+
}
23+
if mcp.NewServer == nil {
24+
panic(fmt.Sprintf("mcp %s new server is required", mcp.ID))
25+
}
26+
if mcp.ConfigTemplates != nil {
27+
if err := CheckConfigTemplatesExample(mcp.ConfigTemplates); err != nil {
28+
panic(fmt.Sprintf("mcp %s config templates example is invalid: %v", mcp.ID, err))
29+
}
30+
}
31+
if _, ok := servers[mcp.ID]; ok {
32+
panic(fmt.Sprintf("mcp %s already registered", mcp.ID))
33+
}
34+
servers[mcp.ID] = mcp
35+
}
36+
37+
func GetServer(id string, config map[string]string, reusingConfig map[string]string) (*server.MCPServer, error) {
38+
embedServer, ok := servers[id]
39+
if !ok {
40+
return nil, fmt.Errorf("mcp %s not found", id)
41+
}
42+
if len(embedServer.ConfigTemplates) == 0 {
43+
return getNoConfigServer(embedServer)
44+
}
45+
if err := ValidateConfigTemplates(embedServer.ConfigTemplates, config, reusingConfig); err != nil {
46+
return nil, fmt.Errorf("mcp %s config is invalid: %w", id, err)
47+
}
48+
return embedServer.NewServer(config, reusingConfig)
49+
}
50+
51+
func getNoConfigServer(embedServer EmbeddingMcp) (*server.MCPServer, error) {
52+
mcpServerCacheLock.RLock()
53+
server, ok := mcpServerCache[embedServer.ID]
54+
mcpServerCacheLock.RUnlock()
55+
if ok {
56+
return server, nil
57+
}
58+
59+
mcpServerCacheLock.Lock()
60+
defer mcpServerCacheLock.Unlock()
61+
server, ok = mcpServerCache[embedServer.ID]
62+
if ok {
63+
return server, nil
64+
}
65+
66+
server, err := embedServer.NewServer(nil, nil)
67+
if err != nil {
68+
return nil, fmt.Errorf("mcp %s new server is invalid: %w", embedServer.ID, err)
69+
}
70+
mcpServerCache[embedServer.ID] = server
71+
return server, nil
72+
}
73+
74+
func Range(f func(id string, mcp EmbeddingMcp) bool) {
75+
for id, mcp := range servers {
76+
if !f(id, mcp) {
77+
break
78+
}
79+
}
80+
}

‎core/model/publicmcp.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ const (
2222
PublicMCPTypeProxyStreamable PublicMCPType = "mcp_proxy_streamable"
2323
PublicMCPTypeGitRepo PublicMCPType = "mcp_git_repo" // read only
2424
PublicMCPTypeOpenAPI PublicMCPType = "mcp_openapi"
25+
PublicMCPTypeEmbed PublicMCPType = "mcp_embed"
2526
)
2627

2728
type ParamType string
@@ -91,6 +92,10 @@ type MCPOpenAPIConfig struct {
9192
Authorization string `json:"authorization,omitempty"`
9293
}
9394

95+
type MCPEmbeddingConfig struct {
96+
Init map[string]string `json:"init"`
97+
}
98+
9499
type PublicMCP struct {
95100
ID string `gorm:"primaryKey" json:"id"`
96101
CreatedAt time.Time `gorm:"index" json:"created_at"`
@@ -102,11 +107,11 @@ type PublicMCP struct {
102107
ReadmeURL string `json:"readme_url"`
103108
Readme string `gorm:"type:text" json:"readme"`
104109
Tags []string `gorm:"serializer:fastjson;type:text" json:"tags,omitempty"`
105-
Author string `json:"author"`
106110
LogoURL string `json:"logo_url"`
107111
Price MCPPrice `gorm:"embedded" json:"price"`
108112
ProxyConfig *PublicMCPProxyConfig `gorm:"serializer:fastjson;type:text" json:"proxy_config,omitempty"`
109113
OpenAPIConfig *MCPOpenAPIConfig `gorm:"serializer:fastjson;type:text" json:"openapi_config,omitempty"`
114+
EmbeddingConfig *MCPEmbeddingConfig `gorm:"serializer:fastjson;type:text" json:"embedding_config,omitempty"`
110115
}
111116

112117
func (p *PublicMCP) BeforeCreate(_ *gorm.DB) error {

0 commit comments

Comments
 (0)