nofx/mcp/x402.go

package mcp

import (
	"bytes"
	"crypto/ecdsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"github.com/ethereum/go-ethereum/crypto"
)

// ── Shared x402 types ────────────────────────────────────────────────────────

// x402v2PaymentRequired is the structure of the Payment-Required header (x402 v2).
type x402v2PaymentRequired struct {
	X402Version int              `json:"x402Version"`
	Accepts     []x402AcceptOption `json:"accepts"`
	Resource    *x402Resource    `json:"resource"`
}

// x402AcceptOption is a payment option from the x402 v2 header.
type x402AcceptOption struct {
	Scheme            string            `json:"scheme"`
	Network           string            `json:"network"`
	Amount            string            `json:"amount"`
	Asset             string            `json:"asset"`
	PayTo             string            `json:"payTo"`
	MaxTimeoutSeconds int               `json:"maxTimeoutSeconds"`
	Extra             map[string]string `json:"extra"`
}

// x402Resource describes the resource being paid for.
type x402Resource struct {
	URL         string `json:"url"`
	Description string `json:"description"`
	MimeType    string `json:"mimeType"`
}

// x402SignFunc is a callback that signs an x402 payment header and returns the
// base64-encoded payment signature.
type x402SignFunc func(paymentHeaderB64 string) (string, error)

// ── Shared x402 helpers ──────────────────────────────────────────────────────

// x402DecodeHeader decodes a base64-encoded x402 Payment-Required header,
// trying RawStdEncoding first then StdEncoding as fallback.
func x402DecodeHeader(b64 string) ([]byte, error) {
	decoded, err := base64.RawStdEncoding.DecodeString(b64)
	if err != nil {
		decoded, err = base64.StdEncoding.DecodeString(b64)
		if err != nil {
			return nil, fmt.Errorf("failed to base64-decode payment header: %w", err)
		}
	}
	return decoded, nil
}

// signBasePaymentHeader decodes a base64 x402 header, parses it, and signs with
// EIP-712 (USDC TransferWithAuthorization). Shared by BlockRunBase and Claw402.
func signBasePaymentHeader(privateKey *ecdsa.PrivateKey, paymentHeaderB64 string, providerName string) (string, error) {
	if privateKey == nil {
		return "", fmt.Errorf("no private key set for %s wallet", providerName)
	}

	decoded, err := x402DecodeHeader(paymentHeaderB64)
	if err != nil {
		return "", err
	}

	var req x402v2PaymentRequired
	if err := json.Unmarshal(decoded, &req); err != nil {
		return "", fmt.Errorf("failed to parse x402 v2 payment header: %w", err)
	}
	if len(req.Accepts) == 0 {
		return "", fmt.Errorf("no payment options in x402 response")
	}

	senderAddr := crypto.PubkeyToAddress(privateKey.PublicKey).Hex()
	return signX402Payment(privateKey, senderAddr, req.Accepts[0], req.Resource)
}

// doX402Request executes an HTTP request and handles the x402 v2 payment flow.
// On a 402 response it reads the Payment-Required (or X-Payment-Required) header,
// signs via signFn, retries with Payment-Signature, and logs the Payment-Response
// header (tx hash) on success.
func doX402Request(
	httpClient *http.Client,
	buildReqFn func() (*http.Request, error),
	signFn x402SignFunc,
	providerTag string,
	logger Logger,
) ([]byte, error) {
	req, err := buildReqFn()
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("failed to send request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusPaymentRequired {
		paymentHeader := resp.Header.Get("Payment-Required")
		if paymentHeader == "" {
			paymentHeader = resp.Header.Get("X-Payment-Required")
		}
		if paymentHeader == "" {
			body, _ := io.ReadAll(resp.Body)
			return nil, fmt.Errorf("received 402 but no Payment-Required header found. Body: %s", string(body))
		}

		// Drain 402 body to allow HTTP connection reuse.
		_, _ = io.Copy(io.Discard, resp.Body)

		paymentSig, err := signFn(paymentHeader)
		if err != nil {
			return nil, fmt.Errorf("failed to sign x402 payment: %w", err)
		}

		req2, err := buildReqFn()
		if err != nil {
			return nil, fmt.Errorf("failed to build retry request: %w", err)
		}
		req2.Header.Set("X-Payment", paymentSig)
		req2.Header.Set("Payment-Signature", paymentSig)

		resp2, err := httpClient.Do(req2)
		if err != nil {
			return nil, fmt.Errorf("failed to send payment retry: %w", err)
		}
		defer resp2.Body.Close()

		body2, err := io.ReadAll(resp2.Body)
		if err != nil {
			return nil, fmt.Errorf("failed to read payment retry response: %w", err)
		}
		if resp2.StatusCode != http.StatusOK {
			return nil, fmt.Errorf("%s payment retry failed (status %d): %s", providerTag, resp2.StatusCode, string(body2))
		}

		if txHash := resp2.Header.Get("Payment-Response"); txHash != "" {
			logger.Infof("💰 [%s] Payment tx: %s", providerTag, txHash)
		}

		return body2, nil
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("%s API error (status %d): %s", providerTag, resp.StatusCode, string(body))
	}
	return body, nil
}

// x402BuildRequest creates a POST request with Content-Type but no auth header.
func x402BuildRequest(url string, jsonData []byte) (*http.Request, error) {
	req, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, fmt.Errorf("fail to build request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")
	return req, nil
}

// x402SetAuthHeader is a no-op — x402 providers authenticate via payment signing.
func x402SetAuthHeader(_ http.Header) {}

// x402Call handles the x402 payment flow for the simple CallWithMessages path.
func x402Call(c *Client, signFn x402SignFunc, tag string, systemPrompt, userPrompt string) (string, error) {
	c.logger.Infof("📡 [%s] Request AI Server: %s", tag, c.BaseURL)

	requestBody := c.hooks.buildMCPRequestBody(systemPrompt, userPrompt)
	jsonData, err := c.hooks.marshalRequestBody(requestBody)
	if err != nil {
		return "", err
	}

	body, err := doX402Request(c.httpClient, func() (*http.Request, error) {
		return c.hooks.buildRequest(c.hooks.buildUrl(), jsonData)
	}, signFn, tag, c.logger)
	if err != nil {
		return "", err
	}
	return c.hooks.parseMCPResponse(body)
}

// x402CallFull handles the x402 payment flow for the advanced Request path.
func x402CallFull(c *Client, signFn x402SignFunc, tag string, req *Request) (*LLMResponse, error) {
	if c.APIKey == "" {
		return nil, fmt.Errorf("AI API key not set, please call SetAPIKey first")
	}
	if req.Model == "" {
		req.Model = c.Model
	}

	c.logger.Infof("📡 [%s] Request AI (full): %s", tag, c.BaseURL)

	requestBody := c.hooks.buildRequestBodyFromRequest(req)
	jsonData, err := c.hooks.marshalRequestBody(requestBody)
	if err != nil {
		return nil, err
	}

	body, err := doX402Request(c.httpClient, func() (*http.Request, error) {
		return c.hooks.buildRequest(c.hooks.buildUrl(), jsonData)
	}, signFn, tag, c.logger)
	if err != nil {
		return nil, err
	}
	return c.hooks.parseMCPResponseFull(body)
}