chore: harness tests count tokens - #5648
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds Gemini and Vertex end-to-end coverage for ChangesProvider request accounting and grounding coverage
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Comment |
This stack of pull requests is managed by Graphite. Learn more about stacking. |
|
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/e2e/api/collections/provider-harness.json (1)
43154-43161: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAlign 34.5 with the safer response-parsing pattern used by 34.1–34.4.
Line 43154 dereferences
pm.response.json()directly and skips thereturns totalTokensassertion every sibling case has. Since 400 is intentionally outside the infra guard, a non-JSON error body aborts the rest of the script rather than reporting a clean failure.♻️ Proposed consistency fix
- "var total = pm.response.json().totalTokens;", + "var body = pm.response.json();", + "var total = body && body.totalTokens;", + "pm.test('GenAI Gemini countTokens returns totalTokens', function () {", + " pm.expect(total, 'totalTokens missing: ' + pm.response.text()).to.be.a('number');", + "});", "var withSystem = Number(pm.collectionVariables.get('genaiCountTokensWithSystem'));",🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/e2e/api/collections/provider-harness.json` around lines 43154 - 43161, Update the 34.5 response-handling script around the totalTokens assertion to use the same safe JSON parsing pattern as sibling cases 34.1–34.4, avoiding direct dereferencing of pm.response.json() when the 400 response is non-JSON. Restore the corresponding “returns totalTokens” assertion while preserving the existing conditional fallback token-count comparison and its collection-variable behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/e2e/api/collections/provider-harness.json`:
- Around line 43154-43161: Update the 34.5 response-handling script around the
totalTokens assertion to use the same safe JSON parsing pattern as sibling cases
34.1–34.4, avoiding direct dereferencing of pm.response.json() when the 400
response is non-JSON. Restore the corresponding “returns totalTokens” assertion
while preserving the existing conditional fallback token-count comparison and
its collection-variable behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 37952764-ad68-446c-a90e-c420546f151c
📒 Files selected for processing (1)
tests/e2e/api/collections/provider-harness.json
7579f05 to
46195fe
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
tests/e2e/api/collections/provider-harness.json (3)
44146-44149: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winBoth SSE items keep only the last chunk's
groundingMetadatainstead of merging across chunks. IfgroundingChunks,groundingSupports,webSearchQueries, andsearchEntryPointare split across different SSE chunks, the finalgmis partial — producing spurious in-bounds failures (supports withchunkCount === 0) or silently dropping the fields under test.
tests/e2e/api/collections/provider-harness.json#L44146-L44149: accumulate the grounding fields acrosssseChunksin the gemini streaming item rather than overwritinggm.tests/e2e/api/collections/provider-harness.json#L44240-L44243: apply the same accumulation in the vertex streaming item.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/e2e/api/collections/provider-harness.json` around lines 44146 - 44149, Update the Gemini streaming item at tests/e2e/api/collections/provider-harness.json:44146-44149 and the Vertex streaming item at tests/e2e/api/collections/provider-harness.json:44240-44243 so the sseChunks loop merges groundingMetadata fields into a cumulative gm instead of replacing it with the last chunk. Accumulate groundingChunks, groundingSupports, and webSearchQueries across chunks, and retain the available searchEntryPoint, while preserving the existing candidate traversal and validation behavior.
43492-43504: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFolder 35 hardcodes
gemini-3.6-flashinstead of the collection's model variables. Folder 34 and the rest of the harness parametrize the model via{{genaiModel}}/{{vertexModel}}; folder 35 pins a literal model id in every URL, so the tests break wholesale when that model is retired or isn't provisioned for a given account, and they can't be pointed at a different model per environment.
tests/e2e/api/collections/provider-harness.json#L43492-L43504: replace the literalgemini-3.6-flashwith{{vertexModel}}in bothrawandpathfor the vertexexcludeDomainsitem.tests/e2e/api/collections/provider-harness.json#L43551-L43563: replace it with{{genaiModel}}here, and apply the same substitution to the remaining non-streaming filter/latLng/snake_case items.tests/e2e/api/collections/provider-harness.json#L43942-L43954: use the provider-appropriate variable in the grounding-fidelity items (gemini and vertex).tests/e2e/api/collections/provider-harness.json#L44106-L44124: use the provider-appropriate variable in both SSE streaming items.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/e2e/api/collections/provider-harness.json` around lines 43492 - 43504, The provider harness hardcodes gemini-3.6-flash instead of using environment-configurable model variables. In tests/e2e/api/collections/provider-harness.json:43492-43504, replace it with {{vertexModel}} in both raw and path; at 43551-43563 and the remaining non-streaming filter/latLng/snake_case items, use {{genaiModel}}; at 43942-43954, use the appropriate {{genaiModel}} or {{vertexModel}} for gemini and vertex items; and at 44106-44124, apply the corresponding provider variable to both SSE streaming URLs.
43154-43155: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAlign 34.5 with the defensive
json()pattern used by every sibling case.Every other case in folder 34 does
var body = pm.response.json(); var total = body && body.totalTokens;plus atotalTokenstype test. Here a non-JSON body (plain-text gateway error) makespm.response.json()throw and aborts the script, so the fallbacks invariant reports as a Newman script error rather than a clean assertion failure.♻️ Proposed consistency fix
- "var total = pm.response.json().totalTokens;", + "var body = pm.response.json();", + "var total = body && body.totalTokens;", + "pm.test('GenAI Gemini countTokens returns totalTokens', function () {", + " pm.expect(total, 'totalTokens missing: ' + pm.response.text()).to.be.a('number');", + "});", "var withSystem = Number(pm.collectionVariables.get('genaiCountTokensWithSystem'));",🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/e2e/api/collections/provider-harness.json` around lines 43154 - 43155, Update the 34.5 response-handling script around the total variable to parse the response into a body variable, safely read totalTokens only when the body exists, and add the same totalTokens type validation used by sibling folder 34 cases. Preserve the existing fallback assertion flow so non-JSON responses produce a clean assertion failure instead of a Newman script error.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/e2e/api/collections/provider-harness.json`:
- Around line 43740-43744: The FunctionCallingConfig assertions are unreachable
for provider-rejection responses because the status guard returns first. In
tests/e2e/api/collections/provider-harness.json#L43740-L43744, move the gemini
latLng FunctionCallingConfig test before the code >= 400 guard, and apply the
identical reordering to the vertex latLng item at
tests/e2e/api/collections/provider-harness.json#L43797-L43801.
- Around line 43482-43487: Update all four Folder 35 request header blocks in
tests/e2e/api/collections/provider-harness.json at lines 43482-43487,
43541-43546, 43932-43937, and 44096-44101 to include the x-goog-api-key
authentication header, preserving the existing Content-Type header and ensuring
every request reaches its assertions.
---
Nitpick comments:
In `@tests/e2e/api/collections/provider-harness.json`:
- Around line 44146-44149: Update the Gemini streaming item at
tests/e2e/api/collections/provider-harness.json:44146-44149 and the Vertex
streaming item at tests/e2e/api/collections/provider-harness.json:44240-44243 so
the sseChunks loop merges groundingMetadata fields into a cumulative gm instead
of replacing it with the last chunk. Accumulate groundingChunks,
groundingSupports, and webSearchQueries across chunks, and retain the available
searchEntryPoint, while preserving the existing candidate traversal and
validation behavior.
- Around line 43492-43504: The provider harness hardcodes gemini-3.6-flash
instead of using environment-configurable model variables. In
tests/e2e/api/collections/provider-harness.json:43492-43504, replace it with
{{vertexModel}} in both raw and path; at 43551-43563 and the remaining
non-streaming filter/latLng/snake_case items, use {{genaiModel}}; at
43942-43954, use the appropriate {{genaiModel}} or {{vertexModel}} for gemini
and vertex items; and at 44106-44124, apply the corresponding provider variable
to both SSE streaming URLs.
- Around line 43154-43155: Update the 34.5 response-handling script around the
total variable to parse the response into a body variable, safely read
totalTokens only when the body exists, and add the same totalTokens type
validation used by sibling folder 34 cases. Preserve the existing fallback
assertion flow so non-JSON responses produce a clean assertion failure instead
of a Newman script error.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 6f6ac87d-0a50-43cd-a528-0b7c6d034177
📒 Files selected for processing (2)
core/providers/gemini/payload_ordering_test.gotests/e2e/api/collections/provider-harness.json
💤 Files with no reviewable changes (1)
- core/providers/gemini/payload_ordering_test.go
Merge activity
|
The base branch was changed.
46195fe to
2265798
Compare
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/providers/gemini/payload_ordering_test.go (1)
161-170: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winCover
ToolConfigin the ingress conversion test.The flat-body case verifies
systemInstruction,tools, andgenerationConfig, but omitstoolConfig. A regression that removes or misdecodesToolConfigfromGeminiCountTokensRequestwould therefore still pass, even though this PR explicitly targets tool-configuration preservation. Include atoolConfigfixture and assert that it survives conversion.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/providers/gemini/payload_ordering_test.go` around lines 161 - 170, Extend the “keeps every counted field of a flat body” test and its JSON fixture to include ToolConfig, then assert the converted request preserves the expected tool-configuration data. Use the existing flatten result and ToolConfig symbols, keeping the current assertions for contents, systemInstruction, tools, and generationConfig unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@core/providers/gemini/payload_ordering_test.go`:
- Around line 161-170: Extend the “keeps every counted field of a flat body”
test and its JSON fixture to include ToolConfig, then assert the converted
request preserves the expected tool-configuration data. Use the existing flatten
result and ToolConfig symbols, keeping the current assertions for contents,
systemInstruction, tools, and generationConfig unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 000ba94d-cf18-4eef-b3e6-b313e058ea8b
📒 Files selected for processing (2)
core/providers/gemini/payload_ordering_test.gotests/e2e/api/collections/provider-harness.json
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/e2e/api/collections/provider-harness.json

