feat(ai-proxy-multi): let configured HTTP statuses trigger a fallback (#13852)

Author: nic-6443

apisix/plugins/ai-proxy-multi.lua

@@ -934,7 +934,8 @@ local function retry_on_error(ctx, conf, code, body)
     ctx.server_picker.after_balance(ctx, true)
     if (code == 429 and fallback_strategy_has(conf.fallback_strategy, "http_429")) or
        (code >= 500 and code < 600 and
-       fallback_strategy_has(conf.fallback_strategy, "http_5xx")) then
+       fallback_strategy_has(conf.fallback_strategy, "http_5xx")) or
+       base.is_fallback_http_status(conf, code) then
         -- Slow-failure guard: only retry when the failed attempt finished within
         -- retry_on_failure_within_ms. A slow failure (e.g. a 5xx returned after
         -- minutes) is given back to the client directly, so fallback never doubles

apisix/plugins/ai-proxy/base.lua

@@ -21,6 +21,7 @@ local core = require("apisix.core")
 local require = require
 local pcall   = pcall
 local pairs   = pairs
+local ipairs  = ipairs
 local type    = type
 local table   = table
 local math_floor = math.floor
@@ -48,6 +49,26 @@ local function count_request_tools(body)
 end
 
 
+-- Statuses the user opted into retrying through fallback_http_statuses.
+-- Both the error path below and ai-proxy-multi's retry decision consult this,
+-- and they have to agree: a status that is not diverted to the error path has
+-- its response streamed to the client and never reaches the retry callback.
+function _M.is_fallback_http_status(conf, status)
+    local statuses = conf.fallback_http_statuses
+    if not statuses then
+        return false
+    end
+
+    for _, s in ipairs(statuses) do
+        if s == status then
+            return true
+        end
+    end
+
+    return false
+end
+
+
 local function resolve_cap(cap_entry, key, conf, ctx)
     local val = cap_entry and cap_entry[key]
     if type(val) == "function" then
@@ -387,7 +408,8 @@ function _M.before_proxy(conf, ctx, on_error)
             -- Upstream responded — mark source before any early returns
             core.response.set_response_source(ctx, "upstream")
 
-            if res.status == 429 or (res.status >= 500 and res.status < 600) then
+            if res.status == 429 or (res.status >= 500 and res.status < 600)
+               or _M.is_fallback_http_status(conf, res.status) then
                 -- Read the upstream error body before closing so the provider's
                 -- error details survive: logged on fallback (see retry_on_error)
                 -- and returned to the client when no retry happens.

apisix/plugins/ai-proxy/schema.lua

@@ -453,6 +453,27 @@ _M.ai_proxy_multi_schema = {
               }
             }
         },
