projax/mcp/handler.go

// Package mcp implements a small MCP-protocol server over HTTP. The wire
// format is JSON-RPC 2.0 with the MCP method set (initialize, tools/list,
// tools/call). Designed to be mounted under /mcp/rpc on the projax web
// binary; a stdio bridge (see ~/.claude/mcp/projax.sh) lets standard MCP
// clients talk to it transparently.
package mcp

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"
)

// ProtocolVersion is the MCP wire version this server speaks. Clients that
// initialize with a different version are still answered; they're expected
// to negotiate down.
const ProtocolVersion = "2024-11-05"

// JSON-RPC 2.0 error codes (subset).
const (
	codeParseError     = -32700
	codeInvalidRequest = -32600
	codeMethodNotFound = -32601
	codeInvalidParams  = -32602
	codeInternalError  = -32603
)

// Tool describes one callable tool exposed through tools/list and tools/call.
type Tool struct {
	Name        string          // e.g. "list_items"
	Description string          // one-line description for the client
	InputSchema json.RawMessage // JSON-schema object describing the params
	Handler     ToolHandler
}

// ToolHandler runs the actual work for a tool. params is the raw JSON object
// the client supplied as the tool arguments. Returning a non-nil result is
// wrapped as a structured text content block; returning an error becomes an
// MCP "isError: true" reply.
type ToolHandler func(ctx context.Context, params json.RawMessage) (any, error)

// Server holds the registered tools + the auth token. Mount via Routes() on
// any *http.ServeMux.
type Server struct {
	Name    string
	Version string
	Token   string // Bearer token; empty means "no auth" (tests only)
	Logger  *slog.Logger
	tools   map[string]Tool
}

// New builds an MCP server with no tools registered.
func New(name, version, token string, logger *slog.Logger) *Server {
	if logger == nil {
		logger = slog.Default()
	}
	return &Server{
		Name:    name,
		Version: version,
		Token:   token,
		Logger:  logger,
		tools:   map[string]Tool{},
	}
}

// Register adds a tool. Duplicate names overwrite.
func (s *Server) Register(t Tool) { s.tools[t.Name] = t }

// Routes registers /rpc on the given mux prefix-relative path. Caller mounts
// at /mcp so the resulting URL is /mcp/rpc.
func (s *Server) Routes(mux *http.ServeMux) {
	mux.HandleFunc("POST /rpc", s.handleRPC)
	// Friendly GET for ops smoke-testing — never returns secrets.
	mux.HandleFunc("GET /rpc", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{
			"server":           s.Name,
			"version":          s.Version,
			"protocolVersion":  ProtocolVersion,
			"tools":            s.toolNames(),
			"authRequired":     s.Token != "",
		})
	})
}

func (s *Server) toolNames() []string {
	out := make([]string, 0, len(s.tools))
	for n := range s.tools {
		out = append(out, n)
	}
	return out
}

// jsonRPCReq mirrors the wire shape. Method "notifications/initialized" has
// no id (notification) — we tolerate the missing id field.
type jsonRPCReq struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id,omitempty"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

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

type rpcError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Data    any    `json:"data,omitempty"`
}

func (s *Server) handleRPC(w http.ResponseWriter, r *http.Request) {
	if s.Token != "" {
		if !s.checkAuth(r) {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
	}
	body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 1<<20))
	if err != nil {
		http.Error(w, "request too large", http.StatusRequestEntityTooLarge)
		return
	}
	var req jsonRPCReq
	if err := json.Unmarshal(body, &req); err != nil {
		s.writeErr(w, nil, codeParseError, "invalid JSON: "+err.Error())
		return
	}
	if req.JSONRPC != "2.0" && req.JSONRPC != "" {
		s.writeErr(w, req.ID, codeInvalidRequest, "jsonrpc must be \"2.0\"")
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), 30*time.Second)
	defer cancel()
	switch req.Method {
	case "initialize":
		s.writeOK(w, req.ID, s.initializeResult())
	case "notifications/initialized":
		// Notifications get no response on JSON-RPC, but if the client sent an
		// id we humour it with an empty result.
		if len(req.ID) > 0 {
			s.writeOK(w, req.ID, map[string]any{})
		} else {
			w.WriteHeader(http.StatusNoContent)
		}
	case "tools/list":
		s.writeOK(w, req.ID, s.toolsListResult())
	case "tools/call":
		s.handleToolsCall(ctx, w, req)
	case "ping":
		s.writeOK(w, req.ID, map[string]any{})
	default:
		s.writeErr(w, req.ID, codeMethodNotFound, "unknown method: "+req.Method)
	}
}

