From a54e28b1d9acbf08eec44cc57a13415b0f3d757b Mon Sep 17 00:00:00 2001 From: webmin Date: Thu, 9 Jul 2026 10:06:16 +0800 Subject: [PATCH] fix(model/openai): wrap streaming branch in Flux.defer to make retry re-issue HTTP requests OpenAIChatModel.doStream0's streaming branch returned a Flux tied to a single transport.stream() invocation. For transports whose stream() is backed by a one-shot future (e.g. JdkHttpTransport), retryWhen re- subscriptions replayed the already-failed publisher instead of sending a new HTTP request. This caused 429/5xx retry to exhaust after a single attempt. Wrap the streaming branch in Flux.defer, matching the non-streaming branch, so each retry creates a fresh publisher and re-issues the HTTP request. Add OpenAIChatModelStreamingRetryTest to verify that streaming 429 responses are retried and hit the server maxAttempts times for OkHttpTransport, JdkHttpTransport, and the default transport. --- .../model/openai/OpenAIChatModel.java | 16 +- .../OpenAIChatModelStreamingRetryTest.java | 316 ++++++++++++++++++ 2 files changed, 328 insertions(+), 4 deletions(-) create mode 100644 agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-openai/src/test/java/io/agentscope/extensions/model/openai/OpenAIChatModelStreamingRetryTest.java diff --git a/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-openai/src/main/java/io/agentscope/extensions/model/openai/OpenAIChatModel.java b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-openai/src/main/java/io/agentscope/extensions/model/openai/OpenAIChatModel.java index bf1b14389c..cd568a5700 100644 --- a/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-openai/src/main/java/io/agentscope/extensions/model/openai/OpenAIChatModel.java +++ b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-openai/src/main/java/io/agentscope/extensions/model/openai/OpenAIChatModel.java @@ -161,10 +161,18 @@ protected Flux doStream0( // Make the API call if (stream) { - // Streaming mode - return client.stream(apiKey, baseUrl, request, effectiveOptions) - .map(response -> formatter.parseResponse(response, start)) - .filter(Objects::nonNull); + // Streaming mode: defer creation so that retryWhen re-subscriptions re-issue + // the HTTP request. Without defer, transports whose stream() is tied to a + // single future (e.g. JdkHttpTransport) would replay the same failed stream. + return Flux.defer( + () -> + client.stream(apiKey, baseUrl, request, effectiveOptions) + .map( + response -> + formatter.parseResponse( + response, start)) + .filter(Objects::nonNull)) + .subscribeOn(Schedulers.boundedElastic()); } else { // Non-streaming mode: make a single call and return as Flux return Flux.defer( diff --git a/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-openai/src/test/java/io/agentscope/extensions/model/openai/OpenAIChatModelStreamingRetryTest.java b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-openai/src/test/java/io/agentscope/extensions/model/openai/OpenAIChatModelStreamingRetryTest.java new file mode 100644 index 0000000000..6a9ef7a49c --- /dev/null +++ b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-openai/src/test/java/io/agentscope/extensions/model/openai/OpenAIChatModelStreamingRetryTest.java @@ -0,0 +1,316 @@ +/* + * 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 + * + * http://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.agentscope.extensions.model.openai; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.agentscope.core.message.Msg; +import io.agentscope.core.message.MsgRole; +import io.agentscope.core.message.TextBlock; +import io.agentscope.core.model.ChatResponse; +import io.agentscope.core.model.ExecutionConfig; +import io.agentscope.core.model.GenerateOptions; +import io.agentscope.core.model.transport.HttpTransport; +import io.agentscope.core.model.transport.HttpTransportConfig; +import io.agentscope.core.model.transport.HttpTransportFactory; +import io.agentscope.core.model.transport.JdkHttpTransport; +import io.agentscope.core.model.transport.OkHttpTransport; +import io.agentscope.extensions.model.openai.exception.OpenAIException; +import io.agentscope.extensions.model.openai.formatter.OpenAIChatFormatter; +import java.io.IOException; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import okhttp3.mockwebserver.Dispatcher; +import okhttp3.mockwebserver.MockResponse; +import okhttp3.mockwebserver.MockWebServer; +import okhttp3.mockwebserver.RecordedRequest; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import reactor.test.StepVerifier; + +/** + * Regression tests for streaming retry behavior in {@link OpenAIChatModel}. + * + *

