[wguo7]

What type of PR is this?

/kind bug

What this PR does / why we need it:

The P/D usage rewriter, used by the coordinator decode step and the NIXL sidecar connector, reads only a top-level usage field. The response.completed event of a streamed /v1/responses reply nests usage under response, so it kept decode's cached_tokens (#3577). observe already finds that response.usage block for the prompt-token parity check (#3751), and with this change the cached-token rewrite applies to it too. The Messages message.usage block stays observe-only. response.incomplete carries usage the same way and is covered too. Events whose response.usage is null, such as response.created, are left unchanged. The rewriter already buffers an SSE line split across writes (TestResponseWriter_BuffersStreamedFramesSplitAcrossWrites), so a large response.completed event is parsed whole.

Which issue(s) this PR fixes:

Fixes #3577

Release note (write NONE if no user-facing change):

In P/D, streamed /v1/responses replies report the prefill cached-token count in the response.completed usage.

Test plan:

• Streamed response.completed and response.incomplete events have their cached tokens rewritten (fail on main, pass here), and a response.created event with null usage passes through unchanged.
• go test passes for pkg/common/..., pkg/coordinator/... and pkg/sidecar/proxy (TestServeMetrics_ListenError skipped, since it also hangs on macOS on main), and make lint reports 0 issues. make presubmit was not run.

[blacksoxx]

Checked this at 034f9db.

setCachedTokens only looked for a top-level usage, and a streamed Responses reply carries its usage inside the response object of the response.completed and response.incomplete events, which is also where the EPP's openai parser reads it from. Recursing into response when the frame has no top-level block covers both terminal events, and a response.created frame, whose usage is null, falls out of the type assertion and is left unchanged, as the third row shows. The chat and completions paths are untouched since their frames carry usage at the top level. On the constant: "response" appears nowhere else in the request paths except as the openai parser's struct tag, so FieldResponse is its first named home.

The two terminal-event rows fail on main and pass on the branch; go test -race for pkg/common/usage, pkg/common/request and pkg/coordinator/steps, go build ./..., go vet and golangci-lint v2.14.0 are clean.

I would approve this; the decision is with the maintainers.

[wguo7]

@elevran, could you review this when you have time? It fixes the cached-token count in streamed Responses replies. It has one approval, all checks pass, and the branch still merges cleanly with main.

[github-actions]

@wguo7 This PR has merge conflicts with main. Please rebase on main and push to resolve them:

gh pr checkout 3596
git fetch origin main
git rebase origin/main
# resolve conflicts
git push --force-with-lease

[github-actions]

@blacksoxx @wguo7 has rebased this PR and it is now conflict-free. Please take another look and re-approve when ready.

[wguo7]

Rebased onto main to resolve the conflict with #3636. Since setCachedTokens now receives only the usage block, the lookup of the nested response object moved into a small helper, replaceNestedCachedTokensJSON, which replaceCachedTokensJSON calls only when a frame has no top-level usage. Chat Completions and other replies with top-level usage still take the path from #3636, and a frame with neither a top-level usage nor a response.usage is left unchanged. That is the only change to the fix, and the test cases were not edited.

CI passes, and locally the pkg/common/usage tests pass, including the ones #3636 added. The streamed responses completed event and streamed responses incomplete event cases of TestReplaceCachedTokens fail on main and pass with this change. Could you take another look when you have time?

[wguo7]

This replaces my earlier rebase note. #3751 conflicted, so I rebased again and dropped the replaceNestedCachedTokensJSON helper. #3751 added observe, which finds the usage block for its prompt-token parity check. It already looks in response.usage on Responses stream events.

The PR changes only pkg/common/usage/cached_tokens.go and its test. observe is unexported and has two callers, processJSON and processSSELine, both in that file. Its second return value changes from topLevel bool to the key of the object that holds a nested block: "" for a top-level block, "response" or "message" otherwise. The cached-token rewrite runs only for "" and "response". Dropping an injected usage-only frame still requires a top-level block, so a frame whose usage is only under response is never dropped.

message.usage on the Messages message_start event is only observed, not rewritten. The Messages API reports cached tokens as cache_read_input_tokens, and setCachedTokens leaves a block in that schema unchanged wherever it appears, so rewriting it would have no effect.

The two Responses event cases in TestReplaceCachedTokens are new in this PR. They fail against main with #3751 and pass with this change. The pkg/common/usage and pkg/coordinator tests pass, and golangci-lint reports 0 issues on pkg/common/usage.