gitea-mcp/internal/mcp/server.go

package mcp

import (
	"encoding/json"
	"errors"
	"net/http"

	"gitea.d-ma.be/mathias/gitea-mcp/internal/registry"
)

const ProtocolVersion = "2025-06-18"

type ServerOptions struct {
	Registry        *registry.Registry
	OriginAllowlist []string
	Sessions        *SessionStore
}

type Server struct {
	opts ServerOptions
}

func NewServer(opts ServerOptions) *Server {
	if opts.Sessions == nil {
		opts.Sessions = NewSessionStore()
	}
	return &Server{opts: opts}
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Origin allowlist (no-op when allowlist empty or Origin missing)
	if len(s.opts.OriginAllowlist) > 0 {
		origin := r.Header.Get("Origin")
		if origin != "" {
			ok := false
			for _, a := range s.opts.OriginAllowlist {
				if a == origin {
					ok = true
					break
				}
			}
			if !ok {
				http.Error(w, "origin not allowed", http.StatusForbidden)
				return
			}
		}
	}

	switch r.Method {
	case http.MethodGet:
		s.handleGET(w, r)
	case http.MethodPost:
		s.handlePOST(w, r)
	default:
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
	}
}

func (s *Server) handlePOST(w http.ResponseWriter, r *http.Request) {
	var req Request
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, NewErrorResponse(nil, -32700, "parse error", nil))
		return
	}
	if req.ID == nil {
		// Notification — no response.
		w.WriteHeader(http.StatusAccepted)
		return
	}

	// initialize is the only method allowed without a session.
	if req.Method == "initialize" {
		sid := s.opts.Sessions.Issue()
		w.Header().Set("Mcp-Session-Id", sid)
		writeJSON(w, http.StatusOK, NewResponse(req.ID, map[string]any{
			"protocolVersion": ProtocolVersion,
			"capabilities":    map[string]any{"tools": map[string]any{}},
			"serverInfo":      map[string]any{"name": "gitea-mcp", "version": "0.1.0"},
		}))
		return
	}

	sid := r.Header.Get("Mcp-Session-Id")
	if !s.opts.Sessions.Valid(sid) {
		http.Error(w, "missing or invalid Mcp-Session-Id", http.StatusBadRequest)
		return
	}

	switch req.Method {
	case "tools/list":
		writeJSON(w, http.StatusOK, NewResponse(req.ID, map[string]any{
			"tools": s.opts.Registry.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 {
			writeJSON(w, http.StatusOK,
				NewErrorResponse(req.ID, -32602, "invalid params", nil))
			return
		}
		out, err := s.opts.Registry.Dispatch(r.Context(), p.Name, p.Arguments)
		if err != nil {
			code := -32000
			if errors.Is(err, ErrToolNotFound) {
				code = CodeNotFound
			}
			writeJSON(w, http.StatusOK,
				NewErrorResponse(req.ID, code, err.Error(), nil))
			return
		}
		writeJSON(w, http.StatusOK, NewResponse(req.ID, map[string]any{
			"content": []map[string]any{{"type": "text", "text": string(out)}},
		}))

	default:
		writeJSON(w, http.StatusOK,
			NewErrorResponse(req.ID, -32601, "method not found: "+req.Method, nil))
	}
}

func (s *Server) handleGET(w http.ResponseWriter, r *http.Request) {
	sid := r.Header.Get("Mcp-Session-Id")
	if !s.opts.Sessions.Valid(sid) {
		http.Error(w, "missing or invalid Mcp-Session-Id", http.StatusBadRequest)
		return
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	flusher, _ := w.(http.Flusher)
	// Emit a comment as keepalive; real notifications come via a future channel.
	_, _ = w.Write([]byte(": stream open\n\n"))
	if flusher != nil {
		flusher.Flush()
	}
	<-r.Context().Done()
}

var ErrToolNotFound = errors.New("tool not found")

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}