Skip to content

chore: harness tests count tokens - #5648

Merged
akshaydeo merged 1 commit into
devfrom
07-29-chore_harness_tests_count_tokens
Jul 29, 2026
Merged

chore: harness tests count tokens#5648
akshaydeo merged 1 commit into
devfrom
07-29-chore_harness_tests_count_tokens

Conversation

@TejasGhatte

@TejasGhatte TejasGhatte commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Summary

Fixes a regression where :countTokens for both Gemini (generativelanguage) and Vertex (aiplatform) reported only the contents token count because systemInstruction, generationConfig, and toolConfig were 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:

  • Gemini: systemInstruction, tools, toolConfig, and generationConfig are rejected as Unknown name at the top level and must be wrapped inside a generateContentRequest envelope. Top-level contents/model are silently ignored when the envelope is present.
  • Vertex: These fields are accepted flat; only toolConfig (and a few others) must be stripped. There is no envelope.

The /genai ingress must now parse both shapes and emit whichever the resolved provider requires. The Bifrost-only fallbacks routing field must be honoured for routing but stripped before the upstream call.

Changes

  • Added E2E test suite 34. Gemini/Vertex countTokens full-prompt accounting covering:
    • 34.1 / 34.6 – Contents-only baseline for Gemini and Vertex respectively, recording totalTokens as a collection variable for relative assertions in subsequent cases.
    • 34.2 / 34.7 – Flat systemInstruction is counted (not stripped) for both providers.
    • 34.3 – Flat tools, toolConfig, and generationConfig are wrapped into the Gemini envelope rather than forwarded verbatim (which would produce a hard 400).
    • 34.8 – Vertex keeps tools and generationConfig flat but strips toolConfig.
    • 34.4 / 34.9 – A generateContentRequest envelope is passed through for Gemini and unwrapped for Vertex.
    • 34.5 – The Bifrost fallbacks field is stripped before the upstream call and does not alter the counted token total.
  • Assertions are relative (count must exceed the contents-only baseline) rather than absolute, so they remain valid across model version updates. A dropped field collapses the count back to baseline, which is exactly the regression signature these cases catch.
  • HTTP 400 is intentionally not in the infra guard for the tools/toolConfig/fallbacks cases — a 400 from upstream is itself the regression signature for those scenarios.

Type of change

  • Bug fix

Affected areas

  • Providers/Integrations

How to test

Import the updated provider-harness.json collection into Postman or Newman and run group 34 with valid genaiKey, genaiModel, vertexModel, and baseUrl collection variables set.

newman run tests/e2e/api/collections/provider-harness.json \
  --folder "34. Gemini/Vertex countTokens full-prompt accounting (PR #5620)" \
  --env-var baseUrl=<your-base-url> \
  --env-var genaiKey=<your-api-key> \
  --env-var genaiModel=gemini-2.5-pro \
  --env-var vertexModel=gemini-2.5-pro

Each case should return HTTP 2xx with totalTokens strictly greater than the baseline recorded in 34.1/34.6. Any case that returns an Unknown name error or a token count equal to the baseline indicates the regression has re-appeared.

Breaking changes

  • No

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

  • I read docs/contributing/README.md and followed the guidelines
  • I added/updated tests where appropriate
  • I updated documentation where needed
  • I verified builds succeed (Go and UI)
  • I verified the CI pipeline passes locally if applicable

@coderabbitai

coderabbitai Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Summary by CodeRabbit

  • Improvements
    • Improved token-counting consistency for Gemini and Vertex requests, including system instructions, tools, generation settings, and routing options.
    • Enhanced Google Search grounding support across Gemini and Vertex, including domain exclusions, time filters, location-aware retrieval, and snake_case request aliases.
    • Improved grounding metadata reliability for both streaming and non-streaming responses, including search queries, citations, and grounding references.

Walkthrough

Adds Gemini and Vertex end-to-end coverage for countTokens request accounting and googleSearch grounding filters, retrieval configuration, and metadata fidelity across non-streaming and streaming responses. It also relaxes one cached-token conversion assertion.

Changes

Provider request accounting and grounding coverage

Layer / File(s) Summary
countTokens request-shape coverage
tests/e2e/api/collections/provider-harness.json, core/providers/gemini/payload_ordering_test.go
Adds Gemini and Vertex baseline-relative token-count tests for prompt fields, envelopes, unsupported fields, and routing-only fallbacks; removes the cached audio-token zero assertion.
Grounding filters and retrieval configuration
tests/e2e/api/collections/provider-harness.json
Tests excludeDomains, timeRangeFilter, snake_case aliases, and retrievalConfig.latLng for Gemini and Vertex.
Grounding metadata fidelity
tests/e2e/api/collections/provider-harness.json
Validates grounding queries, citable URIs, bounded and de-duplicated chunk indices, and rendered search content in non-streaming and SSE responses.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

  • maximhq/bifrost#5619: Adds Vertex countTokens field sanitization exercised by the new harness cases.
  • maximhq/bifrost#5620: Adds the Gemini/Vertex countTokens request-wrapping and token-mapping behavior covered by these tests.

