You've already forked agentic-coding-workflow
59 lines
1.7 KiB
Go
59 lines
1.7 KiB
Go
package mcp
|
|
|
|
import (
|
|
"net/http"
|
|
|
|
"github.com/go-chi/chi/v5"
|
|
mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
|
|
)
|
|
|
|
// MountTransports 在 chi router 上挂载 MCP transport 端点：
|
|
// - POST/GET /mcp Streamable HTTP（2025-03-26+，主推）
|
|
// - POST/GET /mcp/sse legacy HTTP+SSE（2024-11-05，SSE handler 自行根据 ?sessionid 分派）
|
|
//
|
|
// 中间件链：Auth → RateLimit → injectTokenID → SDK handler。
|
|
func MountTransports(
|
|
r chi.Router,
|
|
srv *mcpsdk.Server,
|
|
ts TokenService,
|
|
rl *RateLimiter,
|
|
) {
|
|
streamable := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server {
|
|
return srv
|
|
}, nil)
|
|
sse := mcpsdk.NewSSEHandler(func(*http.Request) *mcpsdk.Server {
|
|
return srv
|
|
}, nil)
|
|
|
|
// 嵌套 group 让中间件作用范围限定到 /mcp/* 子树
|
|
r.Group(func(r chi.Router) {
|
|
r.Use(NewAuthMiddleware(ts))
|
|
r.Use(NewRateLimitMiddleware(rl))
|
|
r.Use(injectTokenIDMiddleware())
|
|
|
|
// Streamable HTTP
|
|
r.Method(http.MethodPost, "/mcp", streamable)
|
|
r.Method(http.MethodGet, "/mcp", streamable)
|
|
|
|
// Legacy HTTP+SSE（SSE handler 自行分派 GET=stream / POST=message via ?sessionid）
|
|
r.Method(http.MethodGet, "/mcp/sse", sse)
|
|
r.Method(http.MethodPost, "/mcp/sse", sse)
|
|
})
|
|
}
|
|
|
|
// injectTokenIDMiddleware 在 auth middleware 之后跑，把 AuthSession.TokenID
|
|
// 复制到 ctx 的 mcpTokenIDCtxKey，供 AuditWrap 在 audit metadata 里加 via_mcp 标记。
|
|
func injectTokenIDMiddleware() func(http.Handler) http.Handler {
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
sess, ok := FromContext(r.Context())
|
|
if !ok {
|
|
next.ServeHTTP(w, r)
|
|
return
|
|
}
|
|
ctx := withMcpTokenID(r.Context(), sess.TokenID)
|
|
next.ServeHTTP(w, r.WithContext(ctx))
|
|
})
|
|
}
|
|
}
|