hyperguild/internal/mcp/server.go

package mcp

import (
	"context"
	"crypto/subtle"
	"encoding/json"
	"log/slog"
	"net/http"
	"strings"

	"github.com/mathiasbq/supervisor/internal/auth"
	"github.com/mathiasbq/supervisor/internal/registry"
)

type request struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      any             `json:"id"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params"`
}

type response struct {
	JSONRPC string    `json:"jsonrpc"`
	ID      any       `json:"id,omitempty"`
	Result  any       `json:"result,omitempty"`
	Error   *rpcError `json:"error,omitempty"`
}

type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// Server is an HTTP handler implementing the MCP JSON-RPC protocol.
type Server struct {
	reg       *registry.Registry
	token     string
	validator *auth.Validator
}

// NewServer constructs an MCP HTTP handler. token is the static bearer token
// (empty disables static auth). validator is optional; when non-nil, a valid
// JWT from Dex is accepted in addition to the static token.
func NewServer(reg *registry.Registry, token string, validator *auth.Validator) *Server {
	return &Server{reg: reg, token: token, validator: validator}
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if !s.checkAuth(w, r) {
		return
	}

	// GET opens the SSE stream for server-to-client events (MCP streamable HTTP).
	// claude.ai probes with GET before sending initialize, so accept without a session.
	if r.Method == http.MethodGet {
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Connection", "keep-alive")
		w.Header().Set("X-Accel-Buffering", "no")
		w.WriteHeader(http.StatusOK)
		if f, ok := w.(http.Flusher); ok {
			_, _ = w.Write([]byte(": stream open\n\n"))
			f.Flush()
		}
		<-r.Context().Done()
		return
	}

	var req request
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, nil, -32700, "parse error")
		return
	}

	// JSON-RPC 2.0 notifications (no id) must not receive a response.
	if req.ID == nil {
		return
	}

	var result any
	var rpcErr *rpcError

	switch req.Method {
	case "initialize":
		result = map[string]any{
			"protocolVersion": "2024-11-05",
			"capabilities":    map[string]any{"tools": map[string]any{}},
			"serverInfo":      map[string]any{"name": "supervisor", "version": "0.1.0"},
		}

	case "tools/list":
		result = map[string]any{"tools": s.reg.Tools()}

	case "tools/call":
		var p struct {
			Name      string          `json:"name"`
			Arguments json.RawMessage `json:"arguments"`
		}
		if err := json.Unmarshal(req.Params, &p); err != nil {
			rpcErr = &rpcError{Code: -32602, Message: "invalid params"}
			break
		}
		out, err := s.reg.Dispatch(context.Background(), p.Name, p.Arguments)
		if err != nil {
			rpcErr = &rpcError{Code: -32000, Message: err.Error()}
			break
		}
		result = map[string]any{
			"content": []map[string]any{{"type": "text", "text": string(out)}},
		}

	default:
		rpcErr = &rpcError{Code: -32601, Message: "method not found: " + req.Method}
	}

	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(response{
		JSONRPC: "2.0",
		ID:      req.ID,
		Result:  result,
		Error:   rpcErr,
	})
}

// checkAuth verifies the bearer token. Accepts a valid Dex JWT (when validator
// is configured) or the static token. Returns true if the request may proceed.
// When neither token nor validator is configured, auth is disabled (default).
func (s *Server) checkAuth(w http.ResponseWriter, r *http.Request) bool {
	if s.token == "" && s.validator == nil {
		return true
	}

	rawToken, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ")
	if !ok {
		s.rejectAuth(w, r)
		return false
	}

	if s.validator != nil {
		if _, err := s.validator.Validate(r.Context(), rawToken); err == nil {
			return true
		}
	}

	if s.token != "" && subtle.ConstantTimeCompare([]byte(rawToken), []byte(s.token)) == 1 {
		return true
	}

	s.rejectAuth(w, r)
	return false
}

func (s *Server) rejectAuth(w http.ResponseWriter, r *http.Request) {
	slog.Warn("mcp auth rejected", "remote", r.RemoteAddr, "method", r.Method)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusUnauthorized)
	_ = json.NewEncoder(w).Encode(response{
		JSONRPC: "2.0",
		Error:   &rpcError{Code: -32001, Message: "unauthorized"},
	})
}

func writeError(w http.ResponseWriter, id any, code int, msg string) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(response{
		JSONRPC: "2.0",
		ID:      id,
		Error:   &rpcError{Code: code, Message: msg},
	})
}