Suggested reviewers: akshaydeo, pratham-mishra04, sammaji

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Linked Issues check ⚠️ Warning This PR only adds harness tests and a minor assertion tweak; it does not implement the request-building and response-conversion work required by #5620. Add the core/transport/provider changes for Gemini and Vertex countTokens handling, or retarget the issue to reflect that this PR is test-only.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title matches the main change: adding harness tests for countTokens behavior.
Description check ✅ Passed The PR description includes the required sections and fills in the main implementation, testing, issue, and checklist details.
Out of Scope Changes check ✅ Passed The changed files stay focused on countTokens harness coverage and a related test assertion adjustment, with no clear unrelated additions.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch 07-29-chore_harness_tests_count_tokens

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 @coderabbitai help to get the list of available commands.

@TejasGhatte
TejasGhatte marked this pull request as ready for review July 29, 2026 09:02

TejasGhatte commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author

@CLAassistant

CLAassistant commented Jul 29, 2026

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@coderabbitai coderabbitai Bot 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.

🧹 Nitpick comments (1)
tests/e2e/api/collections/provider-harness.json (1)

43154-43161: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Align 34.5 with the safer response-parsing pattern used by 34.1–34.4.

Line 43154 dereferences pm.response.json() directly and skips the returns totalTokens assertion 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

📥 Commits

Reviewing files that changed from the base of the PR and between 4b9cd7b and 7579f05.

📒 Files selected for processing (1)
  • tests/e2e/api/collections/provider-harness.json

@TejasGhatte
TejasGhatte force-pushed the 07-29-chore_harness_tests_count_tokens branch from 7579f05 to 46195fe Compare July 29, 2026 10:39
@coderabbitai
coderabbitai Bot requested a review from sammaji July 29, 2026 10:40

@coderabbitai coderabbitai Bot 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.

Actionable comments posted: 2

🧹 Nitpick comments (3)
tests/e2e/api/collections/provider-harness.json (3)

44146-44149: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Both SSE items keep only the last chunk's groundingMetadata instead of merging across chunks. If groundingChunks, groundingSupports, webSearchQueries, and searchEntryPoint are split across different SSE chunks, the final gm is partial — producing spurious in-bounds failures (supports with chunkCount === 0) or silently dropping the fields under test.

  • tests/e2e/api/collections/provider-harness.json#L44146-L44149: accumulate the grounding fields across sseChunks in the gemini streaming item rather than overwriting gm.
  • 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 win

Folder 35 hardcodes gemini-3.6-flash instead 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 literal gemini-3.6-flash with {{vertexModel}} in both raw and path for the vertex excludeDomains item.
  • 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 win

Align 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 a totalTokens type test. Here a non-JSON body (plain-text gateway error) makes pm.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

📥 Commits

Reviewing files that changed from the base of the PR and between 7579f05 and 46195fe.

📒 Files selected for processing (2)
  • core/providers/gemini/payload_ordering_test.go
  • tests/e2e/api/collections/provider-harness.json
💤 Files with no reviewable changes (1)
  • core/providers/gemini/payload_ordering_test.go

Comment thread tests/e2e/api/collections/provider-harness.json
Comment thread tests/e2e/api/collections/provider-harness.json
coderabbitai[bot]
coderabbitai Bot previously approved these changes Jul 29, 2026

akshaydeo commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Merge activity

  • Jul 29, 11:07 AM UTC: A user started a stack merge that includes this pull request via Graphite.
  • Jul 29, 11:13 AM UTC: Graphite rebased this pull request as part of a merge.
  • Jul 29, 11:14 AM UTC: @akshaydeo merged this pull request with Graphite.

@akshaydeo
akshaydeo changed the base branch from 07-29-fix_image_search_for_genai_search_tool to graphite-base/5648 July 29, 2026 11:09
@akshaydeo
akshaydeo changed the base branch from graphite-base/5648 to dev July 29, 2026 11:11
@akshaydeo
akshaydeo dismissed coderabbitai[bot]’s stale review July 29, 2026 11:11

The base branch was changed.

@akshaydeo
akshaydeo force-pushed the 07-29-chore_harness_tests_count_tokens branch from 46195fe to 2265798 Compare July 29, 2026 11:12
@akshaydeo
akshaydeo merged commit 4a07c3f into dev Jul 29, 2026
12 of 14 checks passed
@akshaydeo
akshaydeo deleted the 07-29-chore_harness_tests_count_tokens branch July 29, 2026 11:14

@coderabbitai coderabbitai Bot 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.

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 win

Cover ToolConfig in the ingress conversion test.

The flat-body case verifies systemInstruction, tools, and generationConfig, but omits toolConfig. A regression that removes or misdecodes ToolConfig from GeminiCountTokensRequest would therefore still pass, even though this PR explicitly targets tool-configuration preservation. Include a toolConfig fixture 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

📥 Commits

Reviewing files that changed from the base of the PR and between 46195fe and 2265798.

📒 Files selected for processing (2)
  • core/providers/gemini/payload_ordering_test.go
  • tests/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

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