func (s *Server) checkAuth(r *http.Request) bool {
	h := r.Header.Get("Authorization")
	if !strings.HasPrefix(h, "Bearer ") {
		return false
	}
	return strings.TrimSpace(h[len("Bearer "):]) == s.Token
}

func (s *Server) initializeResult() map[string]any {
	return map[string]any{
		"protocolVersion": ProtocolVersion,
		"capabilities": map[string]any{
			"tools": map[string]any{},
		},
		"serverInfo": map[string]any{
			"name":    s.Name,
			"version": s.Version,
		},
	}
}

type toolDescriptor struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"inputSchema"`
}

func (s *Server) toolsListResult() map[string]any {
	out := make([]toolDescriptor, 0, len(s.tools))
	for _, t := range s.tools {
		schema := t.InputSchema
		if len(schema) == 0 {
			schema = json.RawMessage(`{"type":"object","properties":{}}`)
		}
		out = append(out, toolDescriptor{
			Name:        t.Name,
			Description: t.Description,
			InputSchema: schema,
		})
	}
	return map[string]any{"tools": out}
}

type toolsCallParams struct {
	Name      string          `json:"name"`
	Arguments json.RawMessage `json:"arguments"`
}

func (s *Server) handleToolsCall(ctx context.Context, w http.ResponseWriter, req jsonRPCReq) {
	var p toolsCallParams
	if err := json.Unmarshal(req.Params, &p); err != nil {
		s.writeErr(w, req.ID, codeInvalidParams, "tools/call params: "+err.Error())
		return
	}
	tool, ok := s.tools[p.Name]
	if !ok {
		s.writeErr(w, req.ID, codeMethodNotFound, "unknown tool: "+p.Name)
		return
	}
	result, err := tool.Handler(ctx, p.Arguments)
	if err != nil {
		// Per MCP convention, tool errors stay inside the result envelope with
		// isError=true so the client sees them as tool failures, not transport
		// failures. JSON-RPC-level errors are reserved for transport problems
		// (auth, parse, unknown method).
		s.Logger.Warn("mcp tool error", "tool", p.Name, "err", err)
		s.writeOK(w, req.ID, map[string]any{
			"content": []map[string]any{{
				"type": "text",
				"text": err.Error(),
			}},
			"isError": true,
		})
		return
	}
	payload, err := json.Marshal(result)
	if err != nil {
		s.writeErr(w, req.ID, codeInternalError, "marshal result: "+err.Error())
		return
	}
	s.writeOK(w, req.ID, map[string]any{
		"content": []map[string]any{{
			"type": "text",
			"text": string(payload),
		}},
		"isError": false,
	})
}

func (s *Server) writeOK(w http.ResponseWriter, id json.RawMessage, result any) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(jsonRPCResp{JSONRPC: "2.0", ID: id, Result: result})
}

func (s *Server) writeErr(w http.ResponseWriter, id json.RawMessage, code int, message string) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(jsonRPCResp{JSONRPC: "2.0", ID: id, Error: &rpcError{Code: code, Message: message}})
}

// ToolError is returned by ToolHandlers for user-visible failures that should
// flow through the tool-result envelope as isError. Errors that do NOT match
// this type get wrapped automatically — this is just a sentinel for callers
// that want to provide structured user-facing data alongside the message.
type ToolError struct {
	Msg  string
	Data any
}

func (e *ToolError) Error() string { return e.Msg }

// AsToolError returns the ToolError if err is one (or wraps one).
func AsToolError(err error) (*ToolError, bool) {
	var te *ToolError
	if errors.As(err, &te) {
		return te, true
	}
	return nil, false
}

var _ = fmt.Sprintf // keep import handy if future code uses fmt