Skip to content

Commit 43519b3

Browse files
zijiren233zijiren233
committed
feat: mcp servers test api
1 parent d1583a5 commit 43519b3

14 files changed

Lines changed: 1319 additions & 27 deletions

File tree

‎core/controller/embedmcp.go‎

Lines changed: 203 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,25 @@
11
package controller
22

33
import (
4+
"context"
5+
"fmt"
46
"maps"
57
"net/http"
8+
"net/url"
69
"slices"
10+
"strings"
711

812
"github.com/gin-gonic/gin"
9-
"github.com/labring/aiproxy/core/embedmcp"
13+
"github.com/labring/aiproxy/core/common/mcpproxy"
14+
statelessmcp "github.com/labring/aiproxy/core/common/stateless-mcp"
15+
"github.com/labring/aiproxy/core/mcpservers"
1016
"github.com/labring/aiproxy/core/middleware"
1117
"github.com/labring/aiproxy/core/model"
18+
"github.com/mark3labs/mcp-go/mcp"
19+
"github.com/mark3labs/mcp-go/server"
1220

1321
// init embed mcp
14-
_ "github.com/labring/aiproxy/core/embedmcp/mcpregister"
22+
_ "github.com/labring/aiproxy/core/mcpservers/mcpregister"
1523
)
1624