Summary
Fixes a regression where
:countTokensfor both Gemini (generativelanguage) and Vertex (aiplatform) reported only thecontentstoken count becausesystemInstruction,generationConfig, andtoolConfigwere unconditionally stripped from the request body. In the reported case, a ~7.4k-token system prompt plus tool declarations counted as only 17 tokens on Vertex.The root cause is that the two endpoints require opposite request shapes:
systemInstruction,tools,toolConfig, andgenerationConfigare rejected asUnknown nameat the top level and must be wrapped inside agenerateContentRequestenvelope. Top-levelcontents/modelare silently ignored when the envelope is present.toolConfig(and a few others) must be stripped. There is no envelope.The
/genaiingress must now parse both shapes and emit whichever the resolved provider requires. The Bifrost-onlyfallbacksrouting field must be honoured for routing but stripped before the upstream call.Changes
totalTokensas a collection variable for relative assertions in subsequent cases.systemInstructionis counted (not stripped) for both providers.tools,toolConfig, andgenerationConfigare wrapped into the Gemini envelope rather than forwarded verbatim (which would produce a hard 400).toolsandgenerationConfigflat but stripstoolConfig.generateContentRequestenvelope is passed through for Gemini and unwrapped for Vertex.fallbacksfield is stripped before the upstream call and does not alter the counted token total.Type of change
Affected areas
How to test
Import the updated
provider-harness.jsoncollection into Postman or Newman and run group 34 with validgenaiKey,genaiModel,vertexModel, andbaseUrlcollection variables set.Each case should return HTTP 2xx with
totalTokensstrictly greater than the baseline recorded in 34.1/34.6. Any case that returns anUnknown nameerror or a token count equal to the baseline indicates the regression has re-appeared.Breaking changes
Related issues
Closes #5620
Security considerations
No auth, secrets, or PII changes. Test requests use a scoped no-op tool declaration (
probe_context) with no real side effects.Checklist
docs/contributing/README.mdand followed the guidelines