FranklinWH-Toolbox/service/agentmcp/http.go

package agentmcp

import (
	"context"
	"crypto/subtle"
	"errors"
	"net/http"
	"strings"
	"time"

	"github.com/mark3labs/mcp-go/server"
	"go.uber.org/zap"

	"fwh-toolbox-server/agentmcp/kit"
)

// Endpoint 一个工具集的 MCP 端点。
type Endpoint struct {
	Set    string            // 工具集，决定路径 EndpointPath(Set)
	Server *server.MCPServer // 注册了该工具集的 MCPServer
	Tasks  []kit.TaskTool    // 其中的持久化长任务工具
	Token  string            // 该端点的 Bearer 令牌（同时用于签名任务 id）
}

// Handler 返回 HTTP 处理器：每个工具集挂在 /mcp/<工具集>（Streamable HTTP，无状态，需 Bearer 令牌），只读工具集另挂 /mcp 兼容旧地址；另有 /healthz。
// 各端点的 MCPServer、并发名额、令牌相互独立，只共用监听端口。
// 任务协议请求由 kit.TaskRouter 在 mcp-go 之前处理（任务 id 以端点令牌签名）。
func Handler(endpoints []Endpoint, log *zap.Logger) http.Handler {
	mux := http.NewServeMux()
	for _, ep := range endpoints {
		path := EndpointPath(ep.Set)
		mcpHandler := server.NewStreamableHTTPServer(ep.Server,
			server.WithStateLess(true),
			server.WithEndpointPath(path),
		)
		h := bearerAuth(ep.Token, kit.NewTaskRouter(mcpHandler, ep.Tasks, ep.Token, log))
		mux.Handle(path, h)
		if ep.Set == SetReadOnly {
			mux.Handle("/mcp", h)
		}
	}
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	})
	return mux
}

// bearerAuth 校验 Authorization: Bearer <token>，常量时间比较。
func bearerAuth(token string, next http.Handler) http.Handler {
	want := []byte("Bearer " + token)
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got := []byte(strings.TrimSpace(r.Header.Get("Authorization")))
		if subtle.ConstantTimeCompare(got, want) != 1 {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusUnauthorized)
			_, _ = w.Write([]byte(`{"jsonrpc":"2.0","error":{"code":-32001,"message":"unauthorized"},"id":null}`))
			return
		}
		next.ServeHTTP(w, r)
	})
}

// Serve 启动 HTTP 服务，ctx 取消时优雅关闭。
func Serve(ctx context.Context, addr string, h http.Handler, log *zap.Logger) error {
	srv := &http.Server{
		Addr:              addr,
		Handler:           h,
		ReadHeaderTimeout: 10 * time.Second,
	}
	errCh := make(chan error, 1)
	go func() { errCh <- srv.ListenAndServe() }()
	log.Info("[agent-mcp] HTTP 服务已启动", zap.String("addr", addr))

	select {
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		return srv.Shutdown(shutdownCtx)
	case err := <-errCh:
		if errors.Is(err, http.ErrServerClosed) {
			return nil
		}
		return err
	}
}