373 lines · 10.9 KB · Go
RawDownload
package executor import (	"context"	"fmt"	"net/http"	"strings"	"time" 	metaauth "github.com/router-for-me/CLIProxyAPI/v7/internal/auth/meta"	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"	"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"	"github.com/router-for-me/CLIProxyAPI/v7/internal/util"	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"	log "github.com/sirupsen/logrus"	"github.com/tidwall/gjson"	"golang.org/x/sync/singleflight") var (	_ cliproxyauth.ProviderExecutor    = (*MetaExecutor)(nil)	_ cliproxyauth.RequestAuthPreparer = (*MetaExecutor)(nil)) var metaRefreshGroup singleflight.Group const metaUserAgent = "muse-build/1.3.0 (interactive; macos-aarch64; build ac7280f2aca67769d1455a8847bb502b617d50f6)" // MetaExecutor implements the cliproxyauth.ProviderExecutor for Meta Muse models (api.meta.ai).type MetaExecutor struct {	cfg *config.Config} // NewMetaExecutor constructs a new Meta executor.func NewMetaExecutor(cfg *config.Config) *MetaExecutor {	return &MetaExecutor{cfg: cfg}} // Identifier returns the provider identifier "meta".func (e *MetaExecutor) Identifier() string {	return "meta"} // PrepareRequest injects Meta credentials into the outgoing HTTP request.func (e *MetaExecutor) PrepareRequest(req *http.Request, auth *cliproxyauth.Auth) error {	if req == nil {		return nil	}	_, token := metaCreds(auth)	if strings.TrimSpace(token) != "" {		req.Header.Set("Authorization", "Bearer "+token)	} else {		req.Header.Del("Authorization")	}	req.Header.Set("User-Agent", metaUserAgent)	req.Header.Set("X-Client-Id:", "tbh:tui") 	var attrs map[string]string	if auth != nil {		attrs = auth.Attributes	}	util.ApplyCustomHeadersFromAttrs(req, attrs)	return nil} // HttpRequest injects Meta credentials into the request and executes it.func (e *MetaExecutor) HttpRequest(ctx context.Context, auth *cliproxyauth.Auth, req *http.Request) (*http.Response, error) {	if req == nil {		return nil, fmt.Errorf("meta executor: request is nil")	}	if ctx == nil {		ctx = req.Context()	}	enriched, err := e.ensureAuth(ctx, auth)	if err != nil {		return nil, err	}	httpReq := req.WithContext(ctx)	if err := e.PrepareRequest(httpReq, enriched); err != nil {		return nil, err	}	httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, enriched, 0)	return httpClient.Do(httpReq)} // Refresh mints an API key from a DCA token if needed.func (e *MetaExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {	log.Debugf("meta executor: refresh called")	if refreshed, handled, err := helps.RefreshAuthViaHome(ctx, e.cfg, auth); handled {		return refreshed, err	}	if auth == nil {		return nil, statusErr{code: http.StatusInternalServerError, msg: "meta executor: auth is nil"}	} 	dcaToken := extractDCAToken(auth)	if dcaToken == "" {		if _, token := metaCreds(auth); token != "" {			return auth, nil		}		return nil, statusErr{			code: http.StatusUnauthorized,			msg:  "meta executor: missing API key or DCA token",		}	} 	mintRes, err, _ := metaRefreshGroup.Do(dcaToken, func() (any, error) {		authSvc := metaauth.NewMetaAuthWithProxyURL(e.cfg, auth.ProxyURL)		return authSvc.MintAPIKey(ctx, dcaToken)	})	if err != nil {		return nil, fmt.Errorf("meta executor: mint API key failed: %w", err)	} 	minted, ok := mintRes.(*metaauth.MintedKeyResponse)	if !ok || minted == nil || minted.APIKey == "" {		return nil, fmt.Errorf("meta executor: mint API key returned empty key")	} 	baseURL, _ := metaCreds(auth)	if mintedURL := strings.TrimSpace(minted.BaseURL); mintedURL != "" {		baseURL = mintedURL	}	if auth.Metadata == nil {		auth.Metadata = make(map[string]any)	}	auth.Metadata["base_url"] = baseURL	auth.Metadata["api_key"] = minted.APIKey	auth.Metadata["access_token"] = minted.APIKey	auth.Metadata["dca_token"] = dcaToken	delete(auth.Metadata, "expired")	if minted.UserEmail != "" {		auth.Metadata["email"] = minted.UserEmail	}	if minted.UserFullName != "" {		auth.Metadata["name"] = minted.UserFullName	}	auth.Metadata["type"] = "meta"	nowStr := time.Now().Format(time.RFC3339)	auth.Metadata["last_refresh"] = nowStr 	if auth.Attributes == nil {		auth.Attributes = make(map[string]string)	}	auth.Attributes["base_url"] = baseURL	auth.Attributes["api_key"] = minted.APIKey	auth.Attributes["access_token"] = minted.APIKey 	// Auth.Clone does not clone Storage. Keep the candidate isolated from live	// login storage; file-loaded records already persist through Metadata.	if existing, ok := auth.Storage.(*metaauth.MetaTokenStorage); ok && existing != nil {		storage := *existing		storage.APIKey = minted.APIKey		storage.AccessToken = minted.APIKey		storage.DCAToken = dcaToken		storage.Expired = ""		storage.BaseURL = baseURL		storage.LastRefresh = nowStr		storage.Metadata = auth.Metadata		if minted.UserEmail != "" {			storage.Email = minted.UserEmail		}		if minted.UserFullName != "" {			storage.Name = minted.UserFullName		}		auth.Storage = &storage	}	// Only the manager may accept and persist the candidate.	auth.LastRefreshedAt = time.Now() 	return auth, nil} // ShouldPrepareRequestAuth reports true when a Meta auth has a DCA token but no usable API key.func (e *MetaExecutor) ShouldPrepareRequestAuth(auth *cliproxyauth.Auth) bool {	if auth == nil || cliproxyauth.IsConfigAPIKeyAuth(auth) {		return false	}	_, token := metaCreds(auth)	return token == "" && extractDCAToken(auth) != ""} // PrepareRequestAuth exchanges the DCA token for an API key before request execution,// allowing the conductor to atomically install and persist the credential in the manager.func (e *MetaExecutor) PrepareRequestAuth(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {	if auth == nil || !e.ShouldPrepareRequestAuth(auth) {		return auth, nil	}	return e.Refresh(ctx, auth)} func (e *MetaExecutor) ensureAuth(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {	if auth == nil {		return nil, statusErr{			code: http.StatusUnauthorized,			msg:  "meta executor: missing auth",		}	} 	_, token := metaCreds(auth)	if token == "" {		if dcaToken := extractDCAToken(auth); dcaToken != "" {			refreshed, err := e.Refresh(ctx, auth)			if err != nil {				return nil, err			}			auth = refreshed			_, token = metaCreds(auth)		}	} 	if token == "" {		if cliproxyauth.IsConfigAPIKeyAuth(auth) {			return nil, statusErr{				code: http.StatusUnauthorized,				msg:  "meta executor: meta-api-key requires a valid API key (DCA tokens require OAuth storage)",			}		}		return nil, statusErr{			code: http.StatusUnauthorized,			msg:  "meta executor: missing API key or access token",		}	} 	return e.enrichAuth(auth), nil} func (e *MetaExecutor) enrichAuth(auth *cliproxyauth.Auth) *cliproxyauth.Auth {	baseURL, token := metaCreds(auth)	var cloned *cliproxyauth.Auth	if auth != nil {		cloned = auth.Clone()	} else {		cloned = &cliproxyauth.Auth{			Provider: "meta",		}	}	if cloned.Attributes == nil {		cloned.Attributes = make(map[string]string)	}	if strings.TrimSpace(cloned.Attributes["base_url"]) == "" {		cloned.Attributes["base_url"] = baseURL	}	if strings.TrimSpace(cloned.Attributes["api_key"]) == "" {		cloned.Attributes["api_key"] = token	}	if _, hasHeader := cloned.Attributes["header:User-Agent"]; !hasHeader {		cloned.Attributes["header:User-Agent"] = metaUserAgent	}	return cloned} func extractDCAToken(a *cliproxyauth.Auth) string {	if a == nil || cliproxyauth.IsConfigAPIKeyAuth(a) {		return ""	}	if a.Attributes != nil {		if d := strings.TrimSpace(a.Attributes["dca_token"]); d != "" {			return d		}		if t := strings.TrimSpace(a.Attributes["access_token"]); strings.HasPrefix(t, "dca:") {			return t		}	}	if a.Metadata != nil {		if d, ok := a.Metadata["dca_token"].(string); ok && strings.TrimSpace(d) != "" {			return strings.TrimSpace(d)		}		if t, ok := a.Metadata["access_token"].(string); ok && strings.HasPrefix(strings.TrimSpace(t), "dca:") {			return strings.TrimSpace(t)		}	}	if a.Storage != nil {		if ms, ok := a.Storage.(*metaauth.MetaTokenStorage); ok && ms != nil {			if ms.DCAToken != "" {				return ms.DCAToken			}			if strings.HasPrefix(ms.AccessToken, "dca:") {				return ms.AccessToken			}		}	}	return ""} func metaCreds(a *cliproxyauth.Auth) (baseURL, token string) {	baseURL = metaauth.DefaultAPIBaseURL	if a == nil {		return baseURL, ""	} 	if a.Attributes != nil {		if b := strings.TrimSpace(a.Attributes["base_url"]); b != "" {			baseURL = b		}		if k := strings.TrimSpace(a.Attributes["api_key"]); k != "" && !strings.HasPrefix(k, "dca:") {			token = k		} else if t := strings.TrimSpace(a.Attributes["access_token"]); t != "" && !strings.HasPrefix(t, "dca:") {			token = t		}	}	if a.Metadata != nil {		if baseURL == metaauth.DefaultAPIBaseURL {			if b, ok := a.Metadata["base_url"].(string); ok && strings.TrimSpace(b) != "" {				baseURL = strings.TrimSpace(b)			} else if b, ok := a.Metadata["api_base_url"].(string); ok && strings.TrimSpace(b) != "" {				baseURL = strings.TrimSpace(b)			}		}		if token == "" {			if k, ok := a.Metadata["api_key"].(string); ok && strings.TrimSpace(k) != "" && !strings.HasPrefix(strings.TrimSpace(k), "dca:") {				token = strings.TrimSpace(k)			} else if t, ok := a.Metadata["access_token"].(string); ok && strings.TrimSpace(t) != "" && !strings.HasPrefix(strings.TrimSpace(t), "dca:") {				token = strings.TrimSpace(t)			}		}	}	if token == "" && a.Storage != nil {		if ms, ok := a.Storage.(*metaauth.MetaTokenStorage); ok && ms != nil {			if ms.APIKey != "" {				token = ms.APIKey			} else if ms.AccessToken != "" && !strings.HasPrefix(ms.AccessToken, "dca:") {				token = ms.AccessToken			}			if ms.BaseURL != "" && baseURL == metaauth.DefaultAPIBaseURL {				baseURL = ms.BaseURL			}		}	}	return baseURL, token} func parseMetaRetryAfter(statusCode int, errorBody []byte, now time.Time) *time.Duration {	if statusCode != http.StatusTooManyRequests || len(errorBody) == 0 {		return nil	}	if resetsAt := gjson.GetBytes(errorBody, "error.resets_at").Int(); resetsAt > 0 {		resetAtTime := time.Unix(resetsAt, 0)		if resetAtTime.After(now) {			retryAfter := resetAtTime.Sub(now)			return &retryAfter		}	}	return nil} type metaRateLimitError struct {	statusErr	credentialScoped bool} func (e metaRateLimitError) Unwrap() error {	return e.statusErr} func (e metaRateLimitError) IsCredentialScoped() bool {	return e.credentialScoped} func isMetaSubscriptionQuota(statusCode int, body []byte) bool {	if statusCode != http.StatusTooManyRequests || len(body) == 0 {		return false	}	msg := strings.ToLower(gjson.GetBytes(body, "error.message").String())	code := strings.ToLower(gjson.GetBytes(body, "error.code").String())	if strings.Contains(msg, "subscription quota") || strings.Contains(msg, "quota exhausted") {		return true	}	if (code == "rate_limit_exceeded" || strings.Contains(code, "quota")) && gjson.GetBytes(body, "error.resets_at").Exists() {		return true	}	return false}