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"
)

// Handler 返回 HTTP 处理器：/mcp（Streamable HTTP，无状态，需 Bearer 令牌）与 /healthz。
func Handler(s *server.MCPServer, token string) http.Handler {
	mcpHandler := server.NewStreamableHTTPServer(s,
		server.WithStateLess(true),
		server.WithEndpointPath("/mcp"),
	)
	mux := http.NewServeMux()
	mux.Handle("/mcp", bearerAuth(token, mcpHandler))
	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
	}
}