基于golang 实现dify的MCP SSE服务器 - AI创想

package main
import (
"context"
"encoding/json"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
"io"
"log"
"net/http"
"os"
)
func main() {
// Create a new MCP server
s := server.NewMCPServer(
"Calculator Demo",
"1.0.0",
server.WithResourceCapabilities(true, true),
server.WithLogging(),
server.WithRecovery(),
)
// 实现tools 工具
calculatorTool := mcp.NewTool("calculate",
mcp.WithDescription("Perform basic arithmetic operations"),
mcp.WithString("operation",
mcp.Required(),
mcp.Description("The operation to perform (add, subtract, multiply, divide)"),
mcp.Enum("add", "subtract", "multiply", "divide"),
),
mcp.WithNumber("x",
mcp.Required(),
mcp.Description("First number"),
),
mcp.WithNumber("y",
mcp.Required(),
mcp.Description("Second number"),
),
)
//资源
s.AddResource(mcp.NewResource("test://static/resource",
"Static Resource",
mcp.WithMIMEType("text/plain"),
), handleReadResource)
//动态模板资源
s.AddResourceTemplate(
mcp.NewResourceTemplate(
"test://dynamic/resource/{id}",
"Dynamic Resource",
),
handleResourceTemplate,
)
// Add the calculator handler
s.AddTool(calculatorTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
op := request.Params.Arguments["operation"].(string)
x := request.Params.Arguments["x"].(float64)
y := request.Params.Arguments["y"].(float64)
var result float64
switch op {
case "add":
result = x + y
case "subtract":
result = x - y
case "multiply":
result = x * y
case "divide":
if y == 0 {
return mcp.NewToolResultError("cannot divide by zero"), nil
}
result = x / y
}
return mcp.NewToolResultText(fmt.Sprintf("%.2f", result)), nil
})
sseServer := server.NewSSEServer(s,
server.WithSSEContextFunc(authFromRequest),
)
if err := sseServer.Start(":8080"); err != nil {
log.Fatalf("Server error: %v", err)
}
}
func handleResourceTemplate(
ctx context.Context,
request mcp.ReadResourceRequest,
) ([]mcp.ResourceContents, error) {
return []mcp.ResourceContents{
mcp.TextResourceContents{
URI: request.Params.URI,
MIMEType: "text/plain",
Text: fmt.Sprintf("动态模板:%+v ;", request),
},
}, nil
}
func handleReadResource(
ctx context.Context,
request mcp.ReadResourceRequest,
) ([]mcp.ResourceContents, error) {
return []mcp.ResourceContents{
mcp.TextResourceContents{
URI: "test://static/resource",
MIMEType: "text/plain",
Text: fmt.Sprintf("模板 %+v ；", request),
},
}, nil
}
// authKey is a custom context key for storing the auth token.
type authKey struct{}
// withAuthKey adds an auth key to the context.
func withAuthKey(ctx context.Context, auth string) context.Context {
return context.WithValue(ctx, authKey{}, auth)
}
// authFromRequest extracts the auth token from the request headers.
func authFromRequest(ctx context.Context, r *http.Request) context.Context {
return withAuthKey(ctx, r.Header.Get("Authorization"))
}
// authFromEnv extracts the auth token from the environment
func authFromEnv(ctx context.Context) context.Context {
return withAuthKey(ctx, os.Getenv("API_KEY"))
}
// tokenFromContext extracts the auth token from the context.
// This can be used by tools to extract the token regardless of the
// transport being used by the server.
func tokenFromContext(ctx context.Context) (string, error) {
auth, ok := ctx.Value(authKey{}).(string)
if !ok {
return "", fmt.Errorf("missing auth")
}
return auth, nil
}
type response struct {
Args map[string]interface{} `json:"args"`
Headers map[string]string `json:"headers"`
}
// makeRequest makes a request to httpbin.org including the auth token in the request
// headers and the message in the query string.
func makeRequest(ctx context.Context, message, token string) (*response, error) {
req, err := http.NewRequestWithContext(ctx, "GET", "https://httpbin.org/anything", nil)
if err != nil {
return nil, err
}
req.Header.Set("Authorization", token)
query := req.URL.Query()
query.Add("message", message)
req.URL.RawQuery = query.Encode()
resp, err := http.DefaultClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, err
}
var r *response
if err := json.Unmarshal(body, &r); err != nil {
return nil, err
}
return r, nil
}
// handleMakeAuthenticatedRequestTool is a tool that makes an authenticated request
// using the token from the context.
func handleMakeAuthenticatedRequestTool(
ctx context.Context,
request mcp.CallToolRequest,
) (*mcp.CallToolResult, error) {
message, ok := request.Params.Arguments["message"].(string)
if !ok {
return nil, fmt.Errorf("missing message")
}
token, err := tokenFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("missing token: %v", err)
}
// Now our tool can make a request with the token, irrespective of where it came from.
resp, err := makeRequest(ctx, message, token)
if err != nil {
return nil, err
}
return mcp.NewToolResultText(fmt.Sprintf("%+v", resp)), nil
}

复制代码