feat(mcp): mount Streamable HTTP + legacy SSE transports with auth+ratelimit

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-05-08 10:37:48 +08:00
parent 5fd4eab7ca
commit 6ce21d178e
1 changed files with 58 additions and 0 deletions
@@ -0,0 +1,58 @@
+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))
+		})
+	}
+}