naco-api/internal/x402/x402.go

package x402

import (
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"os"
	"reflect"
	"strings"

	"github.com/google/jsonschema-go/jsonschema"
	"github.com/modelcontextprotocol/go-sdk/mcp"
	log "github.com/nacorid/logger"
	x402 "github.com/x402-foundation/x402/go/v2"
	x402http "github.com/x402-foundation/x402/go/v2/http"
	nethttp "github.com/x402-foundation/x402/go/v2/http/nethttp"
	x402mcp "github.com/x402-foundation/x402/go/v2/mcp"
	evm "github.com/x402-foundation/x402/go/v2/mechanisms/evm/exact/server"
)

type X402Middlewares struct {
	MCPMiddleware  []mcp.Middleware
	HTTPMiddleware []func(http.Handler) http.Handler
}

type ToolInfo struct {
	Tool         mcp.Tool
	IsPayable    bool
	Requirements []x402.PaymentRequirements
	Handler      mcp.ToolHandler
}

type X402Server struct {
	mcpServer          *mcp.Server
	x402ResourceServer *x402.X402ResourceServer
	facilitatorClient  *x402http.HTTPFacilitatorClient
	httpSettlementHook func(w http.ResponseWriter, r *http.Request, resp *x402.SettleResponse)
	mux                *http.ServeMux

	name            string
	version         string
	middlewareChain func(http.Handler) http.Handler
	tools           []ToolInfo
	globalReqLog    func(http.Handler) http.Handler
}

const VERSION1 = "v1/"

func NewX402Server(name, version string, facilitatorURL string, auth x402http.AuthProvider, mws X402Middlewares, globalLoggingMw func(http.Handler) http.Handler, settleHook func(http.ResponseWriter, *http.Request, *x402.SettleResponse)) (*X402Server, error) {
	mcpServer := mcp.NewServer(
		&mcp.Implementation{
			Name:    name,
			Version: version,
		},
		nil,
	)
	for _, mw := range mws.MCPMiddleware {
		mcpServer.AddReceivingMiddleware(mw)
	}

	if facilitatorURL == "" {
		facilitatorURL = "https://facilitator.mogami.tech"
	}
	facilitator := x402http.NewHTTPFacilitatorClient(&x402http.FacilitatorConfig{
		URL:          facilitatorURL,
		AuthProvider: auth,
	})

	resourceServer := x402.Newx402ResourceServer(
		x402.WithFacilitatorClient(facilitator),
	)

	resourceServer.Register("eip155:8453", evm.NewExactEvmScheme())

	s := &X402Server{
		mcpServer:          mcpServer,
		x402ResourceServer: resourceServer,
		facilitatorClient:  facilitator,
		httpSettlementHook: settleHook,
		mux:                http.NewServeMux(),
		name:               name,
		version:            version,
		middlewareChain:    chainHTTPMiddlewares(mws.HTTPMiddleware...),
		tools:              make([]ToolInfo, 0),
		globalReqLog:       globalLoggingMw,
	}

	s.mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})

	s.mux.HandleFunc("GET /openapi.json", s.serveOpenAPI)

	mcpHandler := mcp.NewStreamableHTTPHandler(func(req *http.Request) *mcp.Server {
		return s.mcpServer
	}, nil)
	s.mux.Handle("/mcp", mcpHandler)

	return s, nil
}

func (s *X402Server) AddTool(tool mcp.Tool, handler mcp.ToolHandler) {
	tool.InputSchema = toJSONSchema(tool.InputSchema)
	tool.OutputSchema = toJSONSchema(tool.OutputSchema)

	s.mcpServer.AddTool(&tool, handler)

	s.tools = append(s.tools, ToolInfo{
		Tool:      tool,
		IsPayable: false,
		Handler:   handler,
	})

	restHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s.handleRESTCall(w, r, tool.Name, handler)
	})
	s.mux.Handle("POST /"+VERSION1+tool.Name, s.middlewareChain(restHandler))
}

