Repository navigation
fix(ai): account Responses output_tokens in the AI token limiter - #7937
Conversation
utafrali
left a comment
There was a problem hiding this comment.
The fix is correct: the regex guard (leading quote) reliably blocks max_output_tokens, usageOf() and outputTokensOf() cover both protocols cleanly, and the test suite exercises streaming, split-buffer, complete-body, false-positive, and no-usage cases. The main concern is that RESPONSE = "response" added to the shared Constants interface is an ambiguous name in a file full of HTTP-response-named symbols; it should be scoped to the plugin or given a more specific name.
| String OUTPUT_TOKENS = "output_tokens"; | ||
|
|
||
| /** | ||
| * The response object of an OpenAI Responses stream event, which carries its usage. |
There was a problem hiding this 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.
| */ | ||
| private JsonNode usageOf(final JsonNode node) { | ||
| JsonNode usage = node.get(Constants.USAGE); | ||
| if (Objects.isNull(usage)) { |
There was a problem hiding this 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:
// 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)) {| 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" |
There was a problem hiding this 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:
// 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.
| 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); |
There was a problem hiding this 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.
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 apache#7911
eb4ef55 to
916751f
Compare
|
Thanks for the review. Three of the four points are addressed in 916751f:
On the third point ( Re-ran |
Aias00
left a comment
There was a problem hiding this comment.
LGTM. The fix targets the right layer and the way it is tested is better than what I usually see in this module: the new AiTokenLimiterResponsesUsageTest derives its fixtures from the real ConnectProtocolAdapter conversion instead of hand-writing a Responses body, so if the adapter's output shape drifts the accounting test fails with it.
What I checked
- The leading quote in
OUTPUT_TOKENS_PATTERNis what makes the broadened regex safe."max_completion_tokens"and"max_output_tokens"both have_beforecompletion/output, so"\"(?:completion|output)_tokens\""cannot match them, and"completion_tokens_details"/"output_tokens_details"are rejected by the mandatory closing quote.testResponsesUsageWithoutOutputTokensRecordsNothingis the test that actually pins this: its body has onlymax_output_tokens: 99, goes through the regex path (a non-SSE body never reachesprocessChunk, which requires adata:prefix), and asserts0. Dropping the leading quote would make it record99and fail. BodyWriterkeeps a sliding tail window, not a prefix (writedrops the oldest bytes on overflow,output()returns the last ≤64 KB). That is what makes "the last match is the count" sound for large bodies — the usage sits at the end, so it is never the part that gets evicted.- No double counting.
processChunksetsstreamingUsageRecordedon the first event that yields a count, andfinishResponseonly runs the regex fallback when the flag is still false. Both new tests assertrecorded.size() == 1, which is the assertion that actually guards this. - A Responses stream that carries
"usage": nullinresponse.created/response.in_progressdoes not record a spurious 0:usageOfrequiresusage.isObject(), so aNullNodereturns null rather than falling through toasLong(). - Build wiring is fine.
shenyu-plugin-ai-proxyis a sibling module oftoken-limiterunder the sameshenyu-plugin-aireactor and does not depend back on it, so the test-scoped edge introduces no cycle;ConnectProtocolAdapter(String, String, ProviderProfile)andAiUpstreamResponse(int, Map, Flux<byte[]>)both match the test's call sites. Constants.OUTPUT_TOKENSdoes not collide with an existing constant.
Non-blocking
testResponsesUsageIsNotReadFromMaxOutputTokensdoes not actually discriminate:max_output_tokens: 99appears beforeusage.output_tokens: 12, and the extractor keeps the last match, so the assertion is12whether or not the guard works. The guard is genuinely covered bytestResponsesUsageWithoutOutputTokensRecordsNothing, so this is only about a test whose name promises more than it verifies — movingmax_output_tokensafter theusageobject would make it a real guard test.- The Javadoc on
OUTPUT_TOKENS_PATTERNmentions onlymax_output_tokens; the same leading quote also protectsmax_completion_tokens, which is the Chat-side request echo and the more likely one to appear in a Chat payload. outputTokensOfreturns the node without checking it is numeric, andasLong()on aNullNodeor a text node yields0. That both records a0count and setsstreamingUsageRecorded, which suppresses the regex fallback for the rest of the response. AcanConvertToLong()check that returns null would keep the fallback alive instead of silently pinning the count to zero.- The regex fallback runs over the raw text for every non-SSE body, so it can match inside generated content or a tool-call argument that literally contains
"output_tokens": N. The class of problem predates this PR (it applied tocompletion_tokens), and the last-match rule makes it harmless for the standard field order, but since the pattern now covers two field names it may be worth parsing the buffered tail as JSON first and only falling back to the regex when that fails. - Making the token-limiter module carry a test-scoped dependency on
ai-proxy(which drags inconnect, okhttp and reactor-netty at test scope) is a real cost. I think it is the right trade-off here — a hand-written fixture would have hidden exactly the bug this PR fixes — but it does mean this module's unit tests can no longer be built in isolation from an-plbuild without-am.
Note on CI: build (shenyu-integrated-test-http) is currently red on master because of the CryptorResponsePluginTest.testWhenDoesntExistFiled assertion left behind by #7903 — unrelated to this PR.
|
CI note:
This PR only touches |
Problem
The AI token limiter only recognized the Chat Completions usage layout. The enhanced AI proxy supports OpenAI Responses clients and converts their traffic, but the converted bodies and terminal stream events report usage as
input_tokens/output_tokens, so the observer recorded zero tokens for a successful Responses completion while the same completion was charged 12 tokens in Chat format (issue #7911).Three places read only
completion_tokens:usage.completion_tokensat the event root),responseobject of the terminal event.Fix
responseobject (Responses), in both the streaming observer and the buffered fallback.completion_tokensoroutput_tokensas the output-token count. The fallback pattern keeps the leading quote somax_output_tokens— a request echo — is never mistaken for an accounted count.Tests
AiTokenLimiterResponsesUsageTest(new) drives the observer with the bytes produced by the realConnectProtocolAdapter, not hand-written fixtures, and asserts that Chat and converted Responses bodies/streams record the same 12 tokens and that the converted bytes are forwarded unchanged.AiTokenLimiterPluginTestgains Responses stream (single buffer and split across buffers, buffered and flush paths), complete-body,max_output_tokens-is-not-a-count, and no-usage cases.Both new integration tests fail on the previous code with
expected: <12> but was: <0>and pass with the fix.mvn -pl shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter -am test -Dtest='AiTokenLimiter*Test': 30 tests, 0 failures, checkstyle clean.