From 949e82eabda4dbe7d3be7e1a91e9a8534b239f14 Mon Sep 17 00:00:00 2001 From: hengyuss <1183660933@qq.com> Date: Tue, 3 Nov 2026 16:12:02 +0800 Subject: [PATCH] fix(agent): honor request-selected chat client --- .../agentic/ai/agent/react/ReActAgent.java | 11 +- .../ai/agent/react/ReActAgentTest.java | 173 ++++++++++++++++++ ...ActStreamSubscriptionCancellationTest.java | 40 ++++ 3 files changed, 219 insertions(+), 5 deletions(-) diff --git a/argi-agent/argi-agent-core/src/main/java/io/github/agentic/ai/agent/react/ReActAgent.java b/argi-agent/argi-agent-core/src/main/java/io/github/agentic/ai/agent/react/ReActAgent.java index fe2c16d46..728a489b7 100644 --- a/argi-agent/argi-agent-core/src/main/java/io/github/agentic/ai/agent/react/ReActAgent.java +++ b/argi-agent/argi-agent-core/src/main/java/io/github/agentic/ai/agent/react/ReActAgent.java @@ -469,12 +469,13 @@ private ModelResponse invokeModel(ModelRequest request, UsageAccumulator usage, AtomicReference effectiveTools) { ToolRegistry effective = effectiveToolSet(request); effectiveTools.set(effective); + ChatClient selectedClient = request.getChatClient() != null ? request.getChatClient() : chatClient; ChatRequest chatRequest = new ChatRequest(composeMessages(request), effective.definitions(), request.getOptions(), request.getAttributes()); if (streamModel) { - return ModelResponse.of(collectModelStream(chatRequest, usage)); + return ModelResponse.of(collectModelStream(selectedClient, chatRequest, usage)); } - CompletableFuture future = chatClient.call(chatRequest).toCompletableFuture(); + CompletableFuture future = selectedClient.call(chatRequest).toCompletableFuture(); toolOperation.set(future); try (AutoCloseable ignored = context.control().onLoss(() -> future.cancel(true))) { ChatResponse response = await(future); @@ -490,9 +491,9 @@ private ModelResponse invokeModel(ModelRequest request, UsageAccumulator usage, } } - private ChatResponse collectModelStream(ChatRequest request, UsageAccumulator usage) { - return collectModelStream(Objects.requireNonNull(chatClient.stream(request), "chat client returned no stream"), - usage); + private ChatResponse collectModelStream(ChatClient selectedClient, ChatRequest request, UsageAccumulator usage) { + return collectModelStream( + Objects.requireNonNull(selectedClient.stream(request), "chat client returned no stream"), usage); } private ChatResponse collectModelStream(Flow.Publisher publisher, UsageAccumulator usage) { diff --git a/argi-agent/argi-agent-core/src/test/java/io/github/agentic/ai/agent/react/ReActAgentTest.java b/argi-agent/argi-agent-core/src/test/java/io/github/agentic/ai/agent/react/ReActAgentTest.java index 7d1ec933e..2f7254324 100644 --- a/argi-agent/argi-agent-core/src/test/java/io/github/agentic/ai/agent/react/ReActAgentTest.java +++ b/argi-agent/argi-agent-core/src/test/java/io/github/agentic/ai/agent/react/ReActAgentTest.java @@ -56,6 +56,7 @@ import java.util.Queue; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Flow; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; @@ -291,6 +292,149 @@ public ModelResponse interceptModel(ModelRequest request, ModelCallHandler next) .map(ChatMessage::text).toList(); assertEquals(List.of("required-policy", "first", "reply", "second"), texts); } + + // ------------------------------------------------------------------ + // issue #271: request-selected clients reach the model invocation + // ------------------------------------------------------------------ + + @Test + void interceptorSelectedClientOverridesTheDefaultClient() throws Exception { + ScriptedClient primary = new ScriptedClient( + List.of(response(ChatMessage.assistant("primary"), List.of(), 1))); + ScriptedClient selected = new ScriptedClient( + List.of(response(ChatMessage.assistant("selected"), List.of(), 1))); + ModelInterceptor interceptor = selectingClient(selected); + + AgentOutcome outcome = ReActAgent.builder(primary).interceptors(List.of(interceptor)).build() + .invoke(request(), context()) + .get(3, TimeUnit.SECONDS); + + assertEquals("selected", outcome.messages().get(outcome.messages().size() - 1).text()); + assertEquals(0, primary.requests.size()); + assertEquals(1, selected.requests.size()); + } + + @Test + void streamingRunUsesTheInterceptorSelectedClient() throws Exception { + StreamingClient primary = new StreamingClient(); + StreamingClient selected = new StreamingClient(); + ModelInterceptor interceptor = selectingClient(selected); + + ReActAgent.builder(primary).streamModel(true).interceptors(List.of(interceptor)).build() + .invoke(request(), context()) + .get(3, TimeUnit.SECONDS); + + assertEquals(0, primary.requests.size()); + assertEquals(1, selected.requests.size()); + } + + @Test + void fallbackRetryUsesTheClientSelectedOnTheRetryRequest() throws Exception { + AtomicInteger primaryCalls = new AtomicInteger(); + ChatClient primary = new ChatClient() { + public CompletionStage call(ChatRequest request) { + primaryCalls.incrementAndGet(); + return CompletableFuture.failedFuture(new IllegalStateException("primary unavailable")); + } + + public Flow.Publisher stream(ChatRequest request) { + throw new UnsupportedOperationException(); + } + + public ChatClientCapabilities capabilities() { + return ChatClientCapabilities.basic(); + } + }; + ScriptedClient fallback = new ScriptedClient( + List.of(response(ChatMessage.assistant("fallback"), List.of(), 1))); + ModelInterceptor interceptor = new ModelInterceptor() { + public String getName() { + return "fallback"; + } + + public ModelResponse interceptModel(ModelRequest request, ModelCallHandler next) { + try { + return next.call(request); + } + catch (RuntimeException primaryFailure) { + return next.call(ModelRequest.builder(request).chatClient(fallback).build()); + } + } + }; + + AgentOutcome outcome = ReActAgent.builder(primary).interceptors(List.of(interceptor)).build() + .invoke(request(), context()) + .get(3, TimeUnit.SECONDS); + + assertEquals("fallback", outcome.messages().get(outcome.messages().size() - 1).text()); + assertEquals(1, primaryCalls.get()); + assertEquals(1, fallback.requests.size()); + } + + @Test + void allFailedCandidatesReportTheSelectedClientFailure() { + AtomicInteger primaryCalls = new AtomicInteger(); + AtomicInteger fallbackCalls = new AtomicInteger(); + ChatClient primary = failingClient(primaryCalls, "primary unavailable"); + ChatClient fallback = failingClient(fallbackCalls, "fallback unavailable"); + ModelInterceptor interceptor = new ModelInterceptor() { + public String getName() { + return "fallback"; + } + + public ModelResponse interceptModel(ModelRequest request, ModelCallHandler next) { + try { + return next.call(request); + } + catch (RuntimeException primaryFailure) { + return next.call(ModelRequest.builder(request).chatClient(fallback).build()); + } + } + }; + + ExecutionException failure = assertThrows(ExecutionException.class, + () -> ReActAgent.builder(primary).interceptors(List.of(interceptor)).build() + .invoke(request(), context()) + .get(3, TimeUnit.SECONDS)); + + assertEquals("fallback unavailable", failure.getCause().getMessage()); + assertEquals(1, primaryCalls.get()); + assertEquals(1, fallbackCalls.get()); + } + + @Test + void concurrentRunsKeepTheirClientSelectionsIsolated() throws Exception { + ScriptedClient primary = new ScriptedClient( + List.of(response(ChatMessage.assistant("primary"), List.of(), 1))); + ScriptedClient selected = new ScriptedClient( + List.of(response(ChatMessage.assistant("selected"), List.of(), 1))); + ModelInterceptor interceptor = new ModelInterceptor() { + public String getName() { + return "route-client"; + } + + public ModelResponse interceptModel(ModelRequest request, ModelCallHandler next) { + boolean useSelected = request.getMessages().stream() + .anyMatch(message -> "route-selected".equals(message.text())); + return useSelected + ? next.call(ModelRequest.builder(request).chatClient(selected).build()) + : next.call(request); + } + }; + ReActAgent agent = ReActAgent.builder(primary).interceptors(List.of(interceptor)).build(); + + CompletableFuture primaryRun = agent.invoke( + AgentRequest.messages(List.of(ChatMessage.user("route-primary"))), context()); + CompletableFuture selectedRun = agent.invoke( + AgentRequest.messages(List.of(ChatMessage.user("route-selected"))), context()); + AgentOutcome primaryOutcome = primaryRun.get(3, TimeUnit.SECONDS); + AgentOutcome selectedOutcome = selectedRun.get(3, TimeUnit.SECONDS); + + assertEquals("primary", primaryOutcome.messages().get(primaryOutcome.messages().size() - 1).text()); + assertEquals("selected", selectedOutcome.messages().get(selectedOutcome.messages().size() - 1).text()); + assertEquals(1, primary.requests.size()); + assertEquals(1, selected.requests.size()); + } @Test void emitsStreamingDeltasAndSupportsDynamicTools() throws Exception { @@ -864,6 +1008,35 @@ public CompletionStage execute(ToolCall call, ToolContext context) { }; } + private static ModelInterceptor selectingClient(ChatClient selected) { + return new ModelInterceptor() { + public String getName() { + return "select-client"; + } + + public ModelResponse interceptModel(ModelRequest request, ModelCallHandler next) { + return next.call(ModelRequest.builder(request).chatClient(selected).build()); + } + }; + } + + private static ChatClient failingClient(AtomicInteger calls, String message) { + return new ChatClient() { + public CompletionStage call(ChatRequest request) { + calls.incrementAndGet(); + return CompletableFuture.failedFuture(new IllegalStateException(message)); + } + + public Flow.Publisher stream(ChatRequest request) { + throw new UnsupportedOperationException(); + } + + public ChatClientCapabilities capabilities() { + return ChatClientCapabilities.basic(); + } + }; + } + private static ChatResponse response(ChatMessage message, List calls, long total) { return new ChatResponse(null, "test", List.of(message), calls, FinishReason.STOP, new Usage(1L, 1L, total), Map.of()); diff --git a/argi-agent/argi-agent-core/src/test/java/io/github/agentic/ai/agent/react/ReActStreamSubscriptionCancellationTest.java b/argi-agent/argi-agent-core/src/test/java/io/github/agentic/ai/agent/react/ReActStreamSubscriptionCancellationTest.java index 720c5cd60..dccfb2a56 100644 --- a/argi-agent/argi-agent-core/src/test/java/io/github/agentic/ai/agent/react/ReActStreamSubscriptionCancellationTest.java +++ b/argi-agent/argi-agent-core/src/test/java/io/github/agentic/ai/agent/react/ReActStreamSubscriptionCancellationTest.java @@ -19,6 +19,10 @@ import io.github.agentic.ai.agent.AgentExecutionContext; import io.github.agentic.ai.agent.AgentEvent; import io.github.agentic.ai.agent.AgentRequest; +import io.github.agentic.ai.agent.interceptor.ModelCallHandler; +import io.github.agentic.ai.agent.interceptor.ModelInterceptor; +import io.github.agentic.ai.agent.interceptor.ModelRequest; +import io.github.agentic.ai.agent.interceptor.ModelResponse; import io.github.agentic.ai.api.chat.ChatClient; import io.github.agentic.ai.api.chat.ChatClientCapabilities; import io.github.agentic.ai.api.chat.ChatEvent; @@ -86,6 +90,42 @@ void cancelBeforeModelOnSubscribeNeverRequests() throws Exception { assertTrue(events.isEmpty(), "a cancelled run must not deliver events"); } + @Test + void cancellationTargetsTheClientSelectedByTheInterceptor() throws Exception { + DeferredPublisher defaultModel = new DeferredPublisher( + response(ChatMessage.assistant("default"), List.of(), 1), false); + DeferredPublisher selectedModel = new DeferredPublisher( + response(ChatMessage.assistant("selected"), List.of(), 1), false); + ScriptedPublisherClient defaultClient = new ScriptedPublisherClient(List.of(defaultModel)); + ScriptedPublisherClient selectedClient = new ScriptedPublisherClient(List.of(selectedModel)); + ModelInterceptor interceptor = new ModelInterceptor() { + @Override + public String getName() { + return "select-client"; + } + + @Override + public ModelResponse interceptModel(ModelRequest request, ModelCallHandler next) { + return next.call(ModelRequest.builder(request).chatClient(selectedClient).build()); + } + }; + List events = new CopyOnWriteArrayList<>(); + try (AgentExecutionContext context = context()) { + Flow.Subscription subscription = subscribe(ReActAgent.builder(defaultClient) + .streamModel(true) + .interceptors(List.of(interceptor)) + .build(), events, context); + + assertTrue(selectedModel.subscribed.await(3, TimeUnit.SECONDS), "selected publisher was never subscribed"); + subscription.cancel(); + selectedModel.deliverOnSubscribe(); + + assertTrue(defaultClient.requests.isEmpty(), "the default client must not be invoked"); + assertEquals(0, selectedModel.requests.get()); + assertEquals(1, selectedModel.cancels.get()); + } + } + @Test void normalStreamingCompletionRequestsOnceAndIsNotCancelled() throws Exception { DeferredPublisher model = new DeferredPublisher(response(ChatMessage.assistant("answer"), List.of(), 1), true);