+        fallback_http_statuses = {
+            type = "array",
+            minItems = 1,
+            uniqueItems = true,
+            items = {
+                type = "integer",
+                minimum = 400,
+                maximum = 599,
+            },
+            description = "Additional upstream HTTP status codes that make the "
+                .. "request fall back to another instance, on top of the "
+                .. "http_429 / http_5xx entries of fallback_strategy. Use it "
+                .. "for statuses that mean the instance's credential is "
+                .. "unusable rather than the request being wrong, e.g. [401, "
+                .. "402] when an API key is expired or out of quota. It is "
+                .. "opt-in per status because most 4xx responses are caused by "
+                .. "the request itself and retrying them on another instance "
+                .. "would only burn quota. max_retries and "
+                .. "retry_on_failure_within_ms bound these retries the same "
+                .. "way they bound the fallback_strategy ones.",
+        },
         max_retries = {
             type = "integer",
             minimum = 0,

docs/en/latest/plugins/ai-proxy-multi.md

@@ -68,8 +68,9 @@ When an instance's `provider` is set to `bedrock`, the Plugin expects requests i
 | Name                               | Type            | Required | Default                           | Valid values | Description |
 |------------------------------------|----------------|----------|-----------------------------------|--------------|-------------|
 | fallback_strategy                  | string or array         | False    |  | string: "instance_health_and_rate_limiting", "http_429", "http_5xx"<br />array: ["rate_limiting", "http_429", "http_5xx"] | Fallback strategy. When set, the Plugin will check whether the specified instance's token has been exhausted when a request is forwarded. If so, forward the request to the next instance regardless of the instance priority. When not set, the Plugin will not forward the request to low priority instances when token of the high priority instance is exhausted. |
-| max_retries                        | integer        | False    |                                   | greater or equal to 0 | Maximum number of fallback retries after the initial request fails. Bounds how many additional instances a single request tries, so it does not exhaust every configured instance. Only takes effect together with `fallback_strategy`. When unset, the Plugin retries until an instance succeeds or all are tried. |
-| retry_on_failure_within_ms         | integer        | False    |                                   | greater or equal to 1 | Only fall back to another instance when the upstream fails within this many milliseconds. Fast failures (such as connection errors or quick `429`/`5xx`) are retried, while a slow failure that takes longer than this is returned to the client directly to avoid doubling the wait time. Only takes effect together with `fallback_strategy`. When unset, the Plugin retries regardless of how long the failed attempt took. |
+| fallback_http_statuses             | array[integer] | False    |                                   | between 400 and 599 | Additional upstream HTTP status codes that make the request fall back to another instance, on top of the `http_429` and `http_5xx` entries of `fallback_strategy`. Use it for statuses that mean the instance's credential is unusable rather than the request being wrong, such as `[401, 402]` when an API key is expired or out of quota. It is opt-in per status because most `4xx` responses are caused by the request itself and retrying them on another instance would only burn quota. `max_retries` and `retry_on_failure_within_ms` bound these retries the same way they bound the `fallback_strategy` ones. |
+| max_retries                        | integer        | False    |                                   | greater or equal to 0 | Maximum number of fallback retries after the initial request fails. Bounds how many additional instances a single request tries, so it does not exhaust every configured instance. Only takes effect together with `fallback_strategy` or `fallback_http_statuses`. When unset, the Plugin retries until an instance succeeds or all are tried. |
+| retry_on_failure_within_ms         | integer        | False    |                                   | greater or equal to 1 | Only fall back to another instance when the upstream fails within this many milliseconds. Fast failures (such as connection errors or quick `429`/`5xx`) are retried, while a slow failure that takes longer than this is returned to the client directly to avoid doubling the wait time. Only takes effect together with `fallback_strategy` or `fallback_http_statuses`. When unset, the Plugin retries regardless of how long the failed attempt took. |
 | balancer                           | object         | False    |                                   |              | Load balancing configurations. |
 | balancer.algorithm                 | string         | False    | roundrobin                     | [roundrobin, chash, semantic] | Load balancing algorithm. When set to `roundrobin`, weighted round robin algorithm is used. When set to `chash`, consistent hashing algorithm is used. When set to `semantic`, the Plugin picks an instance by the semantic similarity between the request prompt and each instance's `examples`, and its options are configured under `semantic_opts`. Note that `semantic` does not participate in health checks or `fallback_strategy` / retry — an upstream failure on the chosen instance is returned to the client; it only falls back (to the `semantic_opts.fallback` instance, else the first instance) when no instance clears its threshold or embedding fails. |
 | balancer.hash_on                   | string         | False    |                                   | [vars, headers, cookie, consumer, vars_combinations] | Used when `type` is `chash`. Support hashing on [NGINX variables](https://nginx.org/en/docs/varindex.html), headers, cookie, consumer, or a combination of [NGINX variables](https://nginx.org/en/docs/varindex.html). |
@@ -173,8 +174,8 @@ The setting covers `ai-proxy`, `ai-proxy-multi`, and `ai-request-rewrite`, which
 
 When the selected LLM upstream returns a `429` or `5xx` status, `ai-proxy-multi` reads the upstream error body before deciding whether to fall back:
 
-- If the request is retried on another instance (per `fallback_strategy`, `max_retries`, and `retry_on_failure_within_ms`), the failed instance's error body is recorded in the error log for diagnostics, since a later attempt's response is sent to the client instead.
-- If the request is not retried (no matching `fallback_strategy`, retries exhausted, or the failure took longer than `retry_on_failure_within_ms`), the upstream status code and error body are returned to the client, preserving the upstream `Content-Type`.
+- If the request is retried on another instance (per `fallback_strategy`, `fallback_http_statuses`, `max_retries`, and `retry_on_failure_within_ms`), the failed instance's error body is recorded in the error log for diagnostics, since a later attempt's response is sent to the client instead.
+- If the request is not retried (no matching `fallback_strategy` or `fallback_http_statuses`, retries exhausted, or the failure took longer than `retry_on_failure_within_ms`), the upstream status code and error body are returned to the client, preserving the upstream `Content-Type`.
 
 ## Examples
 

docs/zh/latest/plugins/ai-proxy-multi.md

@@ -68,8 +68,9 @@ import TabItem from '@theme/TabItem';
 | 名称                               | 类型            | 必选项 | 默认值                           | 有效值 | 描述 |
 |------------------------------------|----------------|----------|-----------------------------------|--------------|-------------|
 | fallback_strategy                  | string 或 array         | 否    |  | string: "instance_health_and_rate_limiting", "http_429", "http_5xx"<br />array: ["rate_limiting", "http_429", "http_5xx"] | 故障转移策略。设置后，插件将在转发请求时检查指定实例的令牌是否已耗尽。如果是，则无论实例优先级如何，都将请求转发到下一个实例。未设置时，当高优先级实例的令牌耗尽时，插件不会将请求转发到低优先级实例。 |
-| max_retries                        | integer        | 否    |                                   | 大于或等于 0 | 初始请求失败后允许的最大故障转移重试次数。用于限制单个请求最多尝试多少个额外实例，避免穷举所有已配置的实例。仅在配置 `fallback_strategy` 时生效。未设置时，插件会持续重试直到某个实例成功或所有实例都已尝试。 |
-| retry_on_failure_within_ms         | integer        | 否    |                                   | 大于或等于 1 | 仅当上游在指定毫秒数内失败时才故障转移到其他实例。快速失败（如连接错误、快速返回的 `429`/`5xx`）会触发重试，而耗时超过该值的慢失败会直接将错误返回给客户端，避免客户端等待时间翻倍。仅在配置 `fallback_strategy` 时生效。未设置时，插件无论失败请求耗时多久都会重试。 |
+| fallback_http_statuses             | array[integer] | 否    |                                   | 400 到 599 之间 | 除 `fallback_strategy` 中的 `http_429` 和 `http_5xx` 之外，额外触发故障转移到其他实例的上游 HTTP 状态码。适用于表示该实例凭证不可用（而非请求本身有问题）的状态码，例如 API Key 过期或额度耗尽时的 `[401, 402]`。该能力按状态码显式启用，因为大多数 `4xx` 响应由请求本身导致，在其他实例上重试只会白白消耗额度。`max_retries` 和 `retry_on_failure_within_ms` 对这些重试的约束方式与对 `fallback_strategy` 的重试一致。 |
+| max_retries                        | integer        | 否    |                                   | 大于或等于 0 | 初始请求失败后允许的最大故障转移重试次数。用于限制单个请求最多尝试多少个额外实例，避免穷举所有已配置的实例。仅在配置 `fallback_strategy` 或 `fallback_http_statuses` 时生效。未设置时，插件会持续重试直到某个实例成功或所有实例都已尝试。 |
+| retry_on_failure_within_ms         | integer        | 否    |                                   | 大于或等于 1 | 仅当上游在指定毫秒数内失败时才故障转移到其他实例。快速失败（如连接错误、快速返回的 `429`/`5xx`）会触发重试，而耗时超过该值的慢失败会直接将错误返回给客户端，避免客户端等待时间翻倍。仅在配置 `fallback_strategy` 或 `fallback_http_statuses` 时生效。未设置时，插件无论失败请求耗时多久都会重试。 |
 | balancer                           | object         | 否    |                                   |              | 负载均衡配置。 |
 | balancer.algorithm                 | string         | 否    | roundrobin                     | [roundrobin, chash, semantic] | 负载均衡算法。设置为 `roundrobin` 时，使用加权轮询算法。设置为 `chash` 时，使用一致性哈希算法。设置为 `semantic` 时，插件根据请求提示词与各实例 `examples` 之间的语义相似度选择实例，其相关选项配置在 `semantic_opts` 下。注意：`semantic` 不参与健康检查，也不参与 `fallback_strategy` / 重试——被选中实例的上游失败会直接返回给客户端；只有当没有实例达到其阈值或嵌入请求失败时，才会回退（回退到 `semantic_opts.fallback` 实例，若未配置则回退到第一个实例）。 |
 | balancer.hash_on                   | string         | 否    |                                   | [vars, headers, cookie, consumer, vars_combinations] | 当 `type` 为 `chash` 时使用。支持基于 [NGINX 变量](https://nginx.org/en/docs/varindex.html)、标头、cookie、消费者或 [NGINX 变量](https://nginx.org/en/docs/varindex.html)组合进行哈希。 |