Skip to content

fix(ai): account Responses output_tokens in the AI token limiter - #7937

Merged
Aias00 merged 1 commit into
apache:masterfrom
hanzhe-one:fix/ai-token-limiter-responses-usage-v2
Oct 23, 2026
Merged

Aias00 merged 1 commit into
apache:masterfrom
hanzhe-one:fix/ai-token-limiter-responses-usage-v2

Conversation

@hanzhe-one

Copy link
Copy Markdown
Contributor

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:

  • the streaming observer (usage.completion_tokens at the event root),
  • the buffered fallback pattern,
  • and, for Responses, the usage is not even at the event root — it sits inside the response object of the terminal event.

Fix

  • Read the usage object from the event root (Chat) or from the response object (Responses), in both the streaming observer and the buffered fallback.
  • Accept completion_tokens or output_tokens as the output-token count. The fallback pattern keeps the leading quote so max_output_tokens — a request echo — is never mistaken for an accounted count.
  • Accounting policy is unchanged: the last reported output-token count is charged, and the response bytes are forwarded untouched.

Tests

  • AiTokenLimiterResponsesUsageTest (new) drives the observer with the bytes produced by the real ConnectProtocolAdapter, 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.
  • AiTokenLimiterPluginTest gains 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.

@utafrali utafrali left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)) {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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"

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
@hanzhe-one
hanzhe-one force-pushed the fix/ai-token-limiter-responses-usage-v2 branch from eb4ef55 to 916751f Compare October 21, 2026 08:20
@hanzhe-one

Copy link
Copy Markdown
Contributor Author

Thanks for the review. Three of the four points are addressed in 916751f:

  • Constants.RESPONSE — removed from the shared interface. The field is only meaningful to this observer, so it is now a private RESPONSES_EVENT_OBJECT in AiTokenLimiterPlugin, and shenyu-common is left with just OUTPUT_TOKENS, which mirrors the existing COMPLETION_TOKENS.
  • Implicit Chat invariant — made explicit at the fallback site: "Chat never uses a top-level `response` field, so the fallback cannot shadow a Chat payload."
  • Split-buffer test comment — reworded as suggested: neither buffer holds a complete event, so the streaming JSON-parse path finds nothing and the regex fallback over the accumulated SSE tail accounts the usage.

On the third point (AiTokenLimiterResponsesUsageTest being coupled to the adapter's output format): I left it as is, because the premise does not hold for this test. It does not pin the adapter's output — it feeds the adapter and asserts on whatever bytes come back. CHAT_COMPLETION / CHAT_STREAM are the upstream-side Chat fixtures the test hands to the adapter; if the adapter changes its Responses output format, the observer receives the new format directly and the assertion still exercises the behavior we care about. There is therefore no stale fixture to debug when it fails: a failure means the limiter no longer accounts the adapter's current output, which is exactly the signal we want.

Re-ran mvn -pl shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter -am test -Dtest='AiTokenLimiter*Test': 30 tests, 0 failures, checkstyle clean.

@Aias00 Aias00 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_PATTERN is what makes the broadened regex safe. "max_completion_tokens" and "max_output_tokens" both have _ before completion/output, so "\"(?:completion|output)_tokens\"" cannot match them, and "completion_tokens_details" / "output_tokens_details" are rejected by the mandatory closing quote. testResponsesUsageWithoutOutputTokensRecordsNothing is the test that actually pins this: its body has only max_output_tokens: 99, goes through the regex path (a non-SSE body never reaches processChunk, which requires a data: prefix), and asserts 0. Dropping the leading quote would make it record 99 and fail.
  • BodyWriter keeps a sliding tail window, not a prefix (write drops 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. processChunk sets streamingUsageRecorded on the first event that yields a count, and finishResponse only runs the regex fallback when the flag is still false. Both new tests assert recorded.size() == 1, which is the assertion that actually guards this.
  • A Responses stream that carries "usage": null in response.created/response.in_progress does not record a spurious 0: usageOf requires usage.isObject(), so a NullNode returns null rather than falling through to asLong().
  • Build wiring is fine. shenyu-plugin-ai-proxy is a sibling module of token-limiter under the same shenyu-plugin-ai reactor and does not depend back on it, so the test-scoped edge introduces no cycle; ConnectProtocolAdapter(String, String, ProviderProfile) and AiUpstreamResponse(int, Map, Flux<byte[]>) both match the test's call sites.
  • Constants.OUTPUT_TOKENS does not collide with an existing constant.

Non-blocking

  1. testResponsesUsageIsNotReadFromMaxOutputTokens does not actually discriminate: max_output_tokens: 99 appears before usage.output_tokens: 12, and the extractor keeps the last match, so the assertion is 12 whether or not the guard works. The guard is genuinely covered by testResponsesUsageWithoutOutputTokensRecordsNothing, so this is only about a test whose name promises more than it verifies — moving max_output_tokens after the usage object would make it a real guard test.
  2. The Javadoc on OUTPUT_TOKENS_PATTERN mentions only max_output_tokens; the same leading quote also protects max_completion_tokens, which is the Chat-side request echo and the more likely one to appear in a Chat payload.
  3. outputTokensOf returns the node without checking it is numeric, and asLong() on a NullNode or a text node yields 0. That both records a 0 count and sets streamingUsageRecorded, which suppresses the regex fallback for the rest of the response. A canConvertToLong() check that returns null would keep the fallback alive instead of silently pinning the count to zero.
  4. 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 to completion_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.
  5. Making the token-limiter module carry a test-scoped dependency on ai-proxy (which drags in connect, 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 -pl build 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.

@hanzhe-one

Copy link
Copy Markdown
Contributor Author

CI note: build (shenyu-integrated-test-http) fails in CryptorResponsePluginTest.testWhenDoesntExistFiled (Expected: is "10001" but: was null), but this is a pre-existing regression on master, not something this PR touches.

11a39fda3 (#7903) changed CryptorResponsePlugin.fieldErrorParse to throw CRYPTOR_RESPONSE_ERROR_CONFIGURATION (-116) instead of returning the original body, while the integrated test still asserts the old passthrough contract. The same job fails on the master push for 679a9eb69cb5fc61f1ae242c8f5bdd4011ead8f7 (job 119079963808), which is before this branch existed, and the last green master integration run (7be323f42) predates that commit.

This PR only touches shenyu-common constants and shenyu-plugin-ai-token-limiter; the cryptor plugin and its integrated test are unchanged. Reported as #7968 so it can be fixed on its own. Re-running this job should not change the outcome until #7968 lands.

@Aias00
Aias00 merged commit 9344b5f into apache:master Oct 23, 2026
48 of 51 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants