Skip to content

feat(ai-proxy): harden configurable retry and fallback policy - #7668

Merged
hengyuss merged 4 commits into
apache:masterfrom
hengyuss:feat/ai-proxy-retry-policy
Oct 20, 2026
Merged

hengyuss merged 4 commits into
apache:masterfrom
hengyuss:feat/ai-proxy-retry-policy

Conversation

@hengyuss

Copy link
Copy Markdown
Contributor

Closes #7667.
Related to #7351.
Follow-up to #7532.

What this PR does

  • centralizes upstream HTTP status and classified transport failure decisions in AiProxyRetryPolicy;
  • retries only explicitly eligible failures and forwards non-retryable upstream responses unchanged;
  • retries the primary Provider before attempting fallback and attempts fallback at most once;
  • prevents retry or fallback after the first response body chunk has been emitted;
  • drains retryable upstream error bodies before another attempt so connections can be released;
  • rebuilds fallback requests using the fallback Provider configuration;
  • adds separate retry limits for streaming and non-streaming requests;
  • adds configurable exponential backoff for non-streaming retries while keeping pre-emission streaming retries immediate;
  • binds retry settings under shenyu.plugins.ai.proxy.retry;
  • replaces the remaining legacy transport test fixture with classified AiTransportException.

Retry configuration

shenyu:
  plugins:
    ai:
      proxy:
        retry:
          max-retries: 3
          streaming-max-retries: 1
          initial-backoff: 1s

Behavioral boundaries

  • HTTP 429 and 5xx responses are retry/fallback eligible.
  • Other 4xx responses are forwarded without retry or fallback.
  • Transport failures use an explicit retry/fallback eligibility matrix.
  • Once downstream emission starts, the original failure is propagated without transparent replay.
  • Terminal transport exception mapping to gateway 502/503/504 responses is intentionally out of scope.
  • Connection pool, TLS, HTTP proxy, lifecycle, and observability configuration remain under [Task][AI Proxy] Harden the reactive AI Transport for production #7384.

Tests

  • ./mvnw -pl shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy -Djacoco.skip=true test — 253 tests passed.
  • ./mvnw -pl shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy -am -Dtest=AiProxyPluginConfigurationTest -Dsurefire.failIfNoSpecifiedTests=false -Djacoco.skip=true test — 11 tests passed.
  • Checkstyle passed for all modules participating in both builds.

Make sure that:

  • You have read the contribution guidelines.
  • You submit test cases (unit or integration tests) that back your changes.
  • Your local test passed ./mvnw clean install -Dmaven.javadoc.skip=true.

@Aias00 Aias00 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary

This is a clean, well-structured refactor that centralizes the AI proxy retry/fallback decision in AiProxyRetryPolicy and makes the retry budget configurable via AiProxyRetryProperties. The behavioural boundaries from the PR description are correctly implemented and thoroughly tested.

What is right

  • The split between decideStatus (HTTP status) and decideTransport (transport failure kind) is clear, and the eligibility matrix is internally consistent: 429/5xx retry then fall back; other 4xx are forwarded without retry or fallback even when a fallback is configured. This also fixes the old behaviour where any 4xx with a fallback present would incorrectly trigger a fallback (the previous status >= 400 && (retryableStatus && retries>0 || fallback.isPresent()) condition).
  • attempt() now drains the upstream error body (response.getBody().then()) before emitting UpstreamStatusFailure, which releases the connection before the next attempt.
  • The streaming "no retry/fallback after first chunk emitted" boundary (emitted flag) is preserved and still covered by neverRetriesOrFallsBackAfterBodyEmission.
  • consumeResponse is correctly extracted and reused on both the success and forward paths.
  • Tests are strong: AiProxyRetryPolicyTest covers every matrix cell with CsvSource; AiProxyRetryPropertiesTest covers bounds and exponential backoff; AiProxyEngineOrchestrationTest verifies call counts for the retry whitelist, fallback whitelist, response-too-large, and backoff timing (1s -> 2s -> 4s). I cross-checked the expected call counts against the policy and they all match.
  • Spring binding for shenyu.plugins.ai.proxy.retry.* is covered, including rejection of out-of-range values.

Non-blocking suggestions

  1. decideStatus is invoked twice for a status failure — once in the flatMap (only to detect FORWARD_RESPONSE) and again in onErrorResume. Functionally correct, but the happy-path branch could be simplified to a direct retryableStatus(status) ? emitError : consumeResponse to avoid the redundant decision object.
  2. AiProxyRetryProperties.MAX_EXPONENTIAL_RETRIES = 30 combined with 1L << retryIndex means a configured max-retries of 30 produces a ~6.2-day backoff on the last attempt. Consider capping backoff growth (e.g. Math.min(delay, someMax)) so an extreme config cannot produce absurd waits.

CI

All build (shenyu-integrated-test-*) jobs and pr_build pass. The only failures are e2e / e2e-storage (compose-based infrastructure), unrelated to this change.

Approving.

@wy471x wy471x left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed head 10b6b106ac (base 1d599d90e) locally on JDK 21.

Verification:

  • ./mvnw -pl shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-proxy test -Djacoco.skip=true → 253/253 green, checkstyle green (matches the PR description). The SNAPSHOT dependencies had to be installed from source first (-pl … -am install -DskipTests) because the cached shenyu-plugin-ai-common in the local repository predated AiCommonConfig#getProtocol.
  • ./mvnw -pl shenyu-spring-boot-starter/shenyu-spring-boot-starter-plugin/shenyu-spring-boot-starter-plugin-ai-proxy -am test -Dtest=AiProxyPluginConfigurationTest -Dsurefire.failIfNoSpecifiedTests=false → 11/11 green, checkstyle clean across the whole reactor.
  • Extra probe (throwaway test, deleted afterwards): a request whose payload declares "stream": true while the handle leaves stream unset still gets the streaming budget and the immediate delay. With AiProxyRetryProperties(3, 1, Duration.ZERO) against an always-500 primary there are exactly 2 primary calls, so retryProperties.maxRetries(primary.stream()) and retryDelay(prepared.stream(), …) both read the effective flag produced by prepare() — the two sources cannot disagree.
  • Equivalence of the new FORWARD_RESPONSE branch, which hands the raw transport response to the consumer without provider.adaptResponse / protocol.encodeResponse: behaviour-preserving today, because ConnectProtocolAdapter.convertResponse returns the response unchanged for OPENAI_CHAT clients and for any status outside 2xx, OpenAiChat.encodeResponse returns non-2xx responses unchanged, and AiProxyProvider.adaptResponse defaults to identity. The new parameterized 4xx case (status preserved, X-Upstream: kept) pins the observable part of this.
  • CI on this head: build and every shenyu-integrated-test-* leg pass. The e2e aggregator and e2e-storage failures reproduce on unrelated open PRs (#7663, #7666), so they are environmental rather than caused here.

The shape of the change is right: policy and configuration are separated from orchestration, the decision matrix is table-driven in its own class, and the ordering guarantees ("primary retries before fallback", "fallback at most once", "nothing after emission") are asserted end-to-end instead of only at the unit level.

Behaviour change worth surfacing in the release notes: a non-429 4xx is no longer masked by the configured fallback. Previously any status ≥ 400 with a fallback present switched provider; now 400/401/403/404/408/422 are forwarded to the client unchanged (hence the AiProxyPluginTest expectation moving 400 → 500). That is exactly what #7667 asks for, but a deployment that relied on "fallback swallows every upstream error" will start seeing those 4xx responses. The description bullets already say it; one sentence framed as a migration note would save maintainers the diff archaeology.

Notes (non-blocking):

  • AiProxyException.Kind.TRANSPORT is now unreferenced anywhere in the repository — the engine classifies only AiTransportException kinds. If an out-of-tree implementation still throws the legacy kind, it now loses retry and fallback silently. Either delete the constant or state in AiProxyRetryPolicy's javadoc that transport classification lives solely in AiTransportException.
  • AiProxyEngine(…, Duration retryBackoff) now expands to new AiProxyRetryProperties(3, 1, retryBackoff), so that constructor silently ignores shenyu.plugins.ai.proxy.retry.*. Only test helpers call it today; @Deprecated plus a javadoc pointer (or dropping it and migrating those helpers) prevents two retry sources from drifting apart.
  • The raw-forward branch depends on the invariant "non-2xx responses are never adapted". A one-line comment there would keep a future protocol that legitimately adapts error bodies from being bypassed by accident.
  • MAX_EXPONENTIAL_RETRIES serves both as the configuration range bound and as the shift guard in retryDelay. That is consistent today (a retry only happens while retriesRemaining > 0, so retryIndex ≤ maxRetries - 1), but the two concepts sharing one constant is exactly what breaks if a larger max-retries is ever allowed; two named constants (or a comment) would make the invariant self-evident.

Coverage notes: the transport matrix asserts both the per-kind primary call counts (2 for retryable kinds, 1 otherwise) and the fallback whitelist, and it pins that a fallback attempt never restarts the primary; streaming-max-retries=0 and the default of 1 are both exercised.

@Aias00 Aias00 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review — harden configurable retry and fallback policy

Approve. This is a well-structured refactor with an unusually thorough test suite, and I verified the risky parts against the real dependency sources rather than the description. Two deliberate behavior changes are worth calling out explicitly (both documented, both defensible), and I have four non-blocking notes.

✅ Verified

AiProxyRetryPolicy is safe as a shared singleton. It has no fields, so private static final AiProxyRetryPolicy RETRY_POLICY in AiProxyEngine carries no mutable state across concurrent requests.

The legacy Duration constructor is exactly equivalent to master. new AiProxyRetryProperties(3, 1, retryBackoff) reproduces master's primary.stream() ? 1 : 3 split precisely, and the backoff formula retryIndex = maxRetries - retriesRemaining (shift 0,1,2) is identical to master's 1L << (3 - retries) at the default. All five other test files that use either the 3-arg or the 4-arg-Duration overload (AiProxyHttpCompatibilityTest:221, AiProxyProviderHttpTest:131, AiProxyPluginTest:286, and others) compile untouched — which is the right reason to keep the overload rather than force-migrating callers.

Forwarding the raw response is provably equivalent to master's adapt+encode. This was my main concern on reading the diff, and it turns out to be a clean simplification. FORWARD_RESPONSE is only reachable with status >= 400 (both call sites are inside if (status >= 400) / from UpstreamStatusFailure), and on that path:

  • OpenAiChat.encodeResponse (:240) returns the response unchanged for any status outside 200-299;
  • ConnectProtocolAdapter.convertResponse (:109) does the same;
  • AiProxyProvider.adaptResponse's default is identity and nothing overrides it.

So consumeResponse(response, ...) and master's consumeResponse(adapt(encode(response)), ...) produce the same result. Good — the pipeline removal is behavior-preserving, not a shortcut.

Dropping the AiProxyException(Kind.TRANSPORT) branch is safe. I grepped the whole main source tree: every remaining AiProxyException throw uses REQUEST (OpenAi.java:146, ConnectProtocolAdapter:97, AiProviderCapabilities:131), PROTOCOL (OpenAiChat:247, OpenAiChatSse:95/142, ConnectProtocolAdapter:119/137/155/165), or PROVIDER (AiProxyEngine:236). WebClientAiProxyTransport throws only AiTransportException and covers all ten Kind values. No production path can still raise the legacy type, so the narrowing is complete rather than hopeful.

Terminal mapping is preserved. AiProxyPlugin.handleFailure:143-144 maps AiTransportException → BAD_GATEWAY, which matches master's AiProxyException(TRANSPORT) → BAD_GATEWAY. Your "terminal 502/503/504 mapping is out of scope" note is accurate — it stays a flat 502.

Retry/fallback accounting checks out by hand. testExecuteWithAdminFallback (primary 500, non-streaming, legacy ctor → maxRetries 3): RETRY×3 then FALLBACK = 4 primary + 1 fallback = 5 requests, with requests.get(4) carrying admin.example and Bearer admin-key. fallbackFailureDoesNotRestartThePrimary (streaming, streamingMaxRetries 1): 1 primary + 1 retry + 1 fallback = 3, then propagate. Both match the decision matrix exactly.

The fallback attempt cannot cascade. Passing maxRetries=0, retriesRemaining=0 into the fallback attempt(...) means a 503 from the fallback yields decideStatus(503, 0, false) → FORWARD_RESPONSE, and a transport failure yields PROPAGATE_ERROR. "At most once" is enforced structurally, not by convention.

retriesRemaining - 1 cannot underflow (RETRY requires retriesRemaining > 0), and retryIndex is bounded by maxRetries - 1 ≤ 29, so retryDelay's own range guard is unreachable from the engine — correct, and the guard is still worth having for direct callers.

Config binding mirrors the existing pattern exactly. aiProxyRetryProperties() is a line-for-line analogue of aiProxyTransportProperties() (@Bean + @ConfigurationProperties, same prefix shape), with setter validation and a hasFailed() context test. Documenting the keys in shenyu-bootstrap/application.yml next to the existing transport block is the right place.

Tests are the strongest part. AiProxyRetryPolicyTest pins the entire decision matrix including the 600 → FORWARD_RESPONSE edge; the retry and fallback whitelists are separately parameterized; and nonStreamingRetriesUseExponentialBackoff uses StepVerifier.withVirtualTime to assert the exact 1s/2s/4s schedule deterministically rather than sleeping. clientErrorsAreForwardedWithoutRetryOrFallback was widened from one status to six.


⚠️ Two behavior changes to flag for release notes

Both are stated in your PR body, so this is not a objection — but operators upgrading will feel them, and neither is covered by an e2e or integration test.

1. Non-retryable 4xx no longer triggers fallback. On master, 401 + fallback configured → fallback was used (the old testExecuteWithAdminFallback asserted 2 requests). Now decideStatus(401, …) returns FORWARD_RESPONSE and the client error is passed through with 1 request. Your reading is better than master's — a 400 from the primary is a client error, not a provider outage — but it is a semantic change to the fallback feature.

2. Transport failures are now whitelisted. Master retried and then fell back on any transport failure. Now DNS_FAILURE, TLS_FAILURE, STREAM_DURATION_TIMEOUT and UNKNOWN do neither. Sensible (a DNS/TLS failure is usually configuration, and retrying the same provider won't help), and the RESPONSE_TOO_LARGE split — no same-provider retry, fallback allowed — is the nicest decision in the matrix. Just worth stating plainly.


📝 Non-blocking notes

1. The decision is computed twice. flatMap calls decideStatus to choose drain+error, then onErrorResume calls it again to choose the action. It's a pure function of identical method-local inputs, so the two always agree — but a reader has to prove that rather than see it. Either a one-line comment or carrying the decision inside UpstreamStatusFailure would make it locally obvious. Relatedly, FORWARD_RESPONSE and PROPAGATE_ERROR both fall through to the same return Mono.error(error) at the end of onErrorResume, which is correct but reads as if the distinction matters there.

2. max-retries: 30 × exponential backoff is a footgun. requireRetryCount allows up to 30, and the cumulative wait is initialBackoff × (2^30 − 1) — roughly 34 years at the 1s default. Someone raising max-retries without lowering initial-backoff will produce an effectively hung request rather than a fast failure. Worth documenting the interaction, or capping the cumulative delay independently of the per-attempt count.

3. The legacy constructor now fails fast on a null retryBackoff. setInitialBackoff rejects null, where master would have NPE'd at the first retry. Strictly better; noting only so it isn't a surprise if any out-of-tree caller relied on lazy failure.

4. AiProxyRetryProperties fields are non-volatile, written once by the Spring binder at startup. This is the same pattern as AiProxyTransportProperties, so it's consistent — but it does mean retry settings are not hot-reloadable, which is worth knowing given the plugin already has a runtime-config-refresh story.


CI

Only one job genuinely failed. e2e-storage (shenyu-postgres-compose) failed on DividePluginTest.testDivide:105 with a 502 Bad Gateway body — DividePluginCases uses httpbin.org as a live upstream, so this is the external-site flake, not a regression. The other three storage jobs (h2, mysql, opengauss) show as failed in gh pr checks but their conclusion is cancelled (fail-fast siblings); the h2 job's own log ends in BUILD SUCCESS with Tests run: 8, Failures: 0, and e2e is the 3-second aggregation job. On master, run 38836735684 failed the identical postgres job and the three master runs since (8f3ee32e, 66d7d954, 3af7bfc6) are all green. This PR touches only the AI proxy plugin, its starter, and shenyu-bootstrap/application.yml, none of which are on that path.

Summary

Approve. The retry/fallback eligibility matrix is a real improvement over master's "retry everything, then fall back on everything", the extraction into a stateless policy plus a validated properties bean is the right shape, and the tests pin the matrix and the backoff schedule deterministically. Please surface the two behavior changes above in the release notes.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Task][AI Proxy] Harden configurable retry and fallback policy in AiProxyEngine

3 participants