diff --git a/internal/mcp/handler.go b/internal/mcp/handler.go new file mode 100644 index 0000000..9815421 --- /dev/null +++ b/internal/mcp/handler.go @@ -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)) + }) + } +}