You've already forked agentic-coding-workflow
252 lines
7.0 KiB
Go
252 lines
7.0 KiB
Go
package mcp
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/go-chi/chi/v5"
|
|
"github.com/google/uuid"
|
|
mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
|
|
|
|
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
|
httpx "github.com/yan1h/agent-coding-workflow/internal/transport/http"
|
|
"github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware"
|
|
"github.com/yan1h/agent-coding-workflow/internal/user"
|
|
)
|
|
|
|
// 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))
|
|
})
|
|
}
|
|
}
|
|
|
|
// ===== Admin Token CRUD HTTP Handler =====
|
|
|
|
// AdminHandlerDeps 装配 admin token CRUD HTTP handler 所需依赖。
|
|
type AdminHandlerDeps struct {
|
|
Tokens TokenService
|
|
Users user.Service
|
|
}
|
|
|
|
// MountAdmin 把 admin token CRUD 挂到 /api/v1/mcp/tokens。
|
|
func MountAdmin(r chi.Router, deps AdminHandlerDeps) {
|
|
r.Route("/api/v1/mcp/tokens", func(r chi.Router) {
|
|
r.Post("/", postToken(deps))
|
|
r.Get("/", listTokens(deps))
|
|
r.Delete("/{id}", deleteToken(deps))
|
|
})
|
|
}
|
|
|
|
type postTokenReq struct {
|
|
UserID string `json:"user_id,omitempty"`
|
|
Name string `json:"name"`
|
|
Scope Scope `json:"scope"`
|
|
ExpiresAt time.Time `json:"expires_at"`
|
|
}
|
|
type postTokenResp struct {
|
|
ID string `json:"id"`
|
|
Token string `json:"token"`
|
|
ExpiresAt time.Time `json:"expires_at"`
|
|
}
|
|
|
|
func postToken(deps AdminHandlerDeps) http.HandlerFunc {
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
|
rid := middleware.RequestIDFromContext(r.Context())
|
|
caller, isAdmin, err := resolveSessionUser(r.Context(), deps.Users, r)
|
|
if err != nil {
|
|
httpx.WriteError(w, rid, err)
|
|
return
|
|
}
|
|
if !isAdmin {
|
|
httpx.WriteError(w, rid, errs.New(errs.CodeForbidden, "admin only"))
|
|
return
|
|
}
|
|
var req postTokenReq
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
httpx.WriteError(w, rid, errs.Wrap(err, errs.CodeInvalidInput, "decode body"))
|
|
return
|
|
}
|
|
var targetUID uuid.UUID
|
|
if req.UserID == "" {
|
|
targetUID = caller
|
|
} else {
|
|
tid, err := uuid.Parse(req.UserID)
|
|
if err != nil {
|
|
httpx.WriteError(w, rid, errs.Wrap(err, errs.CodeInvalidInput, "user_id parse"))
|
|
return
|
|
}
|
|
targetUID = tid
|
|
}
|
|
res, err := deps.Tokens.IssueAdmin(r.Context(), IssueAdminInput{
|
|
UserID: targetUID,
|
|
Name: req.Name,
|
|
Scope: req.Scope,
|
|
ExpiresAt: req.ExpiresAt,
|
|
CreatedBy: caller,
|
|
})
|
|
if err != nil {
|
|
httpx.WriteError(w, rid, err)
|
|
return
|
|
}
|
|
httpx.WriteJSON(w, http.StatusCreated, postTokenResp{
|
|
ID: res.ID.String(),
|
|
Token: res.Plaintext,
|
|
ExpiresAt: derefTime(res.ExpiresAt),
|
|
})
|
|
}
|
|
}
|
|
|
|
type tokenRow struct {
|
|
ID string `json:"id"`
|
|
Name string `json:"name"`
|
|
Issuer string `json:"issuer"`
|
|
Scope Scope `json:"scope"`
|
|
ExpiresAt *string `json:"expires_at,omitempty"`
|
|
RevokedAt *string `json:"revoked_at,omitempty"`
|
|
LastUsedAt *string `json:"last_used_at,omitempty"`
|
|
CreatedAt string `json:"created_at"`
|
|
}
|
|
|
|
func listTokens(deps AdminHandlerDeps) http.HandlerFunc {
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
|
rid := middleware.RequestIDFromContext(r.Context())
|
|
caller, isAdmin, err := resolveSessionUser(r.Context(), deps.Users, r)
|
|
if err != nil {
|
|
httpx.WriteError(w, rid, err)
|
|
return
|
|
}
|
|
var callerFilter *uuid.UUID
|
|
if !isAdmin {
|
|
callerFilter = &caller
|
|
}
|
|
activeOnly := r.URL.Query().Get("active_only") == "true"
|
|
tokens, err := deps.Tokens.List(r.Context(), callerFilter, activeOnly)
|
|
if err != nil {
|
|
httpx.WriteError(w, rid, err)
|
|
return
|
|
}
|
|
out := make([]tokenRow, 0, len(tokens))
|
|
for _, t := range tokens {
|
|
row := tokenRow{
|
|
ID: t.ID.String(), Name: t.Name,
|
|
Issuer: string(t.Issuer), Scope: t.Scope,
|
|
CreatedAt: t.CreatedAt.Format(time.RFC3339),
|
|
}
|
|
if t.ExpiresAt != nil {
|
|
s := t.ExpiresAt.Format(time.RFC3339)
|
|
row.ExpiresAt = &s
|
|
}
|
|
if t.RevokedAt != nil {
|
|
s := t.RevokedAt.Format(time.RFC3339)
|
|
row.RevokedAt = &s
|
|
}
|
|
if t.LastUsedAt != nil {
|
|
s := t.LastUsedAt.Format(time.RFC3339)
|
|
row.LastUsedAt = &s
|
|
}
|
|
out = append(out, row)
|
|
}
|
|
httpx.WriteJSON(w, http.StatusOK, out)
|
|
}
|
|
}
|
|
|
|
func deleteToken(deps AdminHandlerDeps) http.HandlerFunc {
|
|
return func(w http.ResponseWriter, r *http.Request) {
|
|
rid := middleware.RequestIDFromContext(r.Context())
|
|
caller, _, err := resolveSessionUser(r.Context(), deps.Users, r)
|
|
if err != nil {
|
|
httpx.WriteError(w, rid, err)
|
|
return
|
|
}
|
|
idStr := chi.URLParam(r, "id")
|
|
id, err := uuid.Parse(idStr)
|
|
if err != nil {
|
|
httpx.WriteError(w, rid, errs.Wrap(err, errs.CodeInvalidInput, "id parse"))
|
|
return
|
|
}
|
|
if err := deps.Tokens.Revoke(r.Context(), id, caller); err != nil {
|
|
httpx.WriteError(w, rid, err)
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusNoContent)
|
|
}
|
|
}
|
|
|
|
func resolveSessionUser(ctx context.Context, us user.Service, r *http.Request) (uuid.UUID, bool, error) {
|
|
tok := extractBearer(r)
|
|
uid, ok, err := us.ResolveSession(ctx, tok)
|
|
if err != nil {
|
|
return uuid.Nil, false, errs.Wrap(err, errs.CodeInternal, "resolve session")
|
|
}
|
|
if !ok {
|
|
return uuid.Nil, false, errs.New(errs.CodeUnauthorized, "session invalid")
|
|
}
|
|
u, err := us.Get(ctx, uid)
|
|
if err != nil {
|
|
return uuid.Nil, false, err
|
|
}
|
|
return uid, u.IsAdmin, nil
|
|
}
|
|
|
|
func extractBearer(r *http.Request) string {
|
|
h := r.Header.Get("Authorization")
|
|
if strings.HasPrefix(h, "Bearer ") {
|
|
return strings.TrimPrefix(h, "Bearer ")
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func derefTime(t *time.Time) time.Time {
|
|
if t == nil {
|
|
return time.Time{}
|
|
}
|
|
return *t
|
|
}
|