diff --git a/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/interceptor/modelfallback/ModelFallbackInterceptor.java b/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/interceptor/modelfallback/ModelFallbackInterceptor.java index 8f0fb0c7e..fe0f1f723 100644 --- a/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/interceptor/modelfallback/ModelFallbackInterceptor.java +++ b/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/interceptor/modelfallback/ModelFallbackInterceptor.java @@ -21,21 +21,38 @@ import io.github.agentic.ai.agent.interceptor.ModelResponse; import io.github.agentic.ai.api.chat.ChatClient; +import io.github.agentic.ai.api.chat.ChatEvent; +import io.github.agentic.ai.runtime.execution.ExecutionGuard; import java.util.ArrayList; import java.util.List; import java.util.Objects; import java.util.concurrent.CancellationException; +import java.util.concurrent.atomic.AtomicBoolean; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Flux; +import reactor.core.scheduler.Schedulers; + /** * Automatic fallback to alternative models on errors. * * Retries failed model calls with alternative models in sequence until success or all - * models exhausted. Streaming responses are passed through unchanged. Errors emitted by - * the publisher after subscription are not handled by this interceptor. + * models exhausted. A failure of a streaming response switches to the next candidate + * only when it arrives before any event has been emitted downstream; once a text delta, + * tool call, usage or response event has been observed, the failure terminates the + * stream, so output of two models is never concatenated. A cancellation — however + * deeply it is wrapped in the cause chain — keeps its unwrapped identity on either side + * of that boundary, and it, an interruption, or a run whose {@link ExecutionGuard} is + * no longer active never triggers another candidate. Each subscription of a streaming + * response owns its switch sequence. + * + * When combined with {@code ModelRetryInterceptor}, register this interceptor first + * (outermost), so each model is retried to exhaustion before the sequence fails over to + * the next candidate. * * Example: ModelFallbackInterceptor interceptor = ModelFallbackInterceptor.builder() * .addFallbackClient(primaryAlternative) .addFallbackClient(secondaryAlternative) @@ -45,6 +62,9 @@ public class ModelFallbackInterceptor extends ModelInterceptor { private static final Logger log = LoggerFactory.getLogger(ModelFallbackInterceptor.class); + /** Bounds the cause-chain walks; a cyclic chain must not spin the thread forever. */ + private static final int MAX_CAUSE_DEPTH = 32; + private final List fallbackClients; private ModelFallbackInterceptor(Builder builder) { @@ -59,9 +79,13 @@ public static Builder builder() { public ModelResponse interceptModel(ModelRequest request, ModelCallHandler handler) { Exception lastException = null; + // A cancelled or expired run is terminal: no provider call at all, the primary + // included, mirroring the per-attempt guard check of the retry interceptor. + assertRunActive(request); + // Try primary model first try { - return handler.call(request); + return withSwitchBoundary(request, handler, handler.call(request), 0); } catch (Exception e) { propagateCancellation(e); @@ -71,11 +95,15 @@ public ModelResponse interceptModel(ModelRequest request, ModelCallHandler handl // Try fallback models in sequence for (int i = 0; i < fallbackClients.size(); i++) { + // A cancelled or expired run is terminal: no further provider call. + assertRunActive(request); try { log.info("Trying fallback client {} of {}", i + 1, fallbackClients.size()); // Reuse the request pipeline so system messages, tools and usage are // retained. - return handler.call(ModelRequest.builder(request).chatClient(fallbackClients.get(i)).build()); + ModelResponse fallback = handler + .call(ModelRequest.builder(request).chatClient(fallbackClients.get(i)).build()); + return withSwitchBoundary(request, handler, fallback, i + 1); } catch (Exception e) { propagateCancellation(e); @@ -85,18 +113,197 @@ public ModelResponse interceptModel(ModelRequest request, ModelCallHandler handl } // All models failed - throw new RuntimeException("All clients failed after " + (fallbackClients.size() + 1) + " attempts", + throw allClientsFailed(lastException); + } + + /** + * Attaches the switch boundary to a candidate's streaming response. A + * non-streaming response is returned unchanged. Every subscription of the returned + * publisher owns its own switch state, so a repeated subscription replays the + * captured candidate stream instead of resuming the fallback sequence of an earlier + * one. + */ + private ModelResponse withSwitchBoundary(ModelRequest request, ModelCallHandler handler, ModelResponse response, + int nextClientIndex) { + if (!response.isStreaming()) { + return response; + } + return ModelResponse.stream(JdkFlowAdapter.publisherToFlowPublisher( + switchingFluxDeferred(request, handler, JdkFlowAdapter.flowPublisherToFlux(response.getStream()), + nextClientIndex))); + } + + /** + * The deferred form of the switch boundary: fresh switch state per subscription and + * a subscription-time interrupt check, mirroring the retry interceptor. + */ + private Flux switchingFluxDeferred(ModelRequest request, ModelCallHandler handler, + Flux candidateFlux, int nextClientIndex) { + return Flux.defer(() -> { + // Mirrors the retry interceptor: a subscribing thread that is already + // interrupted must not start (or resume) a provider sequence. + if (Thread.currentThread().isInterrupted()) { + return Flux.error(new CancellationException("Model call interrupted")); + } + return switchingFlux(request, handler, candidateFlux, nextClientIndex); + }); + } + + /** + * Wraps one candidate's stream with the switch boundary: a failure that arrives + * before any event has been emitted downstream switches to the next candidate, + * because no observable effect has escaped yet. Once an event — text delta, tool + * call, usage or response — has been emitted, the failure terminates the stream: a + * second model's output must never be appended to a partially delivered one. A + * cancellation keeps its unwrapped identity on both sides of that boundary. + */ + private Flux switchingFlux(ModelRequest request, ModelCallHandler handler, Flux candidateFlux, + int nextClientIndex) { + AtomicBoolean eventEmitted = new AtomicBoolean(false); + return candidateFlux.doOnNext(event -> eventEmitted.set(true)).onErrorResume(error -> { + if (error instanceof Error) { + return Flux.error(error); + } + // The error signal itself carries the cancellation truth on a reactive + // thread; the thread's interrupt flag belongs to the pool, not to this run, + // and is deliberately not consulted here. + CancellationException cancellation = findWrappedCancellation(error); + if (cancellation != null) { + log.info("Streaming candidate failed with a cancellation; terminating the fallback sequence: {}", + cancellation.getMessage()); + return Flux.error(cancellation); + } + if (eventEmitted.get()) { + return Flux.error(error); + } + // Dispatch off the signal thread: a non-streaming fallback candidate would + // otherwise park the thread that delivered the failure, which for a Netty + // event loop stalls every other stream multiplexed on it. + return Flux.defer(() -> nextClientFlux(request, handler, nextClientIndex, toException(error))) + .subscribeOn(Schedulers.boundedElastic()); + }); + } + + /** + * Streams the next fallback candidate, advancing until one succeeds or the sequence + * is exhausted. A cancellation, however wrapped, and a dead execution guard end the + * sequence instead of dispatching another provider call. + */ + private Flux nextClientFlux(ModelRequest request, ModelCallHandler handler, int clientIndex, + Exception lastException) { + if (clientIndex >= fallbackClients.size()) { + return Flux.error(allClientsFailed(lastException)); + } + try { + assertRunActive(request); + } + catch (Exception lostRun) { + return Flux.error(lostRun); + } + log.info("Trying fallback client {} of {}", clientIndex + 1, fallbackClients.size()); + try { + ModelResponse next = handler + .call(ModelRequest.builder(request).chatClient(fallbackClients.get(clientIndex)).build()); + if (next.isStreaming()) { + return switchingFluxDeferred(request, handler, JdkFlowAdapter.flowPublisherToFlux(next.getStream()), + clientIndex + 1); + } + return Flux.just(ChatEvent.response(next.getChatResponse())); + } + catch (Exception e) { + CancellationException cancellation = findWrappedCancellation(e); + if (cancellation != null) { + return Flux.error(cancellation); + } + log.warn("Fallback model {} failed: {}", clientIndex + 1, e.getMessage()); + return nextClientFlux(request, handler, clientIndex + 1, e); + } + } + + private static void assertRunActive(ModelRequest request) { + request.getExecutionGuard().ifPresent(ExecutionGuard::assertActive); + } + + private RuntimeException allClientsFailed(Exception lastException) { + StringBuilder sequence = new StringBuilder("primary"); + for (int i = 1; i <= fallbackClients.size(); i++) { + sequence.append(", fallback-").append(i); + } + return new RuntimeException( + "All clients failed after " + (fallbackClients.size() + 1) + " attempts (" + sequence + ")", lastException); } + private static Exception toException(Throwable error) { + return error instanceof Exception exception ? exception : new RuntimeException(error); + } + private static void propagateCancellation(Exception failure) { - if (failure instanceof CancellationException cancelled) { - throw cancelled; + CancellationException cancellation = cancellation(failure); + if (cancellation != null) { + throw cancellation; + } + } + + /** + * Finds the cancellation of a failed attempt, however deeply it is wrapped, without + * touching thread state. A wrapped cancellation is still a cancellation: switching + * to another model would dispatch a call the caller already aborted (issue #210). + * The walk is depth-bounded, because a cyclic cause chain (constructible through + * {@code Throwable.initCause} on two nodes) would otherwise spin the thread + * forever. + */ + private static CancellationException findWrappedCancellation(Throwable failure) { + Throwable cause = failure; + for (int depth = 0; cause != null && depth < MAX_CAUSE_DEPTH; depth++) { + if (cause instanceof CancellationException cancelled) { + return cancelled; + } + if (cause instanceof InterruptedException interrupted) { + CancellationException cancelled = new CancellationException("Model call interrupted"); + cancelled.initCause(interrupted); + return cancelled; + } + cause = cause.getCause() == cause ? null : cause.getCause(); + } + return null; + } + + /** + * The cancellation of a synchronous attempt. On the caller's own thread the + * interrupt flag carries the same meaning: an interruption in the chain restores + * the flag so upstream blocking work also terminates, and a bare interrupt flag + * converts to a cancellation. Reactive paths use + * {@link #findWrappedCancellation(Throwable)} instead, because a pooled signal + * thread's flag belongs to the pool, not to this run. + */ + private static CancellationException cancellation(Throwable failure) { + CancellationException wrapped = findWrappedCancellation(failure); + if (wrapped != null) { + if (chainsInterruption(failure)) { + Thread.currentThread().interrupt(); + } + return wrapped; + } + if (!Thread.currentThread().isInterrupted()) { + return null; + } + CancellationException cancelled = new CancellationException("Model call interrupted"); + if (failure != null) { + cancelled.initCause(failure); } - if (Thread.currentThread().isInterrupted() || failure instanceof InterruptedException) { - Thread.currentThread().interrupt(); - throw new CancellationException("Model call interrupted"); + return cancelled; + } + + private static boolean chainsInterruption(Throwable failure) { + Throwable cause = failure; + for (int depth = 0; cause != null && depth < MAX_CAUSE_DEPTH; depth++) { + if (cause instanceof InterruptedException) { + return true; + } + cause = cause.getCause() == cause ? null : cause.getCause(); } + return false; } @Override diff --git a/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackInterceptorStreamingTest.java b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackInterceptorStreamingTest.java index e81219f7b..2ea3e71c6 100644 --- a/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackInterceptorStreamingTest.java +++ b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackInterceptorStreamingTest.java @@ -51,19 +51,28 @@ void streamingPrimaryResponseDoesNotCallFallback() { ModelResponse result = interceptor(fallback).interceptModel(ModelRequest.builder().messages(List.of()).build(), request -> primary); - assertSame(primary, result); + // The response is wrapped by the switch boundary, but its events are delivered + // unchanged and the fallback client is never invoked. + assertEquals(true, result.isStreaming()); + Flux stream = JdkFlowAdapter.flowPublisherToFlux(result.getStream()); + assertEquals("primary", stream.blockLast(Duration.ofSeconds(2)).response().text()); assertEquals(0, fallback.calls); } @Test - void publisherErrorsRemainStreamingErrors() { + void publisherErrorsExhaustingAllCandidatesFailWithAggregate() { CountingModel fallback = new CountingModel("fallback"); RuntimeException failure = new RuntimeException("stream failed"); ModelResponse result = interceptor(fallback).interceptModel(ModelRequest.builder().messages(List.of()).build(), request -> ModelResponse.stream(JdkFlowAdapter.publisherToFlowPublisher(Flux.error(failure)))); Flux stream = JdkFlowAdapter.flowPublisherToFlux(result.getStream()); - assertSame(failure, assertThrows(RuntimeException.class, () -> stream.blockLast(Duration.ofSeconds(2)))); + // The handler fails for every candidate before any event is emitted, so the + // switch boundary walks the whole sequence and reports the aggregate. + RuntimeException aggregate = assertThrows(RuntimeException.class, + () -> stream.blockLast(Duration.ofSeconds(2))); + assertEquals("All clients failed after 2 attempts (primary, fallback-1)", aggregate.getMessage()); + assertSame(failure, aggregate.getCause()); assertEquals(0, fallback.calls); } diff --git a/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackNeutralClientTest.java b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackNeutralClientTest.java index 431dc6a38..d7ec09563 100644 --- a/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackNeutralClientTest.java +++ b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackNeutralClientTest.java @@ -165,7 +165,7 @@ void fallbackExhaustionPreservesOrderAndFinalCause() { assertThatThrownBy(() -> interceptor.interceptModel(ModelRequest.builder().build(), request -> { attempts.add(request.getChatClient()); throw request.getChatClient() == second ? last : new IllegalStateException("Unavailable"); - })).hasMessage("All clients failed after 3 attempts").hasCause(last); + })).hasMessage("All clients failed after 3 attempts (primary, fallback-1, fallback-2)").hasCause(last); assertThat(attempts).containsExactly(null, first, second); assertThatThrownBy(() -> ModelFallbackInterceptor.builder().build()) .isInstanceOf(IllegalArgumentException.class); diff --git a/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackSwitchBoundaryTest.java b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackSwitchBoundaryTest.java new file mode 100644 index 000000000..4d8244454 --- /dev/null +++ b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/ModelFallbackSwitchBoundaryTest.java @@ -0,0 +1,513 @@ +/* + * Copyright 2024-2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.orchestration.interceptors; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import io.github.agentic.ai.agent.interceptor.ModelCallHandler; +import io.github.agentic.ai.agent.interceptor.ModelRequest; +import io.github.agentic.ai.agent.interceptor.ModelResponse; +import io.github.agentic.ai.agent.orchestration.interceptor.modelfallback.ModelFallbackInterceptor; +import io.github.agentic.ai.agent.orchestration.interceptor.modelretry.ModelRetryInterceptor; +import io.github.agentic.ai.api.chat.ChatClient; +import io.github.agentic.ai.api.chat.ChatClientCapabilities; +import io.github.agentic.ai.api.chat.ChatEvent; +import io.github.agentic.ai.api.chat.ChatRequest; +import io.github.agentic.ai.api.chat.ChatResponse; +import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.ToolCall; +import io.github.agentic.ai.api.chat.Usage; +import io.github.agentic.ai.runtime.execution.ExecutionGuard; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Flow; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import org.junit.jupiter.api.Test; +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Flux; +import reactor.test.StepVerifier; + +/** + * Boundary behaviour of {@link ModelFallbackInterceptor}: a cancellation, however it is + * wrapped, must terminate the sequence instead of dispatching another provider call + * (issue #210). + */ +class ModelFallbackSwitchBoundaryTest { + + @Test + void wrappedCancellationNeverAttemptsAnotherClient() { + StubClient fallback = new StubClient("wrong"); + CancellationException cancelled = new CancellationException("stop"); + List attempts = new ArrayList<>(); + assertThatThrownBy(() -> interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + throw new CompletionException(cancelled); + })).isSameAs(cancelled); + assertThat(attempts).hasSize(1); + assertThat(fallback.requests).isEmpty(); + } + + @Test + void executionExceptionWrappedCancellationNeverAttemptsAnotherClient() { + StubClient fallback = new StubClient("wrong"); + CancellationException cancelled = new CancellationException("stop"); + List attempts = new ArrayList<>(); + assertThatThrownBy(() -> interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + // ExecutionException is checked, so it rides inside a runtime wrapper: the + // cause chain still has to be traversed to find the cancellation. + throw new RuntimeException(new ExecutionException("task aborted", cancelled)); + })).isSameAs(cancelled); + assertThat(attempts).hasSize(1); + } + + @Test + void deeplyWrappedCancellationUnwrapsToTheOriginalCancellation() { + StubClient fallback = new StubClient("wrong"); + CancellationException cancelled = new CancellationException("stop"); + List attempts = new ArrayList<>(); + assertThatThrownBy(() -> interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + throw new RuntimeException(new CompletionException(new RuntimeException(cancelled))); + })).isSameAs(cancelled); + assertThat(attempts).hasSize(1); + } + + @Test + void streamingFailureBeforeFirstEventSwitchesToFallback() { + StubClient fallback = new StubClient("fallback"); + RuntimeException primaryFailure = new IllegalStateException("primary stream failed"); + List attempts = new ArrayList<>(); + ModelResponse result = interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + if (request.getChatClient() == fallback) { + return ModelResponse.of(ChatMessage.assistant("fallback")); + } + return ModelResponse.stream(JdkFlowAdapter.publisherToFlowPublisher(Flux.error(primaryFailure))); + }); + + assertThat(result.isStreaming()).isTrue(); + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectNextMatches(event -> event.type() == ChatEvent.Type.RESPONSE + && "fallback".equals(event.response().text())) + .expectComplete() + .verify(Duration.ofSeconds(2)); + assertThat(attempts).hasSize(2); + } + + @Test + void streamingFailureAfterFirstEventTerminatesWithoutSwitching() { + StubClient fallback = new StubClient("fallback"); + RuntimeException failure = new IllegalStateException("stream died mid-output"); + List attempts = new ArrayList<>(); + ModelResponse result = interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + return ModelResponse.stream(JdkFlowAdapter.publisherToFlowPublisher( + Flux.concat(Flux.just(ChatEvent.textDelta("partial ")), Flux.error(failure)))); + }); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectNextMatches(event -> "partial ".equals(event.text())) + .expectErrorMatches(error -> error == failure) + .verify(Duration.ofSeconds(2)); + assertThat(attempts).as("a stream that already emitted output must not be switched mid-output").hasSize(1); + } + + @Test + void streamingCancellationBeforeFirstEventNeverSwitches() { + StubClient fallback = new StubClient("fallback"); + CancellationException cancelled = new CancellationException("stop"); + List attempts = new ArrayList<>(); + ModelResponse result = interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + return ModelResponse.stream(JdkFlowAdapter.publisherToFlowPublisher( + Flux.error(new CompletionException(cancelled)))); + }); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectErrorMatches(error -> error == cancelled) + .verify(Duration.ofSeconds(2)); + assertThat(attempts).hasSize(1); + } + + @Test + void usageAndToolCallEventsCountAsEmittedOutput() { + StubClient fallback = new StubClient("fallback"); + RuntimeException failure = new IllegalStateException("died after usage"); + List attempts = new ArrayList<>(); + ModelResponse result = interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + return ModelResponse.stream(JdkFlowAdapter.publisherToFlowPublisher( + Flux.concat(Flux.just(ChatEvent.usage(new Usage(1L, 2L, 3L))), Flux.error(failure)))); + }); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectNextMatches(event -> event.usage() != null) + .expectErrorMatches(error -> error == failure) + .verify(Duration.ofSeconds(2)); + assertThat(attempts).as("a usage event is observable output, so no switch may follow").hasSize(1); + } + + @Test + void deadRunGuardIsTerminalForBothPaths() { + StubClient first = new StubClient("first"); + StubClient second = new StubClient("second"); + List attempts = new ArrayList<>(); + MutableGuard guard = new MutableGuard(); + + // Synchronous path: the run dies while the primary attempt is in flight, so no + // fallback client may be dispatched after that. + assertThatThrownBy(() -> ModelFallbackInterceptor.builder() + .addFallbackClient(first) + .addFallbackClient(second) + .build() + .interceptModel(ModelRequest.builder().executionGuard(guard).build(), request -> { + attempts.add(request.getChatClient()); + guard.die(); + throw new IllegalStateException("primary failed"); + })).isSameAs(guard.stopped()); + assertThat(attempts).hasSize(1); + + // Streaming path: the run is alive at subscription entry and dies before the + // switch, so the switch itself must refuse to dispatch the next candidate. + guard.revive(); + attempts.clear(); + ModelResponse streaming = ModelFallbackInterceptor.builder() + .addFallbackClient(first) + .addFallbackClient(second) + .build() + .interceptModel(ModelRequest.builder().executionGuard(guard).build(), request -> { + attempts.add(request.getChatClient()); + return ModelResponse.stream( + JdkFlowAdapter.publisherToFlowPublisher(Flux.error(new IllegalStateException("boom")))); + }); + guard.die(); + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(streaming.getStream())) + .expectErrorMatches(error -> error == guard.stopped()) + .verify(Duration.ofSeconds(2)); + assertThat(attempts).as("the streaming switch must not dispatch a candidate on a dead run").hasSize(1); + } + + @Test + void deadGuardAtEntryDispatchesNoProviderCall() { + StubClient fallback = new StubClient("fallback"); + MutableGuard guard = new MutableGuard(); + guard.die(); + List attempts = new ArrayList<>(); + + assertThatThrownBy(() -> interceptor(fallback).interceptModel(ModelRequest.builder().executionGuard(guard) + .build(), request -> { + attempts.add(request.getChatClient()); + return ModelResponse.of(ChatMessage.assistant("never")); + })).isSameAs(guard.stopped()); + assertThat(attempts).as("a dead run must not pay for a primary call either").isEmpty(); + } + + @Test + void eachSubscriptionOwnsItsSwitchSequence() { + StubClient fallback = new StubClient("fallback"); + List attempts = new ArrayList<>(); + ModelResponse result = interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + if (request.getChatClient() == fallback) { + return ModelResponse.of(ChatMessage.assistant("fallback")); + } + return ModelResponse.stream( + JdkFlowAdapter.publisherToFlowPublisher(Flux.error(new IllegalStateException("primary died")))); + }); + + Flux stream = JdkFlowAdapter.flowPublisherToFlux(result.getStream()); + for (int subscription = 1; subscription <= 2; subscription++) { + StepVerifier.create(stream) + .expectNextMatches(event -> event.type() == ChatEvent.Type.RESPONSE + && "fallback".equals(event.response().text())) + .expectComplete() + .verify(Duration.ofSeconds(2)); + } + // The primary handler call happens once in interceptModel; a repeated + // subscription replays the captured publisher and owns fresh switch state, so + // the second subscription walks the same primary failure and switches again, + // mirroring the per-subscription attempt budget of the retry interceptor. + assertThat(attempts).containsExactly(null, fallback, fallback); + } + + @Test + void streamingCandidateFailureAdvancesToTheNextCandidate() { + StubClient second = new StubClient("second"); + List attempts = new ArrayList<>(); + ModelResponse result = ModelFallbackInterceptor.builder() + .addFallbackClient(new StubClient("unused-by-handler")) + .addFallbackClient(second) + .build() + .interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + if (request.getChatClient() == second) { + return ModelResponse.of(ChatMessage.assistant("second")); + } + // Primary and the first fallback candidate both fail their stream + // before emitting any event. + return ModelResponse.stream( + JdkFlowAdapter.publisherToFlowPublisher(Flux.error(new IllegalStateException("died")))); + }); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectNextMatches(event -> event.type() == ChatEvent.Type.RESPONSE + && "second".equals(event.response().text())) + .expectComplete() + .verify(Duration.ofSeconds(2)); + assertThat(attempts).hasSize(3); + } + + @Test + void synchronousFailureDuringAStreamingSwitchAdvancesToTheNextCandidate() { + StubClient second = new StubClient("second"); + List attempts = new ArrayList<>(); + ModelResponse result = ModelFallbackInterceptor.builder() + .addFallbackClient(new StubClient("unused-by-handler")) + .addFallbackClient(second) + .build() + .interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + if (request.getChatClient() == null) { + // The primary fails its stream before any event is emitted. + return ModelResponse.stream( + JdkFlowAdapter.publisherToFlowPublisher(Flux.error(new IllegalStateException("died")))); + } + if (request.getChatClient() == second) { + return ModelResponse.of(ChatMessage.assistant("second")); + } + // The first fallback candidate fails synchronously: the call itself is + // constructed (or rejected) without ever producing a stream. + throw new IllegalStateException("call construction failed"); + }); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectNextMatches(event -> event.type() == ChatEvent.Type.RESPONSE + && "second".equals(event.response().text())) + .expectComplete() + .verify(Duration.ofSeconds(2)); + assertThat(attempts).hasSize(3); + } + + @Test + void recommendedOrderRetriesEachModelBeforeFailingOver() { + StubClient fallback = new StubClient("fallback"); + List modelCalls = new ArrayList<>(); + AtomicInteger primaryAttempts = new AtomicInteger(); + ModelFallbackInterceptor fallbackInterceptor = ModelFallbackInterceptor.builder() + .addFallbackClient(fallback) + .build(); + ModelRetryInterceptor retryInterceptor = ModelRetryInterceptor.builder() + .maxAttempts(2) + .initialDelay(1) + // The composition is under test, not the default message-based + // classification of retryable failures. + .retryableExceptionPredicate(failure -> true) + .build(); + // The agent applies the interceptor list outermost-first, so registering the + // fallback before the retry gives "retry each model, then fail over". + ModelCallHandler model = request -> { + modelCalls.add(String.valueOf(request.getChatClient())); + if (request.getChatClient() == fallback) { + return ModelResponse.of(ChatMessage.assistant("fallback")); + } + primaryAttempts.incrementAndGet(); + return ModelResponse.stream( + JdkFlowAdapter.publisherToFlowPublisher(Flux.error(new IllegalStateException("primary died")))); + }; + ModelCallHandler chain = request -> retryInterceptor.interceptModel(request, model); + + ModelResponse result = fallbackInterceptor.interceptModel(ModelRequest.builder().build(), chain); + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectNextMatches(event -> event.type() == ChatEvent.Type.RESPONSE + && "fallback".equals(event.response().text())) + .expectComplete() + .verify(Duration.ofSeconds(5)); + + assertThat(primaryAttempts).as("the inner retry exhausts the primary before the failover").hasValue(2); + assertThat(modelCalls).hasSize(3); + } + + @Test + void cancellationAfterFirstEventKeepsItsUnwrappedIdentity() { + StubClient fallback = new StubClient("fallback"); + CancellationException cancelled = new CancellationException("stop"); + List attempts = new ArrayList<>(); + ModelResponse result = interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + return ModelResponse.stream(JdkFlowAdapter.publisherToFlowPublisher( + Flux.concat(Flux.just(ChatEvent.textDelta("partial ")), + Flux.error(new CompletionException(cancelled))))); + }); + + // Before or after the first event, the same wrapped cancellation surfaces as + // the same bare cancellation, so downstream cancellation handling cannot + // misclassify it as a model failure by testing the top-level type only. + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectNextMatches(event -> "partial ".equals(event.text())) + .expectErrorMatches(error -> error == cancelled) + .verify(Duration.ofSeconds(2)); + assertThat(attempts).hasSize(1); + } + + @Test + void toolCallEventCountsAsEmittedOutput() { + StubClient fallback = new StubClient("fallback"); + RuntimeException failure = new IllegalStateException("died after a tool call"); + List attempts = new ArrayList<>(); + ModelResponse result = interceptor(fallback).interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + return ModelResponse.stream(JdkFlowAdapter.publisherToFlowPublisher( + Flux.concat(Flux.just(ChatEvent.toolCall(new ToolCall("call-1", "lookup", "{}"))), + Flux.error(failure)))); + }); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectNextMatches(event -> event.type() == ChatEvent.Type.TOOL_CALL) + .expectErrorMatches(error -> error == failure) + .verify(Duration.ofSeconds(2)); + assertThat(attempts).as("a delivered tool call is observable output, so no switch may follow").hasSize(1); + } + + @Test + void postEventFailureAtAFallbackCandidateDoesNotAdvanceFurther() { + StubClient second = new StubClient("second"); + RuntimeException failure = new IllegalStateException("fallback stream died mid-output"); + List attempts = new ArrayList<>(); + ModelResponse result = ModelFallbackInterceptor.builder() + .addFallbackClient(new StubClient("unused-by-handler")) + .addFallbackClient(second) + .build() + .interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + if (request.getChatClient() == null) { + return ModelResponse.stream( + JdkFlowAdapter.publisherToFlowPublisher(Flux.error(new IllegalStateException("died")))); + } + // The first fallback candidate emits output and then fails. + return ModelResponse.stream(JdkFlowAdapter.publisherToFlowPublisher( + Flux.concat(Flux.just(ChatEvent.textDelta("fallback partial ")), Flux.error(failure)))); + }); + + StepVerifier.create(JdkFlowAdapter.flowPublisherToFlux(result.getStream())) + .expectNextMatches(event -> "fallback partial ".equals(event.text())) + .expectErrorMatches(error -> error == failure) + .verify(Duration.ofSeconds(2)); + assertThat(attempts).as("the boundary applies per candidate, not only to the primary").hasSize(2); + } + + @Test + void wrappedCancellationAtAFallbackAttemptStopsTheSequence() { + StubClient first = new StubClient("first"); + StubClient second = new StubClient("second"); + CancellationException cancelled = new CancellationException("stop"); + List attempts = new ArrayList<>(); + assertThatThrownBy(() -> ModelFallbackInterceptor.builder() + .addFallbackClient(first) + .addFallbackClient(second) + .build() + .interceptModel(ModelRequest.builder().build(), request -> { + attempts.add(request.getChatClient()); + if (request.getChatClient() == null) { + throw new IllegalStateException("primary failed"); + } + if (request.getChatClient() == first) { + throw new CompletionException(cancelled); + } + return ModelResponse.of(ChatMessage.assistant("second")); + })).isSameAs(cancelled); + assertThat(attempts).containsExactly(null, first); + } + + private static ModelFallbackInterceptor interceptor(ChatClient fallback) { + return ModelFallbackInterceptor.builder().addFallbackClient(fallback).build(); + } + + /** A run guard the test can kill and revive to place the run's death precisely. */ + private static final class MutableGuard implements ExecutionGuard { + + private final AtomicBoolean active = new AtomicBoolean(true); + + private final CancellationException stopped = new CancellationException("run stopped"); + + void die() { + active.set(false); + } + + void revive() { + active.set(true); + } + + CancellationException stopped() { + return stopped; + } + + @Override + public void assertActive() { + if (!active.get()) { + throw stopped; + } + } + + @Override + public AutoCloseable onLoss(Runnable cancellation) { + return () -> { + }; + } + + } + + private static final class StubClient implements ChatClient { + + private final String text; + + private final List requests = new ArrayList<>(); + + private StubClient(String text) { + this.text = text; + } + + @Override + public CompletionStage call(ChatRequest request) { + requests.add(request); + return CompletableFuture.completedFuture(new ChatResponse("id", "test-model", + List.of(ChatMessage.assistant(text)), List.of(), null, new Usage(2L, 3L, 5L), Map.of())); + } + + @Override + public Flow.Publisher stream(ChatRequest request) { + throw new UnsupportedOperationException("Streaming is unavailable"); + } + + @Override + public ChatClientCapabilities capabilities() { + return ChatClientCapabilities.basic(); + } + + } + +}