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 dc078f08f8 fix(ai): bound token usage response buffer (#7124)
dc078f08f8 is described below
commit dc078f08f879a8717fb92ba457648ec90af6dff4
Author: Liming Deng <[email protected]>
AuthorDate: Sun Sep 20 13:01:02 2026 +0800
fix(ai): bound token usage response buffer (#7124)
Co-authored-by: aias00 <[email protected]>
---
.../ai/token/limiter/AiTokenLimiterPlugin.java | 67 +++++++++++-----------
.../ai/token/limiter/AiTokenLimiterPluginTest.java | 54 +++++++++++++++++
2 files changed, 89 insertions(+), 32 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 0765ed9695..74ae710783 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
@@ -52,10 +52,7 @@ import reactor.core.publisher.Mono;
import reactor.util.annotation.NonNull;
import java.io.ByteArrayOutputStream;
-import java.io.IOException;
import java.nio.ByteBuffer;
-import java.nio.channels.Channels;
-import java.nio.channels.WritableByteChannel;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.Objects;
@@ -348,46 +345,52 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
static class BodyWriter {
- private final ByteArrayOutputStream stream = new
ByteArrayOutputStream();
+ private static final int MAX_BUFFER_SIZE = 64 * 1024;
- private final WritableByteChannel channel =
Channels.newChannel(stream);
+ private final byte[] bytes;
+
+ private int size;
private final AtomicBoolean isClosed = new AtomicBoolean(false);
- void write(final ByteBuffer buffer) {
- if (!isClosed.get()) {
- try {
- channel.write(buffer);
- } catch (IOException e) {
- isClosed.compareAndSet(false, true);
- LOG.error("Parse Failed.", e);
- }
+ BodyWriter() {
+ this(MAX_BUFFER_SIZE);
+ }
+
+ BodyWriter(final int maxBufferSize) {
+ if (maxBufferSize <= 0) {
+ throw new IllegalArgumentException("maxBufferSize must be
greater than zero");
+ }
+ this.bytes = new byte[maxBufferSize];
+ }
+
+ void write(final ByteBuffer source) {
+ if (isClosed.get() || !source.hasRemaining()) {
+ return;
+ }
+ int incomingSize = source.remaining();
+ if (incomingSize >= bytes.length) {
+ source.position(source.limit() - bytes.length);
+ source.get(bytes);
+ size = bytes.length;
+ return;
+ }
+ int overflow = Math.max(0, size + incomingSize - bytes.length);
+ if (overflow > 0) {
+ System.arraycopy(bytes, overflow, bytes, 0, size - overflow);
+ size -= overflow;
}
+ source.get(bytes, size, incomingSize);
+ size += incomingSize;
}
boolean isEmpty() {
- return stream.size() == 0;
+ return size == 0;
}
String output() {
- try {
- isClosed.compareAndSet(false, true);
- return stream.toString(StandardCharsets.UTF_8);
- } catch (Exception e) {
- LOG.error("Write failed: ", e);
- return "Write failed: " + e.getMessage();
- } finally {
- try {
- stream.close();
- } catch (IOException e) {
- LOG.error("Close stream error: ", e);
- }
- try {
- channel.close();
- } catch (IOException e) {
- LOG.error("Close channel error: ", e);
- }
- }
+ isClosed.compareAndSet(false, true);
+ return new String(bytes, 0, size, StandardCharsets.UTF_8);
}
}
}
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
new file mode 100644
index 0000000000..d929206d80
--- /dev/null
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPluginTest.java
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.shenyu.plugin.ai.token.limiter;
+
+import org.junit.jupiter.api.Test;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Test cases for {@link AiTokenLimiterPlugin}.
+ */
+class AiTokenLimiterPluginTest {
+
+ @Test
+ void testBodyWriterRetainsOnlyTailBytes() {
+ AiTokenLimiterPlugin.BodyWriter writer = new
AiTokenLimiterPlugin.BodyWriter(8);
+
+ assertTrue(writer.isEmpty());
+
writer.write(ByteBuffer.wrap("012345".getBytes(StandardCharsets.UTF_8)));
+ writer.write(ByteBuffer.wrap("6789".getBytes(StandardCharsets.UTF_8)));
+
+ assertFalse(writer.isEmpty());
+ assertEquals("23456789", writer.output());
+ }
+
+ @Test
+ void testBodyWriterKeepsTailOfLargeChunk() {
+ AiTokenLimiterPlugin.BodyWriter writer = new
AiTokenLimiterPlugin.BodyWriter(8);
+
+
writer.write(ByteBuffer.wrap("0123456789".getBytes(StandardCharsets.UTF_8)));
+
+ assertEquals("23456789", writer.output());
+ }
+}