func (s *X402Server) AddPayableTool(
	tool mcp.Tool,
	handler mcp.ToolHandler,
	requirements ...x402.PaymentRequirements,
) {

	if len(requirements) == 0 {
		log.Fatalf("tool %s requires at least one payment requirement", tool.Name)
	}

	tool.InputSchema = toJSONSchema(tool.InputSchema)
	tool.OutputSchema = toJSONSchema(tool.OutputSchema)

	wrapper := x402mcp.NewPaymentWrapper(s.x402ResourceServer, x402mcp.PaymentWrapperConfig{
		Accepts: requirements,
		Resource: &x402mcp.ResourceInfo{
			URL:         "mcp://tool/" + tool.Name,
			Description: tool.Description,
		},
	})

	s.mcpServer.AddTool(&tool, wrapper.Wrap(handler))

	s.tools = append(s.tools, ToolInfo{
		Tool:         tool,
		IsPayable:    true,
		Requirements: requirements,
		Handler:      handler,
	})

	var options x402http.PaymentOptions
	for _, req := range requirements {
		options = append(options, x402http.PaymentOption{
			Scheme:  req.Scheme,
			PayTo:   req.PayTo,
			Price:   req.Amount,
			Network: x402.Network(req.Network),
		})
	}

	paymentMiddleware := nethttp.X402Payment(nethttp.Config{
		Routes: x402http.RoutesConfig{
			"POST /" + VERSION1 + tool.Name: {
				Accepts:     options,
				Description: tool.Description,
			},
		},
		Facilitator: s.facilitatorClient,
		Schemes: []nethttp.SchemeConfig{
			{Network: "eip155:8453", Server: evm.NewExactEvmScheme()},
		},
		SettlementHandler: s.httpSettlementHook,
	})

	restHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s.handleRESTCall(w, r, tool.Name, handler)
	})
	s.mux.Handle("POST /"+VERSION1+tool.Name, s.middlewareChain(paymentMiddleware(restHandler)))
}

func (s *X402Server) handleRESTCall(w http.ResponseWriter, r *http.Request, toolName string, handler mcp.ToolHandler) {
	bodyBytes, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "Failed to read body", http.StatusBadRequest)
		return
	}

	log.DebugWithContext(r.Context(), "RESTful API request.", "RemoteAddr", r.RemoteAddr, "X-Forwarded-For", r.Header.Get("X-Forwarded-For"))

	callReq := &mcp.CallToolRequest{
		Params: &mcp.CallToolParamsRaw{
			Name:      toolName,
			Arguments: json.RawMessage(bodyBytes),
		},
	}

	resp, err := handler(r.Context(), callReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)

	if len(resp.Content) == 1 {
		if textContent, ok := resp.Content[0].(*mcp.TextContent); ok {
			var js json.RawMessage
			if json.Unmarshal([]byte(textContent.Text), &js) == nil {
				w.Write([]byte(textContent.Text))
				return
			}
		}
	}

	json.NewEncoder(w).Encode(resp)
}

func (s *X402Server) OnAfterSettle(hook func(ctx x402.SettleResultContext) error) {
	s.x402ResourceServer.OnAfterSettle(hook)
}

