This is an automated email from the ASF dual-hosted git repository.
Aias00 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shenyu.git
The following commit(s) were added to refs/heads/master by this push:
new 059748ce25 fix(ai): prevent token usage loss across streaming flush
boundaries (#7730)
059748ce25 is described below
commit 059748ce25ad8dd532412657ad66f2964a3ab0d0
Author: lymerin <[email protected]>
AuthorDate: Mon Oct 19 09:33:44 2026 +0800
fix(ai): prevent token usage loss across streaming flush boundaries (#7730)
---
.../ai/token/limiter/AiTokenLimiterPlugin.java | 52 ++++++++-----
.../ai/token/limiter/AiTokenLimiterPluginTest.java | 86 ++++++++++++++++++++++
2 files changed, 118 insertions(+), 20 deletions(-)
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/main/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPlugin.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/main/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPlugin.java
index dd4e595217..3631130958 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/main/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPlugin.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/main/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPlugin.java
@@ -217,27 +217,38 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
@Override
@NonNull
public Mono<Void> writeWith(@NonNull final Publisher<? extends
DataBuffer> body) {
- return super.writeWith(appendResponse(body));
+ return Mono.defer(() -> {
+ final BodyWriter writer = new BodyWriter();
+ final GzipStreamDecoder decoder = createDecoder();
+ return super.writeWith(appendResponse(body, writer, decoder))
+ .doFinally(signal -> finishResponse(writer, decoder));
+ });
}
@Override
@NonNull
public Mono<Void> writeAndFlushWith(@NonNull final Publisher<? extends
Publisher<? extends DataBuffer>> body) {
- Flux<? extends Publisher<? extends DataBuffer>> intercepted =
Flux.from(body)
- .map(this::appendResponse);
- return super.writeAndFlushWith(intercepted);
+ return Mono.defer(() -> {
+ final BodyWriter writer = new BodyWriter();
+ final GzipStreamDecoder decoder = createDecoder();
+ Flux<? extends Publisher<? extends DataBuffer>> intercepted =
Flux.from(body)
+ .map(inner -> appendResponse(inner, writer, decoder));
+ return super.writeAndFlushWith(intercepted)
+ .doFinally(signal -> finishResponse(writer, decoder));
+ });
}
- @NonNull
- private Flux<? extends DataBuffer> appendResponse(final Publisher<?
extends DataBuffer> body) {
- BodyWriter writer = new BodyWriter();
+ private GzipStreamDecoder createDecoder() {
HttpHeaders headers = serverHttpResponse.getHeaders();
boolean isGzip = headers.containsKey(Constants.CONTENT_ENCODING)
&& headers.getFirst(Constants.CONTENT_ENCODING)
.contains(Constants.HTTP_ACCEPT_ENCODING_GZIP);
+ return isGzip ? new GzipStreamDecoder() : null;
+ }
- final GzipStreamDecoder decoder = isGzip ? new GzipStreamDecoder()
: null;
-
+ @NonNull
+ private Flux<? extends DataBuffer> appendResponse(final Publisher<?
extends DataBuffer> body,
+ final BodyWriter
writer, final GzipStreamDecoder decoder) {
return Flux.<DataBuffer>from(body)
.doOnNext(buffer -> {
try (DataBuffer.ByteBufferIterator it =
buffer.readableByteBuffers()) {
@@ -246,7 +257,7 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
byte[] inBytes = new byte[ro.remaining()];
ro.get(inBytes);
- byte[] processedBytes = isGzip ?
decoder.decode(inBytes) : inBytes;
+ byte[] processedBytes =
Objects.nonNull(decoder) ? decoder.decode(inBytes) : inBytes;
if (processedBytes.length > 0) {
processChunk(processedBytes, writer);
}
@@ -254,19 +265,20 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
} catch (Exception e) {
LOG.error("read dataBuffer error", e);
}
- })
- .doFinally(signal -> {
- if (Objects.nonNull(decoder)) {
- decoder.close();
- }
- if (!streamingUsageRecorded.get()) {
- String sse = writer.output();
- long usageTokens = extractUsageTokensFromSse(sse);
- tokensRecorder.accept(usageTokens);
- }
});
}
+ private void finishResponse(final BodyWriter writer, final
GzipStreamDecoder decoder) {
+ if (Objects.nonNull(decoder)) {
+ decoder.close();
+ }
+ if (!streamingUsageRecorded.get()) {
+ String sse = writer.output();
+ long usageTokens = extractUsageTokensFromSse(sse);
+ tokensRecorder.accept(usageTokens);
+ }
+ }
+
private void processChunk(final byte[] processedBytes, final
BodyWriter writer) {
String chunk = new String(processedBytes, StandardCharsets.UTF_8);
for (String line : SSE_LINE_PATTERN.split(chunk)) {
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPluginTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPluginTest.java
index 277427a091..994152d7e7 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPluginTest.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPluginTest.java
@@ -23,11 +23,23 @@ import
org.apache.shenyu.infra.redis.serializer.ShenyuRedisSerializationContext;
import
org.apache.shenyu.plugin.ai.token.limiter.handler.AiTokenLimiterPluginHandler;
import org.apache.shenyu.plugin.base.utils.CacheKeyUtils;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.Mockito;
+import org.reactivestreams.Publisher;
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.data.redis.core.ReactiveValueOperations;
import org.springframework.data.redis.serializer.RedisSerializationContext;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.MediaType;
+import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
+import org.springframework.mock.web.server.MockServerWebExchange;
+import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
+import reactor.core.publisher.Sinks;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
@@ -35,8 +47,12 @@ import java.lang.reflect.Method;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
import java.util.zip.GZIPOutputStream;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -47,6 +63,53 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
*/
class AiTokenLimiterPluginTest {
+ private static final String USAGE_EVENT = "data:
{\"usage\":{\"completion_tokens\":12}}\n\n";
+
+ @ParameterizedTest
+ @CsvSource({"false, field", "true, field", "true, number"})
+ void testResponseUsageAcrossBuffers(final boolean flush, final String
boundary) {
+ byte[] content = USAGE_EVENT.getBytes(StandardCharsets.UTF_8);
+ int split = USAGE_EVENT.indexOf("field".equals(boundary) ? "tokens" :
"2");
+ byte[][] chunks = {Arrays.copyOfRange(content, 0, split),
Arrays.copyOfRange(content, split, content.length)};
+ assertResponseUsage(content, chunks, flush, false);
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testResponseUsageAcrossGzipBuffers(final boolean flush) throws
IOException {
+ byte[] compressed = compressGzip(USAGE_EVENT);
+ assertResponseUsage(compressed, splitIntoThreeChunks(compressed),
flush, true);
+ }
+
+ @Test
+ void testFlushUsageIsFinalizedAfterOuterPublisherCompletes() {
+ MockServerWebExchange exchange =
MockServerWebExchange.from(MockServerHttpRequest.get("/"));
+ ByteArrayOutputStream forwarded = new ByteArrayOutputStream();
+ exchange.getResponse().setWriteHandler(body -> body.doOnNext(buffer ->
{
+ byte[] bytes = new byte[buffer.readableByteCount()];
+ buffer.read(bytes);
+ forwarded.writeBytes(bytes);
+ DataBufferUtils.release(buffer);
+ }).then());
+ List<Long> recorded = new ArrayList<>();
+ List<Throwable> errors = new ArrayList<>();
+ AiTokenLimiterPlugin.AiStatisticServerHttpResponse response = new
AiTokenLimiterPlugin.AiStatisticServerHttpResponse(
+ exchange, exchange.getResponse(), recorded::add);
+ Sinks.Many<Publisher<? extends DataBuffer>> body =
Sinks.many().unicast().onBackpressureBuffer();
+ response.writeAndFlushWith(body.asFlux()).subscribe(unused -> { },
errors::add);
+ int split = USAGE_EVENT.indexOf("tokens");
+
body.tryEmitNext(Mono.just(exchange.getResponse().bufferFactory().wrap(USAGE_EVENT.substring(0,
split).getBytes(StandardCharsets.UTF_8)))).orThrow();
+ assertEquals(USAGE_EVENT.substring(0, split),
forwarded.toString(StandardCharsets.UTF_8));
+ assertTrue(recorded.isEmpty());
+
body.tryEmitNext(Mono.just(exchange.getResponse().bufferFactory().wrap(USAGE_EVENT.substring(split).getBytes(StandardCharsets.UTF_8)))).orThrow();
+ assertEquals(USAGE_EVENT, forwarded.toString(StandardCharsets.UTF_8));
+ assertTrue(recorded.isEmpty());
+
+ body.tryEmitComplete().orThrow();
+ assertTrue(errors.isEmpty());
+ assertEquals(List.of(12L), recorded);
+ }
+
@Test
void testBodyWriterRetainsOnlyTailBytes() {
AiTokenLimiterPlugin.BodyWriter writer = new
AiTokenLimiterPlugin.BodyWriter(8);
@@ -242,6 +305,29 @@ class AiTokenLimiterPluginTest {
assertEquals("0", valuePair.read(valuePair.write("0")));
}
+ private void assertResponseUsage(final byte[] expected, final byte[][]
chunks, final boolean flush, final boolean gzip) {
+ MockServerWebExchange exchange =
MockServerWebExchange.from(MockServerHttpRequest.get("/"));
+
exchange.getResponse().getHeaders().setContentType(MediaType.TEXT_EVENT_STREAM);
+ if (gzip) {
+
exchange.getResponse().getHeaders().set(HttpHeaders.CONTENT_ENCODING, "gzip");
+ }
+ List<Long> recorded = new ArrayList<>();
+ AiTokenLimiterPlugin.AiStatisticServerHttpResponse response = new
AiTokenLimiterPlugin.AiStatisticServerHttpResponse(
+ exchange, exchange.getResponse(), recorded::add);
+ Flux<DataBuffer> body =
Flux.fromArray(chunks).map(exchange.getResponse().bufferFactory()::wrap);
+ Mono<Void> completion = flush ?
response.writeAndFlushWith(body.map(Mono::just)) : response.writeWith(body);
+ completion.block();
+
+ assertEquals(List.of(12L), recorded);
+ byte[] forwarded =
DataBufferUtils.join(exchange.getResponse().getBody()).map(buffer -> {
+ byte[] bytes = new byte[buffer.readableByteCount()];
+ buffer.read(bytes);
+ DataBufferUtils.release(buffer);
+ return bytes;
+ }).block();
+ assertArrayEquals(expected, forwarded);
+ }
+
@SuppressWarnings("unchecked")
private ReactiveValueOperations<String, String> mockValueOperations() {
return Mockito.mock(ReactiveValueOperations.class);
