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 9344b5fbaf fix(ai): account Responses output_tokens in the AI token
limiter (#7937)
9344b5fbaf is described below
commit 9344b5fbafc929213b90b8d66e50febb2f5fdc14
Author: hanzhe-one <[email protected]>
AuthorDate: Fri Oct 23 11:21:58 2026 +0800
fix(ai): account Responses output_tokens in the AI token limiter (#7937)
The token limiter's response observer only recognized the Chat Completions
usage layout. ShenYu's enhanced AI proxy converts a Responses client's
traffic, and its converted bodies and terminal stream events report usage
as input_tokens/output_tokens, so a successful Responses completion was
forwarded intact but recorded as zero tokens.
Read the output tokens under either protocol: look for usage at the event
root (Chat) or inside the response object of a Responses event, accept
completion_tokens or output_tokens, and let the buffered fallback match
both names. The pattern keeps the leading quote so max_output_tokens, a
request echo, is never mistaken for a count. Accounting policy is
unchanged: the last reported output-token count is what is charged, and
the response bytes are still forwarded untouched.
Verified with the actual ConnectProtocolAdapter output: the new
integration test records 12 tokens for both the Chat and the converted
Responses body and stream, and fails on the previous code with
"expected: <12> but was: <0>".
Fixes #7911
---
.../apache/shenyu/common/constant/Constants.java | 5 +
.../shenyu-plugin-ai-token-limiter/pom.xml | 7 ++
.../ai/token/limiter/AiTokenLimiterPlugin.java | 57 +++++++--
.../ai/token/limiter/AiTokenLimiterPluginTest.java | 59 +++++++++-
.../limiter/AiTokenLimiterResponsesUsageTest.java | 127 +++++++++++++++++++++
5 files changed, 246 insertions(+), 9 deletions(-)
diff --git
a/shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java
b/shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java
index 82716875f4..f9c024fcd3 100644
---
a/shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java
+++
b/shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java
@@ -1056,6 +1056,11 @@ public interface Constants {
*/
String COMPLETION_TOKENS = "completion_tokens";
+ /**
+ * The output tokens reported by the OpenAI Responses protocol, the
counterpart of COMPLETION_TOKENS.
+ */
+ String OUTPUT_TOKENS = "output_tokens";
+
/**
* The constant AI_MODEL.
*/
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/pom.xml
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/pom.xml
index dd1bfe3d34..e1a96254e7 100644
--- a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/pom.xml
+++ b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/pom.xml
@@ -32,6 +32,13 @@
<artifactId>shenyu-infra-redis</artifactId>
<version>${project.version}</version>
</dependency>
+ <!-- The observer is covered against the bytes the AI proxy actually
produces. -->
+ <dependency>
+ <groupId>org.apache.shenyu</groupId>
+ <artifactId>shenyu-plugin-ai-proxy</artifactId>
+ <version>${project.version}</version>
+ <scope>test</scope>
+ </dependency>
</dependencies>
</project>
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 3631130958..f12eeefc57 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
@@ -71,7 +71,17 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
private static final String REDIS_KEY_PREFIX = "SHENYU:AI:TOKENLIMIT:";
- private static final Pattern COMPLETION_TOKENS_PATTERN =
Pattern.compile("\"completion_tokens\"\\s*:\\s*(\\d+)");
+ /**
+ * The envelope of an OpenAI Responses terminal event, which carries the
usage of a response. Kept local because a
+ * bare {@code response} name in the shared {@link Constants} would read
as an HTTP response.
+ */
+ private static final String RESPONSES_EVENT_OBJECT = "response";
+
+ /**
+ * Output tokens of the Chat and Responses protocols. The leading quote
keeps {@code max_output_tokens}, a request
+ * echo, from being read as an accounted count.
+ */
+ private static final Pattern OUTPUT_TOKENS_PATTERN =
Pattern.compile("\"(?:completion|output)_tokens\"\\s*:\\s*(\\d+)");
private static final Pattern SSE_LINE_PATTERN = Pattern.compile("\\r?\\n");
@@ -294,10 +304,10 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
}
try {
JsonNode node = MAPPER.readTree(payload);
- JsonNode usage = node.get(Constants.USAGE);
- if (Objects.nonNull(usage) &&
usage.has(Constants.COMPLETION_TOKENS)) {
- long c =
usage.get(Constants.COMPLETION_TOKENS).asLong();
- tokensRecorder.accept(c);
+ JsonNode usage = usageOf(node);
+ JsonNode outputTokens = Objects.isNull(usage) ? null :
outputTokensOf(usage);
+ if (Objects.nonNull(outputTokens)) {
+ tokensRecorder.accept(outputTokens.asLong());
streamingUsageRecorded.set(true);
}
} catch (Exception e) {
@@ -307,11 +317,42 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
writer.write(ByteBuffer.wrap(processedBytes));
}
+ /**
+ * Locate the usage object of a payload: Chat payloads carry it at the
root, while a Responses terminal event
+ * carries it inside its {@code response} object.
+ *
+ * @param node parsed payload of a stream event or a complete body
+ * @return the usage object, or null when the payload reports none
+ */
+ private JsonNode usageOf(final JsonNode node) {
+ JsonNode usage = node.get(Constants.USAGE);
+ if (Objects.isNull(usage)) {
+ // Chat never uses a top-level "response" field, so the
fallback cannot shadow a Chat payload.
+ JsonNode responsesEvent = node.get(RESPONSES_EVENT_OBJECT);
+ usage = Objects.isNull(responsesEvent) ? null :
responsesEvent.get(Constants.USAGE);
+ }
+ return Objects.nonNull(usage) && usage.isObject() ? usage : null;
+ }
+
+ /**
+ * Read the output tokens of a usage object under either protocol's
field name.
+ *
+ * @param usage usage object of a Chat or Responses payload
+ * @return the reported output tokens, or null when the object has
neither field
+ */
+ private JsonNode outputTokensOf(final JsonNode usage) {
+ if (usage.has(Constants.COMPLETION_TOKENS)) {
+ return usage.get(Constants.COMPLETION_TOKENS);
+ }
+ return usage.has(Constants.OUTPUT_TOKENS) ?
usage.get(Constants.OUTPUT_TOKENS) : null;
+ }
+
private long extractUsageTokensFromSse(final String sse) {
- Matcher m = COMPLETION_TOKENS_PATTERN.matcher(sse);
+ // The usage of a body is reported once, at its end, so the last
match of either field is the count.
+ Matcher matcher = OUTPUT_TOKENS_PATTERN.matcher(sse);
long last = 0L;
- while (m.find()) {
- last = Long.parseLong(m.group(1));
+ while (matcher.find()) {
+ last = Long.parseLong(matcher.group(1));
}
return last;
}
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 994152d7e7..1477e58dc1 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
@@ -65,6 +65,53 @@ class AiTokenLimiterPluginTest {
private static final String USAGE_EVENT = "data:
{\"usage\":{\"completion_tokens\":12}}\n\n";
+ private static final String RESPONSES_COMPLETED_EVENT = "event:
response.completed\n"
+ + "data:
{\"type\":\"response.completed\",\"sequence_number\":4,\"response\":{\"id\":\"resp_1\",\"object\":\"response\","
+ +
"\"status\":\"completed\",\"usage\":{\"input_tokens\":3,\"output_tokens\":12,\"total_tokens\":15}}}\n\n";
+
+ private static final String RESPONSES_BODY =
"{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"completed\","
+ +
"\"output\":[],\"usage\":{\"input_tokens\":3,\"output_tokens\":12,\"total_tokens\":15}}";
+
+ @Test
+ void testResponsesStreamEventRecordsOutputTokens() {
+ // Issue #7911: a Responses terminal event carries its usage inside
the response object.
+ assertResponseUsage(RESPONSES_COMPLETED_EVENT, new byte[][]
{RESPONSES_COMPLETED_EVENT.getBytes(StandardCharsets.UTF_8)},
+ false, false, 12L);
+ }
+
+ @Test
+ void testResponsesStreamEventAcrossBuffersRecordsOutputTokens() {
+ // Split inside the output_tokens key: neither buffer holds a complete
event, so the streaming JSON-parse path
+ // finds nothing and the regex fallback over the accumulated SSE tail
is what accounts the usage.
+ byte[] content =
RESPONSES_COMPLETED_EVENT.getBytes(StandardCharsets.UTF_8);
+ int split = RESPONSES_COMPLETED_EVENT.indexOf("\"output_tokens\"");
+ byte[][] chunks = {Arrays.copyOfRange(content, 0, split),
Arrays.copyOfRange(content, split, content.length)};
+ for (boolean flush : new boolean[] {false, true}) {
+ assertResponseUsage(RESPONSES_COMPLETED_EVENT, chunks, flush,
false, 12L);
+ }
+ }
+
+ @Test
+ void testCompleteResponsesBodyRecordsOutputTokens() {
+ // A converted, non-streaming Responses body is not SSE framed: only
the buffered fallback sees it.
+ byte[] content = RESPONSES_BODY.getBytes(StandardCharsets.UTF_8);
+ assertResponseUsage(RESPONSES_BODY, new byte[][] {content}, false,
false, 12L);
+ }
+
+ @Test
+ void testResponsesUsageIsNotReadFromMaxOutputTokens() {
+ // max_output_tokens is an echo of the request limit, never an
accounted count.
+ String body =
"{\"id\":\"resp_1\",\"object\":\"response\",\"max_output_tokens\":99,\"usage\":{\"input_tokens\":3,"
+ + "\"output_tokens\":12,\"total_tokens\":15}}";
+ assertResponseUsage(body, new byte[][]
{body.getBytes(StandardCharsets.UTF_8)}, false, false, 12L);
+ }
+
+ @Test
+ void testResponsesUsageWithoutOutputTokensRecordsNothing() {
+ String body =
"{\"id\":\"resp_1\",\"object\":\"response\",\"status\":\"incomplete\",\"max_output_tokens\":99}";
+ assertResponseUsage(body, new byte[][]
{body.getBytes(StandardCharsets.UTF_8)}, false, false, 0L);
+ }
+
@ParameterizedTest
@CsvSource({"false, field", "true, field", "true, number"})
void testResponseUsageAcrossBuffers(final boolean flush, final String
boundary) {
@@ -306,6 +353,16 @@ class AiTokenLimiterPluginTest {
}
private void assertResponseUsage(final byte[] expected, final byte[][]
chunks, final boolean flush, final boolean gzip) {
+ assertResponseUsage(expected, chunks, flush, gzip, 12L);
+ }
+
+ private void assertResponseUsage(final String expected, final byte[][]
chunks, final boolean flush, final boolean gzip,
+ final long expectedTokens) {
+ assertResponseUsage(expected.getBytes(StandardCharsets.UTF_8), chunks,
flush, gzip, expectedTokens);
+ }
+
+ private void assertResponseUsage(final byte[] expected, final byte[][]
chunks, final boolean flush, final boolean gzip,
+ final long expectedTokens) {
MockServerWebExchange exchange =
MockServerWebExchange.from(MockServerHttpRequest.get("/"));
exchange.getResponse().getHeaders().setContentType(MediaType.TEXT_EVENT_STREAM);
if (gzip) {
@@ -318,7 +375,7 @@ class AiTokenLimiterPluginTest {
Mono<Void> completion = flush ?
response.writeAndFlushWith(body.map(Mono::just)) : response.writeWith(body);
completion.block();
- assertEquals(List.of(12L), recorded);
+ assertEquals(List.of(expectedTokens), recorded);
byte[] forwarded =
DataBufferUtils.join(exchange.getResponse().getBody()).map(buffer -> {
byte[] bytes = new byte[buffer.readableByteCount()];
buffer.read(bytes);
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterResponsesUsageTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterResponsesUsageTest.java
new file mode 100644
index 0000000000..4e35e59a57
--- /dev/null
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterResponsesUsageTest.java
@@ -0,0 +1,127 @@
+/*
+ * 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.apache.shenyu.plugin.ai.proxy.enhanced.model.AiUpstreamResponse;
+import
org.apache.shenyu.plugin.ai.proxy.enhanced.protocol.ConnectProtocolAdapter;
+import org.apache.shenyu.plugin.ai.proxy.enhanced.protocol.OpenAiResponses;
+import org.junit.jupiter.api.Test;
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferUtils;
+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 java.io.ByteArrayOutputStream;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * Usage accounting on the exact bytes ShenYu's AI proxy converts to, for both
client protocols.
+ */
+class AiTokenLimiterResponsesUsageTest {
+
+ private static final Duration TIMEOUT = Duration.ofSeconds(5);
+
+ private static final String CHAT_COMPLETION =
"{\"id\":\"chatcmpl-1\",\"object\":\"chat.completion\",\"model\":\"model\","
+ +
"\"choices\":[{\"index\":0,\"message\":{\"role\":\"assistant\",\"content\":\"hello\"},\"finish_reason\":\"stop\"}],"
+ +
"\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":12,\"total_tokens\":15}}";
+
+ private static final String CHAT_STREAM = "data:
{\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"model\":\"model\","
+ +
"\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hello\"},\"finish_reason\":\"stop\"}]}\n\n"
+ + "data:
{\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"model\":\"model\",\"choices\":[],"
+ +
"\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":12,\"total_tokens\":15}}\n\n"
+ + "data: [DONE]\n\n";
+
+ @Test
+ void recordsTheSameUsageForChatAndConvertedResponsesBodies() {
+ // Issue #7911: the Responses body is produced by the real adapter,
not a hand-written fixture.
+ byte[] responses = converted(false);
+ assertEquals(12L, recordedUsage(false,
CHAT_COMPLETION.getBytes(StandardCharsets.UTF_8)));
+ assertEquals(12L, recordedUsage(false, responses));
+ }
+
+ @Test
+ void recordsTheSameUsageForChatAndConvertedResponsesStreams() {
+ byte[] responses = converted(true);
+ assertEquals(12L, recordedUsage(true,
CHAT_STREAM.getBytes(StandardCharsets.UTF_8)));
+ assertEquals(12L, recordedUsage(true, responses));
+ }
+
+ @Test
+ void forwardsTheConvertedBytesUnchanged() {
+ for (boolean stream : new boolean[] {false, true}) {
+ byte[] responses = converted(stream);
+ assertArrayEquals(responses, forwarded(stream, responses));
+ }
+ }
+
+ private byte[] converted(final boolean stream) {
+ ConnectProtocolAdapter adapter = new
ConnectProtocolAdapter(OpenAiResponses.NAME, null, null);
+ AiUpstreamResponse upstream = new AiUpstreamResponse(200,
+ Map.of("Content-Type", List.of(stream ?
MediaType.TEXT_EVENT_STREAM_VALUE : MediaType.APPLICATION_JSON_VALUE)),
+ Flux.just((stream ? CHAT_STREAM :
CHAT_COMPLETION).getBytes(StandardCharsets.UTF_8)));
+ return adapter.convertResponse(upstream, stream).getBody()
+ .reduce(new ByteArrayOutputStream(), (buffer, bytes) -> {
+ buffer.writeBytes(bytes);
+ return buffer;
+ })
+ .map(ByteArrayOutputStream::toByteArray)
+ .block(TIMEOUT);
+ }
+
+ private long recordedUsage(final boolean stream, final byte[] body) {
+ List<Long> recorded = new ArrayList<>();
+ MockServerWebExchange exchange = exchange(stream);
+ AiTokenLimiterPlugin.AiStatisticServerHttpResponse observer = new
AiTokenLimiterPlugin.AiStatisticServerHttpResponse(
+ exchange, exchange.getResponse(), recorded::add);
+ observer.writeWith(Flux.just(buffer(exchange, body))).block(TIMEOUT);
+ assertEquals(1, recorded.size());
+ return recorded.get(0);
+ }
+
+ private byte[] forwarded(final boolean stream, final byte[] body) {
+ MockServerWebExchange exchange = exchange(stream);
+ AiTokenLimiterPlugin.AiStatisticServerHttpResponse observer = new
AiTokenLimiterPlugin.AiStatisticServerHttpResponse(
+ exchange, exchange.getResponse(), tokens -> { });
+ observer.writeWith(Flux.just(buffer(exchange, body))).block(TIMEOUT);
+ return
DataBufferUtils.join(exchange.getResponse().getBody()).map(buffer -> {
+ byte[] bytes = new byte[buffer.readableByteCount()];
+ buffer.read(bytes);
+ DataBufferUtils.release(buffer);
+ return bytes;
+ }).block(TIMEOUT);
+ }
+
+ private MockServerWebExchange exchange(final boolean stream) {
+ MockServerWebExchange exchange =
MockServerWebExchange.from(MockServerHttpRequest.post("/v1/responses"));
+ exchange.getResponse().getHeaders().setContentType(stream ?
MediaType.TEXT_EVENT_STREAM : MediaType.APPLICATION_JSON);
+ return exchange;
+ }
+
+ private DataBuffer buffer(final MockServerWebExchange exchange, final
byte[] body) {
+ return exchange.getResponse().bufferFactory().wrap(body);
+ }
+}
