utafrali commented on code in PR #7937:
URL: https://github.com/apache/shenyu/pull/7937#discussion_r4334419706
##########
shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java:
##########
@@ -1056,6 +1056,16 @@ 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 response object of an OpenAI Responses stream event, which carries
its usage.
Review Comment:
The constant name `RESPONSE = "response"` is too generic for a shared
constants interface that already contains `RESPONSE_MONO`,
`CLIENT_RESPONSE_ATTR`, `CLIENT_RESPONSE_CONN_ATTR`, and
`ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR`. A reader seeing `Constants.RESPONSE` at
a call site has no indication it refers to the `response` envelope object of an
OpenAI Responses stream event.
Consider a more specific name such as `RESPONSES_RESPONSE_OBJECT` or
`AI_RESPONSES_RESPONSE`, mirroring the pattern used for other AI-specific
constants (`USAGE`, `COMPLETION_TOKENS`). Alternatively, define it as a
package-private constant inside the `AiTokenLimiterPlugin` itself, since it is
only referenced there.
##########
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"
Review Comment:
The test drives the observer with bytes produced by
`ConnectProtocolAdapter`, which means a change to the adapter's output format
(field order, whitespace, added fields) will silently break this test even when
the token-limiter logic itself is unchanged. The class-level Javadoc notes the
intent, but it is worth adding an inline comment here flagging that this is the
integration point:
```java
// If this assertion breaks after adapter changes, update CHAT_COMPLETION /
CHAT_STREAM
// to match the new adapter output before debugging the token-limiter logic.
byte[] responses = converted(false);
```
This makes triage faster for contributors who see the failure without
reading the PR description.
##########
shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/main/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPlugin.java:
##########
@@ -307,11 +311,38 @@ private void processChunk(final byte[] processedBytes,
final BodyWriter writer)
writer.write(ByteBuffer.wrap(processedBytes));
}
+ /**
+ * Locate the usage object of a stream event: Chat events carry it at
the root, while Responses events
+ * carry it inside the {@code response} object of the terminal event.
+ */
+ private JsonNode usageOf(final JsonNode node) {
+ JsonNode usage = node.get(Constants.USAGE);
+ if (Objects.isNull(usage)) {
Review Comment:
The fallback path in `usageOf()` checks for a top-level `response` field
only when `usage` is absent at the root. The assumption is that a Chat
Completions event never has a top-level `response` object — which is true
today, but it is an implicit invariant. A brief comment would make the guard's
reasoning explicit and protect a future reader:
```java
// Chat events put usage at the root; Responses terminal events wrap it
inside response.usage.
// Chat Completions never uses a top-level "response" field, so the fallback
is unambiguous.
if (Objects.isNull(usage)) {
```
##########
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() {
+ // The observer parses each buffer on its own, so the whole event has
to arrive in one buffer; the fallback
+ // then accounts the buffered tail. Both paths must agree on the
Responses field.
+ 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);
Review Comment:
The comment "the fallback then accounts the buffered tail" in
`testResponsesStreamEventAcrossBuffersRecordsOutputTokens` is slightly
misleading. The test splits the raw bytes mid-`"output_tokens"` key, so neither
buffer on its own contains a complete SSE event. What actually happens is that
the streaming JSON-parse path in `processChunk` fails to find a complete event
in the first buffer and the regex fallback path (`extractUsageTokensFromSse`)
runs over the accumulated body. Rephrasing to say "neither partial buffer forms
a complete JSON event, so the regex fallback over the accumulated SSE tail is
the code path exercised here" would make the coverage intent clear.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