func (s *X402Server) serveOpenAPI(w http.ResponseWriter, r *http.Request) {
	paths := make(map[string]any)

	log.DebugWithContext(r.Context(), "OpenAPI Spec requested.", "RemoteAddr", r.RemoteAddr, "X-Forwarded-For", r.Header.Get("X-Forwarded-For"))

	for _, info := range s.tools {
		inputSchema := parseSchema(info.Tool.InputSchema)
		outputSchema := parseSchema(info.Tool.OutputSchema)

		postOp := map[string]any{
			"summary":     info.Tool.Description,
			"description": fmt.Sprintf("Exposes the MCP tool '%s' as a RESTful endpoint.", info.Tool.Name),
			"requestBody": map[string]any{
				"required": true,
				"content": map[string]any{
					"application/json": map[string]any{
						"schema": inputSchema,
					},
				},
			},
			"responses": map[string]any{
				"200": map[string]any{
					"description": "Successful operation",
					"content": map[string]any{
						"application/json": map[string]any{
							"schema": outputSchema,
						},
					},
				},
			},
		}

		if info.IsPayable {
			postOp["description"] = fmt.Sprintf("%s\n\n**Paid Route**: Requires payments via x402.", postOp["description"])
			postOp["responses"].(map[string]any)["402"] = map[string]any{
				"description": "Payment Required. Responds with the PAYMENT-REQUIRED header challenge.",
				"headers": map[string]any{
					"PAYMENT-REQUIRED": map[string]any{
						"schema":      map[string]any{"type": "string"},
						"description": "Base64 encoded PaymentRequired challenge matching the x402 specification.",
					},
				},
			}
		}

		paths["/"+VERSION1+info.Tool.Name] = map[string]any{
			"post": postOp,
		}
	}

	openapiDoc := map[string]any{
		"openapi": "3.0.3",
		"info": map[string]any{
			"title":       s.name,
			"version":     s.version,
			"description": "Auto-generated REST API matching active Model Context Protocol (MCP) tools.",
			"contact": map[string]any{
				"email": "x402@naco.li",
			},
		},
		"paths": paths,
	}

	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(openapiDoc)
}

func (s *X402Server) Start(addr string, permissions os.FileMode) error {
	var (
		listener net.Listener
		err      error
	)

	if socketPath, ok := strings.CutPrefix(addr, "unix:"); ok {
		_ = os.RemoveAll(socketPath)

		listener, err = net.Listen("unix", socketPath)
		if err != nil {
			return fmt.Errorf("failed to listen on unix socket: %w", err)
		}
		if err := os.Chmod(socketPath, permissions); err != nil {
			return fmt.Errorf("failed to set socket permissions: %w", err)
		}

		log.Info("Starting X402 MCP Server on unix socket", "socketPath", socketPath)
	} else {
		log.Info("Starting X402 MCP Server on TCP", "addr", addr)

		listener, err = net.Listen("tcp", addr)
		if err != nil {
			return fmt.Errorf("failed to listen on tcp: %w", err)
		}
	}
	defer listener.Close()

	srv := &http.Server{
		Handler:  s.globalReqLog(s.mux),
		ErrorLog: slog.NewLogLogger(slog.Default().Handler(), slog.LevelDebug),
	}
	return srv.Serve(listener)
}

func toJSONSchema(val any) json.RawMessage {
	if val == nil {
		return json.RawMessage(`{"type": "object"}`)
	}
	switch v := val.(type) {
	case json.RawMessage:
		if len(v) == 0 {
			return json.RawMessage(`{"type": "object"}`)
		}
		return v
	case []byte:
		if len(v) == 0 {
			return json.RawMessage(`{"type": "object"}`)
		}
		return json.RawMessage(v)
	case string:
		if v == "" {
			return json.RawMessage(`{"type": "object"}`)
		}
		return json.RawMessage(v)
	}

	t := reflect.TypeOf(val)
	schema, err := jsonschema.ForType(t, &jsonschema.ForOptions{})
	if err != nil {
		return json.RawMessage(`{"type": "object"}`)
	}
	bytes, err := json.Marshal(schema)
	if err != nil {
		return json.RawMessage(`{"type": "object"}`)
	}
	return bytes
}

func parseSchema(schemaAny any) any {
	if schemaAny == nil {
		return map[string]any{"type": "object"}
	}
	switch v := schemaAny.(type) {
	case json.RawMessage:
		var m any
		if err := json.Unmarshal(v, &m); err == nil {
			return m
		}
	case []byte:
		var m any
		if err := json.Unmarshal(v, &m); err == nil {
			return m
		}
	case string:
		var m any
		if err := json.Unmarshal([]byte(v), &m); err == nil {
			return m
		}
	}
	return schemaAny
}

func chainHTTPMiddlewares(middlewares ...func(http.Handler) http.Handler) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		for i := len(middlewares) - 1; i >= 0; i-- {
			next = middlewares[i](next)
		}
		return next
	}
}