Verifies that when a streaming request receives a retryable HTTP error (e.g. 429), + * {@code retryWhen} re-subscribes to a fresh publisher that re-issues the HTTP request, + * rather than replaying the same failed stream. The bug primarily manifests with + * {@link JdkHttpTransport}, whose stream() is backed by a one-shot + * {@link java.util.concurrent.CompletableFuture}. + */ +class OpenAIChatModelStreamingRetryTest { + + private static final int MAX_ATTEMPTS = 3; + private static final Duration BACKOFF = Duration.ofMillis(10); + + private MockWebServer mockServer; + + @BeforeEach + void setUp() throws IOException { + mockServer = new MockWebServer(); + mockServer.start(); + } + + @AfterEach + void tearDown() throws IOException { + mockServer.shutdown(); + } + + @Test + @DisplayName("Streaming 429 exhausts retries and re-issues HTTP request each time") + void streaming429ExhaustsRetriesAndReissuesEachRequest() { + AtomicInteger requestCount = startAlways429Server(); + + OpenAIChatModel model = buildStreamingModel(defaultTransport()); + + StepVerifier.create(model.stream(List.of(singleUserMessage()), null, retryOptions())) + .expectErrorSatisfies(e -> assertInstanceOf(OpenAIException.class, rootCause(e))) + .verify(Duration.ofSeconds(10)); + + assertEquals( + MAX_ATTEMPTS, requestCount.get(), "Each retry should re-issue the HTTP request"); + } + + @Test + @DisplayName("Streaming 429 retries then succeeds on subsequent attempt") + void streaming429RetriesThenSucceeds() { + AtomicInteger requestCount = new AtomicInteger(0); + mockServer.setDispatcher( + new Dispatcher() { + @Override + public MockResponse dispatch(RecordedRequest request) { + int attempt = requestCount.incrementAndGet(); + return attempt < MAX_ATTEMPTS + ? rateLimitResponse() + : successStreamResponse("hello"); + } + }); + + OpenAIChatModel model = buildStreamingModel(defaultTransport()); + + List responses = + model.stream(List.of(singleUserMessage()), null, retryOptions()) + .collectList() + .block(Duration.ofSeconds(10)); + + assertEquals(MAX_ATTEMPTS, requestCount.get()); + assertTrue(responses != null && !responses.isEmpty(), "Expected at least one response"); + assertEquals("hello", firstText(responses.get(0))); + } + + @Test + @DisplayName("Streaming 400 does not retry") + void streaming400DoesNotRetry() { + mockServer.enqueue( + new MockResponse() + .setResponseCode(400) + .setHeader("Content-Type", "application/json") + .setBody("{\"error\":{\"message\":\"bad request\"}}")); + + OpenAIChatModel model = buildStreamingModel(defaultTransport()); + + StepVerifier.create(model.stream(List.of(singleUserMessage()), null, retryOptions())) + .expectErrorSatisfies(e -> assertInstanceOf(OpenAIException.class, rootCause(e))) + .verify(Duration.ofSeconds(5)); + + assertEquals(1, mockServer.getRequestCount(), "400 errors must not be retried"); + } + + @Test + @DisplayName("Streaming 429 with JdkHttpTransport retries and re-issues HTTP requests") + void streaming429WithJdkHttpTransportRetries() { + AtomicInteger requestCount = startAlways429Server(); + + OpenAIChatModel model = + buildStreamingModel( + JdkHttpTransport.builder() + .config( + HttpTransportConfig.builder() + .readTimeout(Duration.ofSeconds(5)) + .build()) + .build()); + + StepVerifier.create(model.stream(List.of(singleUserMessage()), null, retryOptions())) + .expectError() + .verify(Duration.ofSeconds(10)); + + assertEquals( + MAX_ATTEMPTS, + requestCount.get(), + "JdkHttpTransport streaming retry should re-issue HTTP request on each attempt"); + } + + @Test + @DisplayName("Streaming 429 with OkHttpTransport retries and re-issues HTTP requests") + void streaming429WithOkHttpTransportRetries() { + AtomicInteger requestCount = startAlways429Server(); + + OpenAIChatModel model = + buildStreamingModel( + OkHttpTransport.builder() + .config( + HttpTransportConfig.builder() + .readTimeout(Duration.ofSeconds(5)) + .build()) + .build()); + + StepVerifier.create(model.stream(List.of(singleUserMessage()), null, retryOptions())) + .expectError() + .verify(Duration.ofSeconds(10)); + + assertEquals( + MAX_ATTEMPTS, + requestCount.get(), + "OkHttpTransport streaming retry should re-issue HTTP request on each attempt"); + } + + @Test + @DisplayName("Non-streaming 429 retries and re-issues HTTP requests") + void nonStreaming429Retries() { + AtomicInteger requestCount = startAlways429Server(); + + OpenAIChatModel model = + buildModel( + false, + OkHttpTransport.builder() + .config( + HttpTransportConfig.builder() + .readTimeout(Duration.ofSeconds(5)) + .build()) + .build()); + + StepVerifier.create(model.stream(List.of(singleUserMessage()), null, retryOptions(false))) + .expectError() + .verify(Duration.ofSeconds(10)); + + assertEquals( + MAX_ATTEMPTS, + requestCount.get(), + "Non-streaming retry should re-issue HTTP request on each attempt"); + } + + // ==================== helpers ==================== + + private HttpTransport defaultTransport() { + HttpTransportFactory.shutdown(); + return HttpTransportFactory.getDefault(); + } + + private AtomicInteger startAlways429Server() { + AtomicInteger requestCount = new AtomicInteger(0); + mockServer.setDispatcher( + new Dispatcher() { + @Override + public MockResponse dispatch(RecordedRequest request) { + requestCount.incrementAndGet(); + return rateLimitResponse(); + } + }); + return requestCount; + } + + private OpenAIChatModel buildStreamingModel(HttpTransport transport) { + return buildModel(true, transport); + } + + private OpenAIChatModel buildModel(boolean stream, HttpTransport transport) { + String baseUrl = mockServer.url("/").toString(); + if (baseUrl.endsWith("/")) { + baseUrl = baseUrl.substring(0, baseUrl.length() - 1); + } + + return OpenAIChatModel.builder().apiKey("test-key").modelName("gpt-4").stream(stream) + .baseUrl(baseUrl) + .formatter(new OpenAIChatFormatter()) + .httpTransport(transport) + .build(); + } + + private GenerateOptions retryOptions() { + return retryOptions(true); + } + + private GenerateOptions retryOptions(boolean stream) { + return GenerateOptions.builder() + .executionConfig( + ExecutionConfig.builder() + .maxAttempts(MAX_ATTEMPTS) + .initialBackoff(BACKOFF) + .maxBackoff(Duration.ofMillis(50)) + .build()) + .stream(stream) + .build(); + } + + private Msg singleUserMessage() { + return Msg.builder() + .role(MsgRole.USER) + .content(TextBlock.builder().text("hello").build()) + .build(); + } + + private static MockResponse rateLimitResponse() { + return new MockResponse() + .setResponseCode(429) + .setHeader("Content-Type", "application/json") + .setBody( + "{\"error\":{\"message\":\"rate" + + " limit\",\"type\":\"rate_limit_exceeded\"}}"); + } + + private static MockResponse successStreamResponse(String content) { + return new MockResponse() + .setResponseCode(200) + .setHeader("Content-Type", "text/event-stream") + .setBody(sseChunk(content) + sseDone()); + } + + private static String sseChunk(String content) { + return "data: {\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"" + + content + + "\"},\"finish_reason\":null}]}\n\n"; + } + + private static String sseDone() { + return "data: [DONE]\n\n"; + } + + private static String firstText(ChatResponse response) { + return response.getContent().stream() + .filter(TextBlock.class::isInstance) + .map(TextBlock.class::cast) + .map(TextBlock::getText) + .findFirst() + .orElse(""); + } + + /** + * Walks the cause chain to the innermost non-null cause. Reactor's retryWhen wraps the + * final failure in a RetryExhaustedException, so assertions on the root cause are more + * stable than matching the wrapper. + */ + private static Throwable rootCause(Throwable t) { + Throwable cursor = t; + while (cursor.getCause() != null && cursor.getCause() != cursor) { + cursor = cursor.getCause(); + } + return cursor; + } +}