naco-api/internal/x402/mcpwrapper.go

package x402

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	log "github.com/nacorid/logger"
	x402 "github.com/x402-foundation/x402/go/v2"
	"github.com/x402-foundation/x402/go/v2/types"
)

// Payment meta keys exchanged over MCP, mirroring the x402 SDK contract.
const (
	PaymentMetaKey         = "x402/payment"
	PaymentResponseMetaKey = "x402/payment-response"
)

// SettleAmountResolver computes the actual settlement amount for a completed
// tool call. It returns ("", false) to settle the amount advertised in the
// payment requirements as-is.
type SettleAmountResolver func(result *mcp.CallToolResult) (string, bool)

// PaymentWrapperConfig configures a PaymentWrapper. Unlike the SDK's wrapper,
// it accepts a SettleAmountResolver so upto-scheme tools can charge the actual
// computed cost after the handler runs (the same behavior the REST path gets
// via nethttp.SetSettlementOverrides).
type PaymentWrapperConfig struct {
	Accepts    []x402.PaymentRequirements
	Resource   *types.ResourceInfo
	Extensions map[string]interface{}
	// SettleAmountResolver is invoked after a successful handler call. When it
	// returns a non-empty amount, that amount is charged instead of the
	// advertised maximum (which must therefore be <= the accept amount).
	SettleAmountResolver SettleAmountResolver
}

// PaymentWrapper gates an MCP tool behind x402 payments and settles them.
type PaymentWrapper struct {
	server *x402.X402ResourceServer
	config PaymentWrapperConfig
}

// NewPaymentWrapper creates a PaymentWrapper for the given resource server.
func NewPaymentWrapper(server *x402.X402ResourceServer, config PaymentWrapperConfig) *PaymentWrapper {
	if len(config.Accepts) == 0 {
		panic("PaymentWrapperConfig.Accepts must have at least one payment requirement")
	}
	return &PaymentWrapper{server: server, config: config}
}

// Wrap returns a ToolHandler that requires a valid x402 payment before invoking
// the wrapped handler and settles it afterwards.
func (w *PaymentWrapper) Wrap(handler mcp.ToolHandler) mcp.ToolHandler {
	return func(ctx context.Context, request *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
		paymentData := extractPaymentFromRequest(request)
		if paymentData == nil {
			return w.paymentRequiredResult("Payment Required"), nil
		}

		payloadBytes, err := json.Marshal(paymentData)
		if err != nil {
			return w.paymentRequiredResult(fmt.Sprintf("Invalid payment payload: %v", err)), nil
		}
		var payload types.PaymentPayload
		if err := json.Unmarshal(payloadBytes, &payload); err != nil {
			return w.paymentRequiredResult(fmt.Sprintf("Invalid payment payload: %v", err)), nil
		}

		matched := w.server.FindMatchingRequirements(w.config.Accepts, payload)
		if matched == nil {
			return w.paymentRequiredResult("No matching payment requirements found"), nil
		}
		requirements := *matched

		verifyResp, err := w.server.VerifyPayment(ctx, payload, requirements)
		if err != nil {
			return w.paymentRequiredResult(fmt.Sprintf("Payment verification error: %v", err)), nil
		}
		if !verifyResp.IsValid {
			return w.paymentRequiredResult(fmt.Sprintf("Payment verification failed: %s", verifyResp.InvalidReason)), nil
		}

		result, err := handler(ctx, request)
		if err != nil {
			return nil, err
		}
		if result.IsError {
			return result, nil
		}

		var overrides *x402.SettlementOverrides
		if w.config.SettleAmountResolver != nil {
			if amount, ok := w.config.SettleAmountResolver(result); ok && amount != "" {
				overrides = &x402.SettlementOverrides{Amount: amount}
				log.Debug("MCP settlement override", "amount", amount)
			}
		}

		settleResp, err := w.server.SettlePayment(ctx, payload, requirements, overrides)
		if err != nil {
			return w.settlementFailedResult(fmt.Sprintf("Settlement error: %v", err)), nil
		}
		if !settleResp.Success {
			return w.settlementFailedResult(fmt.Sprintf("Settlement failed: %s", settleResp.ErrorReason)), nil
		}

		if result.Meta == nil {
			result.Meta = mcp.Meta{}
		}
		result.Meta[PaymentResponseMetaKey] = settleResp
		return result, nil
	}
}

func extractPaymentFromRequest(request *mcp.CallToolRequest) interface{} {
	if request == nil || request.Params.Meta == nil {
		return nil
	}
	return request.Params.Meta[PaymentMetaKey]
}

func (w *PaymentWrapper) paymentRequiredResult(errorMessage string) *mcp.CallToolResult {
	var resource *types.ResourceInfo
	if w.config.Resource != nil {
		resource = &types.ResourceInfo{
			URL:         w.config.Resource.URL,
			Description: w.config.Resource.Description,
		}
	}
	payment := types.PaymentRequired{
		X402Version: 2,
		Error:       errorMessage,
		Resource:    resource,
		Accepts:     w.config.Accepts,
		Extensions:  w.config.Extensions,
	}
	b, err := json.Marshal(payment)
	if err != nil {
		panic(err)
	}
	return &mcp.CallToolResult{
		IsError: true,
		Content: []mcp.Content{&mcp.TextContent{Text: string(b)}},
	}
}

func (w *PaymentWrapper) settlementFailedResult(errorMessage string) *mcp.CallToolResult {
	return &mcp.CallToolResult{
		IsError: true,
		Content: []mcp.Content{
			&mcp.TextContent{Text: fmt.Sprintf("{\"x402Version\":2,\"error\":%q}", errorMessage)},
		},
	}
}