1725
type EmbedMCPConfigTemplate struct {
@@ -21,18 +29,18 @@ type EmbedMCPConfigTemplate struct {
2129
Description string `json:"description,omitempty"`
2230
}
2331

24-
func newEmbedMCPConfigTemplate(template embedmcp.ConfigTemplate) EmbedMCPConfigTemplate {
32+
func newEmbedMCPConfigTemplate(template mcpservers.ConfigTemplate) EmbedMCPConfigTemplate {
2533
return EmbedMCPConfigTemplate{
2634
Name: template.Name,
27-
Required: template.Required == embedmcp.ConfigRequiredTypeInitOnly,
35+
Required: template.Required == mcpservers.ConfigRequiredTypeInitOnly,
2836
Example: template.Example,
2937
Description: template.Description,
3038
}
3139
}
3240

3341
type EmbedMCPConfigTemplates = map[string]EmbedMCPConfigTemplate
3442

35-
func newEmbedMCPConfigTemplates(templates embedmcp.ConfigTemplates) EmbedMCPConfigTemplates {
43+
func newEmbedMCPConfigTemplates(templates mcpservers.ConfigTemplates) EmbedMCPConfigTemplates {
3644
emcpTemplates := make(EmbedMCPConfigTemplates, len(templates))
3745
for key, template := range templates {
3846
emcpTemplates[key] = newEmbedMCPConfigTemplate(template)
@@ -49,7 +57,7 @@ type EmbedMCP struct {
4957
ConfigTemplates EmbedMCPConfigTemplates `json:"config_templates"`
5058
}
5159

52-
func newEmbedMCP(mcp *embedmcp.EmbedMcp, enabled bool) *EmbedMCP {
60+
func newEmbedMCP(mcp *mcpservers.EmbedMcp, enabled bool) *EmbedMCP {
5361
emcp := &EmbedMCP{
5462
ID: mcp.ID,
5563
Enabled: enabled,
@@ -72,7 +80,7 @@ func newEmbedMCP(mcp *embedmcp.EmbedMcp, enabled bool) *EmbedMCP {
7280
// @Success 200 {array} EmbedMCP
7381
// @Router /api/embedmcp/ [get]
7482
func GetEmbedMCPs(c *gin.Context) {
75-
embeds := embedmcp.Servers()
83+
embeds := mcpservers.Servers()
7684
enabledMCPs, err := model.GetPublicMCPsEnabled(slices.Collect(maps.Keys(embeds)))
7785
if err != nil {
7886
middleware.ErrorResponse(c, http.StatusInternalServerError, err.Error())
@@ -111,13 +119,13 @@ func SaveEmbedMCP(c *gin.Context) {
111119
return
112120
}
113121

114-
emcp, ok := embedmcp.GetEmbedMCP(req.ID)
122+
emcp, ok := mcpservers.GetEmbedMCP(req.ID)
115123
if !ok {
116124
middleware.ErrorResponse(c, http.StatusNotFound, "embed mcp not found")
117125
return
118126
}
119127

120-
pmcp, err := embedmcp.ToPublicMCP(emcp, req.InitConfig, req.Enabled)
128+
pmcp, err := mcpservers.ToPublicMCP(emcp, req.InitConfig, req.Enabled)
121129
if err != nil {
122130
middleware.ErrorResponse(c, http.StatusInternalServerError, err.Error())
123131
return
@@ -130,3 +138,189 @@ func SaveEmbedMCP(c *gin.Context) {
130138

131139
middleware.SuccessResponse(c, nil)
132140
}
141+
142+
type testEmbedMcpEndpointProvider struct {
143+
key string
144+
}
145+
146+
func newTestEmbedMcpEndpoint(key string) mcpproxy.EndpointProvider {
147+
return &testEmbedMcpEndpointProvider{
148+
key: key,
149+
}
150+
}
151+
152+
func (m *testEmbedMcpEndpointProvider) NewEndpoint(session string) (newEndpoint string) {
153+
endpoint := fmt.Sprintf("/api/test-embedmcp/message?sessionId=%s&key=%s", session, m.key)
154+
return endpoint
155+
}
156+
157+
func (m *testEmbedMcpEndpointProvider) LoadEndpoint(endpoint string) (session string) {
158+
parsedURL, err := url.Parse(endpoint)
159+
if err != nil {
160+
return ""
161+
}
162+
return parsedURL.Query().Get("sessionId")
163+
}
164+
165+
// query like: /api/test-embedmcp/aiproxy-openapi/sse?key=adminkey&config[key1]=value1&config[key2]=value2&reusing[key3]=value3
166+
func getConfigFromQuery(c *gin.Context) (map[string]string, map[string]string) {
167+
initConfig := make(map[string]string)
168+
reusingConfig := make(map[string]string)
169+
170+
queryParams := c.Request.URL.Query()
171+
172+
for paramName, paramValues := range queryParams {
173+
if len(paramValues) == 0 {
174+
continue
175+
}
176+
177+
paramValue := paramValues[0]
178+
179+
if strings.HasPrefix(paramName, "config[") && strings.HasSuffix(paramName, "]") {
180+
key := paramName[7 : len(paramName)-1]
181+
if key != "" {
182+
initConfig[key] = paramValue
183+
}
184+
}
185+
186+
if strings.HasPrefix(paramName, "reusing[") && strings.HasSuffix(paramName, "]") {
187+
key := paramName[8 : len(paramName)-1]
188+
if key != "" {
189+
reusingConfig[key] = paramValue
190+
}
191+
}
192+
}
193+
194+
return initConfig, reusingConfig
195+
}
196+
197+
// TestEmbedMCPSseServer godoc
198+
//
199+
// @Summary Test Embed MCP SSE Server
200+
// @Description Test Embed MCP SSE Server
201+
// @Tags embedmcp
202+
// @Security ApiKeyAuth
203+
// @Param id path string true "MCP ID"
204+
// @Param config[key] query string false "Initial configuration parameters (e.g., config[host]=http://localhost:3000)"
205+
// @Param reusing[key] query string false "Reusing configuration parameters (e.g., reusing[authorization]=apikey)"
206+
// @Success 200 {object} nil
207+
// @Failure 400 {object} nil
208+
// @Router /api/test-embedmcp/{id}/sse [get]
209+
func TestEmbedMCPSseServer(c *gin.Context) {
210+
id := c.Param("id")
211+
if id == "" {
212+
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
213+
mcp.NewRequestId(nil),
214+
mcp.INVALID_REQUEST,
215+
"mcp id is required",
216+
))
217+
return
218+
}
219+
220+
initConfig, reusingConfig := getConfigFromQuery(c)
221+
emcp, err := mcpservers.GetMCPServer(id, initConfig, reusingConfig)
222+
if err != nil {
223+
middleware.ErrorResponse(c, http.StatusBadRequest, err.Error())
224+
return
225+
}
226+
227+
handleTestEmbedMCPServer(c, emcp)
228+
}
229+
230+
const (
231+
testEmbedMcpType = "test-embedmcp"
232+
)
233+
234+
func handleTestEmbedMCPServer(c *gin.Context, s *server.MCPServer) {
235+
token := middleware.GetToken(c)
236+
237+
// Store the session
238+
store := getStore()
239+
newSession := store.New()
240+
241+
newEndpoint := newTestEmbedMcpEndpoint(token.Key).NewEndpoint(newSession)
242+
server := statelessmcp.NewSSEServer(
243+
s,
244+
statelessmcp.WithMessageEndpoint(newEndpoint),
245+
)
246+
247+
store.Set(newSession, testEmbedMcpType)
248+
defer func() {
249+
store.Delete(newSession)
250+
}()
251+
252+
ctx, cancel := context.WithCancel(c.Request.Context())
253+
defer cancel()
254+
255+
// Start message processing goroutine
256+
go processMCPSseMpscMessages(ctx, newSession, server)
257+
258+
// Handle SSE connection
259+
server.HandleSSE(c.Writer, c.Request)
260+
}
261+
262+
// TestEmbedMCPMessage godoc
263+
//
264+
// @Summary Test Embed MCP Message
265+
// @Description Send a message to the test embed MCP server
266+
// @Tags embedmcp
267+
// @Security ApiKeyAuth
268+
// @Param sessionId query string true "Session ID"
269+
// @Accept json
270+
// @Produce json
271+
// @Success 200 {object} nil
272+
// @Failure 400 {object} nil
273+
// @Router /api/test-embedmcp/message [post]
274+
func TestEmbedMCPMessage(c *gin.Context) {
275+
sessionID, _ := c.GetQuery("sessionId")
276+
if sessionID == "" {
277+
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
278+
mcp.NewRequestId(nil),
279+
mcp.INVALID_REQUEST,
280+
"missing sessionId",
281+
))
282+
return
283+
}
284+
285+
sendMCPSSEMessage(c, testEmbedMcpType, sessionID)
286+
}
287+
288+
// TestEmbedMCPStreamable godoc
289+
//
290+
// @Summary Test Embed MCP Streamable Server
291+
// @Description Test Embed MCP Streamable Server with various HTTP methods
292+
// @Tags embedmcp
293+
// @Security ApiKeyAuth
294+
// @Param id path string true "MCP ID"
295+
// @Param config[key] query string false "Initial configuration parameters (e.g., config[host]=http://localhost:3000)"
296+
// @Param reusing[key] query string false "Reusing configuration parameters (e.g., reusing[authorization]=apikey)"
297+
// @Accept json
298+
// @Produce json
299+
// @Success 200 {object} nil
300+
// @Failure 400 {object} nil
301+
// @Router /api/test-embedmcp/{id}/streamable [get]
302+
// @Router /api/test-embedmcp/{id}/streamable [post]
303+
// @Router /api/test-embedmcp/{id}/streamable [delete]
304+
func TestEmbedMCPStreamable(c *gin.Context) {
305+
id := c.Param("id")
306+
if id == "" {
307+
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
308+
mcp.NewRequestId(nil),
309+
mcp.INVALID_REQUEST,
310+
"mcp id is required",
311+
))
312+
return
313+
}
314+
315+
initConfig, reusingConfig := getConfigFromQuery(c)
316+
server, err := mcpservers.GetMCPServer(id, initConfig, reusingConfig)
317+
if err != nil {
318+
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
319+
mcp.NewRequestId(nil),
320+
mcp.INVALID_REQUEST,
321+
err.Error(),
322+
))
323+
return
324+
}
325+
handleGroupStreamableMCPServer(c, server)
326+
}

‎core/controller/groupmcp-server.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ func handleGroupProxySSE(c *gin.Context, config *model.GroupMCPProxyConfig) {
125125
c.Writer,
126126
c.Request,
127127
getStore(),
128-
newPublicMcpEndpoint(token.Key, model.PublicMCPTypeProxySSE),
128+
newGroupMcpEndpoint(token.Key, model.GroupMCPTypeProxySSE),
129129
backendURL.String(),
130130
headers,
131131
)
@@ -167,6 +167,7 @@ func handleGroupMCPServer(c *gin.Context, s *server.MCPServer, mcpType model.Gro
167167
// @Router /mcp/group/message [post]
168168
func GroupMCPMessage(c *gin.Context) {
169169
token := middleware.GetToken(c)
170+
170171
mcpTypeStr, _ := c.GetQuery("type")
171172
if mcpTypeStr == "" {
172173
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
@@ -177,6 +178,7 @@ func GroupMCPMessage(c *gin.Context) {
177178
return
178179
}
179180
mcpType := model.GroupMCPType(mcpTypeStr)
181+
180182
sessionID, _ := c.GetQuery("sessionId")
181183
if sessionID == "" {
182184
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(

‎core/controller/publicmcp-server.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +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"
20+
"github.com/labring/aiproxy/core/mcpservers"
2121
"github.com/labring/aiproxy/core/middleware"
2222
"github.com/labring/aiproxy/core/model"
2323
"github.com/labring/aiproxy/openapi-mcp/convert"
@@ -174,7 +174,7 @@ func handlePublicEmbedMCP(c *gin.Context, mcpID string, config *model.MCPEmbeddi
174174
}
175175
reusingConfig = param.ReusingParams
176176
}
177-
server, err := embedmcp.GetMCPServer(mcpID, config.Init, reusingConfig)
177+
server, err := mcpservers.GetMCPServer(mcpID, config.Init, reusingConfig)
178178
if err != nil {
179179
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
180180
mcp.NewRequestId(nil),
@@ -517,7 +517,7 @@ func handlePublicEmbedStreamable(c *gin.Context, mcpID string, config *model.MCP
517517
}
518518
reusingConfig = param.ReusingParams
519519
}
520-
server, err := embedmcp.GetMCPServer(mcpID, config.Init, reusingConfig)
520+
server, err := mcpservers.GetMCPServer(mcpID, config.Init, reusingConfig)
521521
if err != nil {
522522
c.JSON(http.StatusBadRequest, CreateMCPErrorResponse(
523523
mcp.NewRequestId(nil),

0 commit comments

Comments
 (0)