hyperguild/internal/mcp/server.go

package mcp

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

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

// NewServer constructs an MCP HTTP handler. If token is non-empty, every
// request must carry "Authorization: Bearer <token>" or it is rejected with
// HTTP 401 and JSON-RPC error -32001. Empty token disables auth (default).
func NewServer(reg *registry.Registry, token string) *Server {
	return &Server{reg: reg, token: token}
}

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 when one is configured. Returns true if
// the request may proceed, false if it has been rejected (401 already written).
func (s *Server) checkAuth(w http.ResponseWriter, r *http.Request) bool {
	if s.token == "" {
		return true
	}

	const prefix = "Bearer "
	hdr := r.Header.Get("Authorization")
	if !strings.HasPrefix(hdr, prefix) ||
		subtle.ConstantTimeCompare([]byte(hdr[len(prefix):]), []byte(s.token)) != 1 {
		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"},
		})
		return false
	}
	return true
}

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},
	})
}