feat: implement apiToken failover mechanism

plugins/wasm-go/extensions/ai-proxy/config/config.go

@@ -1,9 +1,8 @@
 package config
 
 import (
-	"github.com/tidwall/gjson"
-
 	"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-proxy/provider"
+	"github.com/tidwall/gjson"
 )
 
 // @Name ai-proxy

plugins/wasm-go/extensions/ai-proxy/main.go

@@ -47,6 +47,10 @@ func parseGlobalConfig(json gjson.Result, pluginConfig *config.PluginConfig, log
 	if err := pluginConfig.Complete(); err != nil {
 		return err
 	}
+
+	providerConfig := pluginConfig.GetProviderConfig()
+	providerConfig.SetApiTokensFailover(log)
+
 	return nil
 }
 
@@ -62,6 +66,7 @@ func parseOverrideRuleConfig(json gjson.Result, global config.PluginConfig, plug
 	if err := pluginConfig.Complete(); err != nil {
 		return err
 	}
+
 	return nil
 }
 
@@ -88,8 +93,23 @@ func onHttpRequestHeader(ctx wrapper.HttpContext, pluginConfig config.PluginConf
 	ctx.SetContext(ctxKeyApiName, apiName)
 
 	if handler, ok := activeProvider.(provider.RequestHeadersHandler); ok {
-		// Disable the route re-calculation since the plugin may modify some headers related to  the chosen route.
+		// Disable the route re-calculation since the plugin may modify some headers related to the chosen route.
 		ctx.DisableReroute()
+
+		providerConfig := pluginConfig.GetProviderConfig()
+		apiTokenInUse := providerConfig.GetRandomToken()
+		if providerConfig.IsFailoverEnabled() {
+			// Use the health check token if it is a health check request.
+			if apiTokenHealthCheck, _ := proxywasm.GetHttpRequestHeader("ApiToken-Health-Check"); apiTokenHealthCheck != "" {
+				apiTokenInUse = apiTokenHealthCheck
+			} else {
+				// if enable apiToken failover, only use available apiToken
+				apiTokenInUse = providerConfig.GetGlobalRandomToken(log)
+			}
+		}
+		log.Debugf("[onHttpRequestHeader] use apiToken %s to send request", apiTokenInUse)
+		ctx.SetContext(provider.ApiTokenInUse, apiTokenInUse)
+
 		hasRequestBody := wrapper.HasRequestBody()
 		action, err := handler.OnRequestHeaders(ctx, apiName, log)
 		if err == nil {
@@ -101,6 +121,7 @@ func onHttpRequestHeader(ctx wrapper.HttpContext, pluginConfig config.PluginConf
 			}
 			return action
 		}
+
 		_ = util.SendResponse(500, "ai-proxy.proc_req_headers_failed", util.MimeTypeTextPlain, fmt.Sprintf("failed to process request headers: %v", err))
 		return types.ActionContinue
 	}
@@ -156,14 +177,32 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, pluginConfig config.PluginCo
 	log.Debugf("[onHttpResponseHeaders] provider=%s", activeProvider.GetProviderType())
 
 	status, err := proxywasm.GetHttpResponseHeader(":status")
+	apiTokenInUse := ctx.GetContext(provider.ApiTokenInUse).(string)
 	if err != nil || status != "200" {
 		if err != nil {
 			log.Errorf("unable to load :status header from response: %v", err)
 		}
 		ctx.DontReadResponseBody()
+
+		providerConfig := pluginConfig.GetProviderConfig()
+		// If apiToken failover is enabled and the request is not a health check request, handle unavailable apiToken.
+		if providerConfig.IsFailoverEnabled() && ctx.GetContext(provider.ApiTokenHealthCheck) == nil {
+			providerConfig.HandleUnavailableApiToken(apiTokenInUse, log)
+		}
+
 		return types.ActionContinue
 	}
 
+	// Reset ctxApiTokenRequestFailureCount if the request is successful,
+	// the apiToken is removed only when the number of consecutive request failures exceeds the threshold.
+	failureApiTokenRequestCount, _, err := provider.GetApiTokenRequestCount(provider.CtxApiTokenRequestFailureCount)
+	if err != nil {
+		log.Errorf("failed to get failureApiTokenRequestCount: %v", err)
+	}
+	if _, ok := failureApiTokenRequestCount[apiTokenInUse]; ok {
+		provider.ResetApiTokenRequestCount(provider.CtxApiTokenRequestFailureCount, apiTokenInUse, log)
+	}
+
 	if handler, ok := activeProvider.(provider.ResponseHeadersHandler); ok {
 		apiName, _ := ctx.GetContext(ctxKeyApiName).(provider.ApiName)
 		action, err := handler.OnResponseHeaders(ctx, apiName, log)

plugins/wasm-go/extensions/ai-proxy/provider/ai360.go

@@ -49,7 +49,7 @@ func (m *ai360Provider) OnRequestHeaders(ctx wrapper.HttpContext, apiName ApiNam
 	_ = util.OverwriteRequestHost(ai360Domain)
 	_ = proxywasm.RemoveHttpRequestHeader("Accept-Encoding")
 	_ = proxywasm.RemoveHttpRequestHeader("Content-Length")
-	_ = proxywasm.ReplaceHttpRequestHeader("Authorization", m.config.GetRandomToken())
+	_ = proxywasm.ReplaceHttpRequestHeader("Authorization", getApiTokenInUse(ctx))
 	// Delay the header processing to allow changing streaming mode in OnRequestBody
 	return types.HeaderStopIteration, nil
 }

plugins/wasm-go/extensions/ai-proxy/provider/baichuan.go

@@ -49,7 +49,7 @@ func (m *baichuanProvider) OnRequestHeaders(ctx wrapper.HttpContext, apiName Api
 	}
 	_ = util.OverwriteRequestPath(baichuanChatCompletionPath)
 	_ = util.OverwriteRequestHost(baichuanDomain)
-	_ = util.OverwriteRequestAuthorization("Bearer " + m.config.GetRandomToken())
+	_ = util.OverwriteRequestAuthorization("Bearer " + getApiTokenInUse(ctx))
 	_ = proxywasm.RemoveHttpRequestHeader("Content-Length")
 	return types.ActionContinue, nil
 }

plugins/wasm-go/extensions/ai-proxy/provider/baidu.go

@@ -83,7 +83,7 @@ func (b *baiduProvider) OnRequestBody(ctx wrapper.HttpContext, apiName ApiName,
 			return types.ActionContinue, errors.New("request model is empty")
 		}
 		// 根据模型重写requestPath
-		path := b.getRequestPath(request.Model)
+		path := b.getRequestPath(ctx, request.Model)
 		_ = util.OverwriteRequestPath(path)
 
 		if b.config.context == nil {
@@ -126,7 +126,7 @@ func (b *baiduProvider) OnRequestBody(ctx wrapper.HttpContext, apiName ApiName,
 	}
 	request.Model = mappedModel
 	ctx.SetContext(ctxKeyFinalRequestModel, request.Model)
-	path := b.getRequestPath(mappedModel)
+	path := b.getRequestPath(ctx, mappedModel)
 	_ = util.OverwriteRequestPath(path)
 
 	if b.config.context == nil {
@@ -226,13 +226,13 @@ type baiduTextGenRequest struct {
 	UserId          string        `json:"user_id,omitempty"`
 }
 
-func (b *baiduProvider) getRequestPath(baiduModel string) string {
+func (b *baiduProvider) getRequestPath(ctx wrapper.HttpContext, baiduModel string) string {
 	// https://cloud.baidu.com/doc/WENXINWORKSHOP/s/clntwmv7t
 	suffix, ok := baiduModelToPathSuffixMap[baiduModel]
 	if !ok {
 		suffix = baiduModel
 	}
-	return fmt.Sprintf("/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/%s?access_token=%s", suffix, b.config.GetRandomToken())
+	return fmt.Sprintf("/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/%s?access_token=%s", suffix, getApiTokenInUse(ctx))
 }
 
 func (b *baiduProvider) setSystemContent(request *baiduTextGenRequest, content string) {

plugins/wasm-go/extensions/ai-proxy/provider/claude.go

@@ -108,7 +108,7 @@ func (c *claudeProvider) OnRequestHeaders(ctx wrapper.HttpContext, apiName ApiNa
 
 	_ = util.OverwriteRequestPath(claudeChatCompletionPath)
 	_ = util.OverwriteRequestHost(claudeDomain)
-	_ = proxywasm.ReplaceHttpRequestHeader("x-api-key", c.config.GetRandomToken())
+	_ = proxywasm.ReplaceHttpRequestHeader("x-api-key", getApiTokenInUse(ctx))
 
 	if c.config.claudeVersion == "" {
 		c.config.claudeVersion = defaultVersion

plugins/wasm-go/extensions/ai-proxy/provider/cloudflare.go

@@ -49,7 +49,7 @@ func (c *cloudflareProvider) OnRequestHeaders(ctx wrapper.HttpContext, apiName A
 	}
 	_ = util.OverwriteRequestPath(strings.Replace(cloudflareChatCompletionPath, "{account_id}", c.config.cloudflareAccountId, 1))
 	_ = util.OverwriteRequestHost(cloudflareDomain)
-	_ = util.OverwriteRequestAuthorization("Bearer " + c.config.GetRandomToken())
+	_ = util.OverwriteRequestAuthorization("Bearer " + getApiTokenInUse(ctx))
 
 	_ = proxywasm.RemoveHttpRequestHeader("Accept-Encoding")
 	_ = proxywasm.RemoveHttpRequestHeader("Content-Length")

plugins/wasm-go/extensions/ai-proxy/provider/cohere.go

@@ -59,7 +59,7 @@ func (m *cohereProvider) OnRequestHeaders(ctx wrapper.HttpContext, apiName ApiNa
 	}
 	_ = util.OverwriteRequestHost(cohereDomain)
 	_ = util.OverwriteRequestPath(chatCompletionPath)
-	_ = util.OverwriteRequestAuthorization("Bearer " + m.config.GetRandomToken())
+	_ = util.OverwriteRequestAuthorization("Bearer " + getApiTokenInUse(ctx))
 	_ = proxywasm.RemoveHttpRequestHeader("Content-Length")
 	return types.ActionContinue, nil
 }

plugins/wasm-go/extensions/ai-proxy/provider/deepl.go

@@ -79,7 +79,7 @@ func (d *deeplProvider) OnRequestHeaders(ctx wrapper.HttpContext, apiName ApiNam
 		return types.ActionContinue, errUnsupportedApiName
 	}
 	_ = util.OverwriteRequestPath(deeplChatCompletionPath)
-	_ = util.OverwriteRequestAuthorization("DeepL-Auth-Key " + d.config.GetRandomToken())
+	_ = util.OverwriteRequestAuthorization("DeepL-Auth-Key " + getApiTokenInUse(ctx))
 	_ = proxywasm.RemoveHttpRequestHeader("Content-Length")
 	_ = proxywasm.RemoveHttpRequestHeader("Accept-Encoding")
 	return types.HeaderStopIteration, nil

plugins/wasm-go/extensions/ai-proxy/provider/deepseek.go

@@ -49,7 +49,7 @@ func (m *deepseekProvider) OnRequestHeaders(ctx wrapper.HttpContext, apiName Api
 	}
 	_ = util.OverwriteRequestPath(deepseekChatCompletionPath)
 	_ = util.OverwriteRequestHost(deepseekDomain)
-	_ = util.OverwriteRequestAuthorization("Bearer " + m.config.GetRandomToken())
+	_ = util.OverwriteRequestAuthorization("Bearer " + getApiTokenInUse(ctx))
 	_ = proxywasm.RemoveHttpRequestHeader("Content-Length")
 	return types.ActionContinue, nil
 }

plugins/wasm-go/extensions/ai-proxy/provider/doubao.go

@@ -42,7 +42,7 @@ func (m *doubaoProvider) GetProviderType() string {
 
 func (m *doubaoProvider) OnRequestHeaders(ctx wrapper.HttpContext, apiName ApiName, log wrapper.Log) (types.Action, error) {
 	_ = util.OverwriteRequestHost(doubaoDomain)
-	_ = util.OverwriteRequestAuthorization("Bearer " + m.config.GetRandomToken())
+	_ = util.OverwriteRequestAuthorization("Bearer " + getApiTokenInUse(ctx))
 	_ = proxywasm.RemoveHttpRequestHeader("Content-Length")
 	if m.config.protocol == protocolOriginal {
 		ctx.DontReadRequestBody()