Conversation
SnapshotInference was called only from the outbound recorders — forwardproxy's request/response sites and extproc's recordOutbound* pair. No inbound recorder called it, so neither reverseproxy nor extproc's recordInbound* pair carried token counts, and cost/usage reads every figure it reports off SessionEvent.Inference (usage.go's foldInto). A turn served through the reverse proxy therefore reached /v1/usage with a whole, correctly priced cost and no counts at all: the cost record is published into Extensions.Custom by settle.Publish and snapshotted by every recorder via SnapshotPlugins, so the money arrives while the counts — which live only on the inference extension — do not. pricedRequests == priceableRequests, empty unpricedBy and empty incompleteBy all report the figure whole, and it is; the tokens are simply absent. Unreported traffic reads exactly like free traffic. The omission is not one listener forgetting a field. The design assumed inference traffic is egress, so a reverse proxy in front of a model endpoint — inbound inference, which is all reverseproxy serves — was never covered. Fill in the five inbound recording sites. Snapshot rather than the live pointer for the reason the helper already documents: response-phase assignments to the token fields would otherwise appear on the request event. Guard it in the parity suite, which was green throughout because observation compares the cost record (through PluginEventJSON, the route that works) and never compared Inference — the same root cause as #936. A pairwise check alone would not have caught this: both inbound listeners shared the gap, so they agreed with each other while reporting nothing. So observation gains Inference for the split case, and cost_parity_test.go gains an absolute per-fixture token expectation for the shared one. Before this change that expectation fails on four fixtures across both inbound listeners and passes on every outbound one. Deliberately unchanged: the SessionDenied recorders, since no listener records protocol extensions on a deny and a denied request has no counts; and the recording gates, so nothing newly appears in the session stream — this only fills a field on events that were already recorded. Fixes #1165 Signed-off-by: Rong Chang <rong@us.ibm.com>
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (6)
Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthroughInbound reverse proxy and ext_proc session events now include inference snapshots. Parity observations and fixture assertions now check inference token reports alongside cost records. ChangesInbound inference reporting
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Bug fix · Severity of issue fixed: Medium Suggested reviewers: Merge Risk: ⚪ Minimal · up to The supplied evidence identifies no actionable issue with inbound token reporting or its parity checks; the PR appears mergeable after normal checks. Security Architecture ReviewSecurity architecture risk: 🟡 Moderate · up to Inbound model requests can now place prompts and completions in session records, not just token counts. The session endpoint relies on restricted network access rather than authentication. The actual deployment exposure is unconfirmed. Retained concerns
Security review detailsSecurity Blast Radius
Security Findings and Attack Paths
Trust Boundaries and Controls
Resilience and Maintainability Implications
Hardening Proposals
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
huang195
left a comment
huang195
left a comment
There was a problem hiding this comment.
The production change is correct and safe — I verified the mechanism end to end, including the two things that could have gone wrong and didn't: usage.go:1043 returns before foldInto for SessionRequest, so the new request-phase snapshots cannot double-count in /v1/usage; and modifyResponse returns at :488 after installing the streaming body, so the buffered and onClose recorders never both fire. foldInto really does take every token figure off e.Inference and nothing else, so the diagnosis holds.
One must-fix, in the guard rather than the fix. assertTokens runs only on l.run(t, cf.fixture, pipeline.SessionResponse), so of the five sites this PR adds, the two request-phase ones are covered by observationDiff's pairwise check alone — which is exactly the check parity_test.go:439 documents as unable to catch a gap a direction's whole listener set shares. Measured: dropping Inference: from both inbound request recorders (deletion proved by git diff --stat) leaves ./listener/parity/ at EXIT:0, and core/listener/extproc and core/listener/reverseproxy's own packages pass too.
Mutation gate: 10 mutated, 9 killed, 1 survived. Every response-phase site is genuinely pinned, in both directions, and the wantTokens: nil pins on the /v1/embeddings fixtures are live rather than vacuous. Red-state re-derivation matches the body exactly — 4 inbound fixtures, 8 carried NO token report lines, 0 outbound failures, 0 pairwise Inference: diffs. gofmt/vet/tidy and 60/60 core packages confirmed as stated.
Also: snapshot.go:77's "EVERY RECORDER MUST CALL THIS" is absolute while 5 of the 14 non-test recorders deliberately don't, and the body's "both noted in the code" isn't true of the four *Reject ones — the reason lives only in the PR description, which is the one place the next recorder author won't read.
| // And the counts the figure was computed from, on the same event. Absolute | ||
| // for the reason above, sharpened: the listeners of ONE DIRECTION shared | ||
| // this gap, so they agreed with each other while reporting nothing. | ||
| assertTokens(t, l.name, obs.Inference, cf.wantTokens) |
There was a problem hiding this comment.
must-fix — this absolute check runs on one phase only, so two of the five sites this PR adds are unguarded against the exact shape it was written for.
obs comes from l.run(t, cf.fixture, pipeline.SessionResponse) at :352, so the three response-phase additions are pinned and the two request-phase ones — extproc/server.go:360 and reverseproxy/server.go:440 — are covered only by observationDiff's new pairwise comparison. parity_test.go:439 says of that check: "a gap shared by a direction's whole listener set passes here … The absolute expectation in cost_parity_test.go covers the shared case." True for the response event; not for the request event.
Measured, with the deletion proved before the run:
removed Inference: from extproc.recordInboundSession AND reverseproxy.handleRequest
core/listener/extproc/server.go | 1 -
core/listener/reverseproxy/server.go | 1 -
GOWORK=off go test -count=1 ./listener/parity/ EXIT:0 ** SURVIVED **
Nothing else covers them: with all five sites reverted, core/listener/extproc and core/listener/reverseproxy's own packages still pass — only parity fails.
Either remedy works. (a) extend this loop with a request-phase pass — the fixtures already carry wantTokens, and a request event should assert Model and PresentKinds present with the counts still zero, which is what SnapshotInference's original doc comment promises. (b) narrow the claim at parity_test.go:439 to "…covers the shared case for the response event" and say the request half rides on pairwise alone. (a) is worth more, since (b) leaves the new rule in snapshot.go unenforceable for two of its own call sites — but (b) is honest and is not a blocker by itself.
| // extension during OnResponse; without snapshotting, the request event's | ||
| // view would contain the eventual response's token counts and completion. | ||
| // | ||
| // EVERY RECORDER MUST CALL THIS, IN EITHER DIRECTION. It is the only route by which |
There was a problem hiding this comment.
suggestion — the rule is absolute and five recorders in this tree are exempt, four of them silently.
Enumerating every non-test pipeline.SessionEvent{ site: nine call SnapshotInference; forwardproxy.recordTunnelOpened and the four *Reject recorders do not. The tunnel one already carries the right kind of note in place ("MCP/Inference snapshots are nil by definition (the bytes are opaque)"). The four deny recorders carry none — I grepped each for infer|token|count|extension|protocol|snapshot.
The PR body says they are "Deliberately unchanged, both noted in the code"; the deny half isn't — that reason exists only in the description. It matters more than usual here because this comment's stated purpose is "so the next recorder cannot repeat it by simply not thinking of it", and someone who checks the rule against recordInboundReject finds a contradiction with nothing in the tree to resolve it.
Cheapest fix is to carry the body's own sentence up here: except the SessionDenied recorders and recordTunnelOpened, which record no protocol extensions at all.
|
|
||
| // anthropicTokens is the count report for the turn every fixture below sends, buffered or | ||
| // streamed: the same tokens the rate table turns into modelledWholeUSD. Shared so that a | ||
| // fixture changing its counters without changing its expectation is a compile-time edit here |
There was a problem hiding this comment.
nit — "a compile-time edit here" is a runtime failure. The counters live in JSON string literals, so nothing couples them to anthropicTokens at compile time. Verified by changing one:
cache_read_input_tokens 200 -> 250
go vet ./listener/parity/ EXIT:0 (compiles fine)
go test ... EXIT:1 "CacheReadTokens = 250, want 200"
Sharing the var is still the right call and the failure is loud — it is just a failing expectation, not a compile error. "one failing expectation here" would be accurate.
| // but the status code and plugin invocations are always meaningful. | ||
| plugins := pipeline.SnapshotPlugins(pctx.Extensions.Custom) | ||
| // Always pair every inbound request with a response row (carries StatusCode). | ||
| // Inference is the token report, and it is what a cost consumer reads: without it |
There was a problem hiding this comment.
nit — this note sits above the if s.Sessions != nil guard, twelve lines from the Inference: field it explains. The sibling note in recordInboundResponseEvent (:743) sits directly above its Append, which reads better.
3 participants
Fixes #1165.
The problem
pipeline.SnapshotInferencewas called only from the outbound recorders —forwardproxy'srequest/response sites and
extproc'srecordOutbound*pair. No inbound recorder called it, soreverseproxyandextproc.recordInbound*published session events with no token counts, andcost/usagereads every figure it reports offSessionEvent.Inference(foldInto).A turn served through the reverse proxy therefore reaches
/v1/usagewith a whole, correctlypriced cost and zero tokens: the cost record is published into
Extensions.Custombysettle.Publishand snapshotted by every recorder throughSnapshotPlugins— so the moneyarrives — while the counts live only on the inference extension and do not.
pricedRequests == priceableRequests, emptyunpricedBy, emptyincompleteBy: every completeness signal says thefigure is whole, and it is. The measurement that is missing is the one nothing reports on.
It is not one listener forgetting a field. The design assumed inference traffic is egress, so a
reverse proxy in front of a model endpoint — inbound inference, which is all
reverseproxyserves — was never covered.
forwardproxyrequest / responseextproc.recordOutbound{,Response}Sessionextproc.recordInbound{,Response}Sessionreverseproxyrequest / buffered / streamingonCloseThe solution
Five inbound recording sites gain
Inference: pipeline.SnapshotInference(pctx.Extensions.Inference). A snapshot rather than thelive pointer for the reason the helper already documents: response-phase assignments to the
token fields would otherwise appear on the already-appended request event.
SnapshotInference's doc comment gains the normative rule — every recorder must call it, ineither direction — and why the omission does not look like one, so the next recorder cannot
repeat it by simply not thinking of it.
Then the guard. The parity suite was green throughout because
observationcompares the costrecord (through
PluginEventJSON, the route that works) and never comparedInference— thesame root cause #936 diagnoses for itself. A pairwise check alone could not have caught this:
both inbound listeners shared the gap, so they agreed with each other while reporting
nothing. So:
observationgainsInference(flattened to the fields a cost or usage consumer reads,PresentKindsincluded, since that is what separates "used no cache" from "reported nocache"), which covers the case where listeners diverge;
cost_parity_test.gogains an absolute per-fixture token expectation, which covers the casewhere they share a gap. The two
/v1/embeddingsfixtures assertwantTokens: nilexplicitly— a believed cost with no counts behind it is the one shape where money without tokens is
honest, and saying so is what stops the new check from being satisfied by an absence.
Deliberately unchanged, both noted in the code:
SessionDeniedrecorders — no listener records protocol extensions on a deny, and adenied request has no counts;
Customexactly as before, so nothing newly appears in the session stream. This change onlyfills in a field on events that were already recorded. (
cost_parity_test.go's own headerdocuments a test that depends on the current gate divergence.)
Worth recording for consumers, because it is the tempting workaround for the released behaviour:
do not divide a cost by its rate to recover the counts. That yields a restatement of the rate
table rather than a measurement — it agrees with itself even on a mispriced turn, and nothing
downstream can distinguish a derived count from an observed one.
Testing
Red first, on the unmodified recorders, with the new expectation in place:
Four inference fixtures × both inbound listeners fail; every
outbound/…subtest passes; andno pairwise
Inference:diff appears anywhere — the suite demonstrating in its own terms whythe pairwise check could not have found this. Green after the fix.
The expectations are pinned from what the known-good outbound listeners actually record, not
composed by hand.
Then, in
core/withGOWORK=off:go test ./...— all packages passgo vet ./...— cleangofmt -l .— no new entries (five pre-existing offenders unchanged, confirmed by stashing)go mod tidy -diff— cleango build ./cmd/authbridge-proxy/...against the modified core — okFound by a third-party harness metering agent spend through the reverse proxy on v0.7.0. Its
probe printed "the cost is whole: $0.084471 for None token(s)" and reported success — the
arithmetic in #1165 shows the counts existed at pricing time and only the snapshot was missing.
🤖 Generated with Claude Code
Summary by CodeRabbit