fix: add graceful fallback for unsupported count_tokens

[0xPixelNinja]

Summary

Add a graceful fallback for count_tokens when the selected provider returns unsupported_operation.

This avoids failing outright for unsupported providers by returning a best-effort response.input_tokens estimate derived from the normalized Responses request, while still preserving configured provider fallbacks and explicit operation allowlists.

Changes

• add a core fallback path for unsupported count_tokens requests
• preserve configured cross-provider fallbacks instead of short-circuiting them
• preserve explicit AllowedRequests.count_tokens = false behavior
• estimate modality-aware token details for text, image, and audio inputs
• add regression tests for unsupported providers, fallback preservation, disallowed operations, and mixed-modality inputs
• update Azure integration expectations to match the new fallback behavior

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

cd core
go test .
go test . -run "TestHandleProviderRequest_CountTokens|TestBuildGracefulCountTokensFallbackResponse|TestHandleProviderRequest_OCROperationNotAllowed"

No new configs or environment variables were added.

Screenshots/Recordings

Breaking changes

• Yes
• No

If yes, describe impact and migration instructions.

Related issues

Relates to #2902

Security considerations

No new auth, secret, or sandboxing behavior was introduced.

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

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
8 out of 12 committers have signed the CLA.

✅ Madhuvod
✅ G-XD
✅ jeremym-tanium
✅ R-droid101
✅ akshaydeo
✅ 0xPixelNinja
✅ AdityaPainuli
✅ impoiler
❌ danpiths
❌ Pratham-Mishra04
❌ TejasGhatte
❌ roroghost17
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

• New Features
  • Added token-count estimation fallback when a provider does not support native token counting.
  • Added per-request controls through the compatibility header, including wildcard and individual feature overrides.
  • Added a “Count Tokens Fallback” setting in the compatibility configuration and UI; it is disabled by default.
  • Fallback estimates account for text, images, audio, tools, and other request content.
• Bug Fixes
  • Preserves configured provider fallbacks and respects providers that explicitly disallow token-count estimation.

Walkthrough

Adds an opt-in CountTokensFallback compat flag, persists it, estimates tokens from request content, synthesizes responses for unsupported operations, and wires plugin hooks, HTTP overrides, UI, schema, and tests.

Changes

Count Tokens Graceful Fallback

Layer / File(s)	Summary
Configuration schema and persistence layer core/schemas/bifrost.go, framework/configstore/..., transports/config.schema.json, ui/lib/types/config.ts, ui/app/workspace/config/views/compatibilityView.tsx	Adds the opt-in flag, database storage and migration, config hash support, schema definitions, and UI switch.
Token estimation from request structure plugins/compat/counttokensfallback.go	Estimates text, image, file, and audio tokens across request messages, tools, instructions, and content blocks.
Fallback decision and response building plugins/compat/counttokensfallback.go	Checks provider restrictions, errors, fallback state, and attempt order before building a BifrostCountTokensResponse.
Plugin configuration and hook wiring plugins/compat/main.go, transports/bifrost-http/server/plugins.go, transports/bifrost-http/handlers/config.go, plugins/compat/hooks_test.go	Enables the feature in compat configuration, passes the account to Init, captures request state, synthesizes responses, and forwards reload settings.
Per-request context override support core/schemas/bifrost.go, transports/bifrost-http/lib/ctx.go	Parses x-bf-compat values for count_tokens_fallback and the wildcard override.
Fallback behavior validation plugins/compat/counttokensfallback_test.go	Tests unsupported-operation recovery, fallback ordering, provider disallowance, and mixed-content token accounting.

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

Sequence Diagram(s)

sequenceDiagram
  participant Client
  participant HTTP as Bifrost HTTP
  participant PreLLMHook
  participant Provider
  participant PostLLMHook

  Client->>HTTP: CountTokensRequest with optional x-bf-compat
  HTTP->>PreLLMHook: Request and context
  PreLLMHook->>Provider: Forward request
  Provider->>PostLLMHook: unsupported_operation error
  PostLLMHook->>PostLLMHook: Estimate tokens and build response
  PostLLMHook->>Client: CountTokensResponse

Loading

Possibly related PRs

• maximhq/bifrost#4027: Adds provider-specific count-tokens handling related to this fallback flow.

Suggested reviewers: akshaydeo, danpiths, pratham-mishra04

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 19.15% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: graceful fallback handling for unsupported count_tokens requests.
Description check	✅ Passed	The description covers the purpose, changes, tests, breaking changes, security, and related issue, but it omits some affected areas and verification details.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

core/count_tokens_fallback.go (1)

137-143: ⚡ Quick win

Honor ResponsesMessageContent source exclusivity in token estimation.

When ContentStr is non-nil, ContentBlocks should not be counted in the same message. Using both can over-estimate tokens and diverges from the schema contract.

♻️ Proposed fix

 	if msg.Content != nil {
 		if msg.Content.ContentStr != nil {
 			estimate = estimate.withText(estimateTokensFromText(*msg.Content.ContentStr))
-		}
-		for _, block := range msg.Content.ContentBlocks {
-			estimate = estimate.add(estimateCountTokensFromContentBlock(block))
+		} else {
+			for _, block := range msg.Content.ContentBlocks {
+				estimate = estimate.add(estimateCountTokensFromContentBlock(block))
+			}
 		}
 	}

Based on learnings: schemas.ResponsesMessageContent treats ContentStr and ContentBlocks as mutually exclusive; only use ContentBlocks when ContentStr is nil.

🤖 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/count_tokens_fallback.go` around lines 137 - 143, The token estimator is
double-counting when both ResponsesMessageContent.ContentStr and ContentBlocks
are present; update the logic in core/count_tokens_fallback.go (the branch that
inspects msg.Content) to treat ResponsesMessageContent as exclusive: if
msg.Content.ContentStr != nil, only call
estimate.withText(estimateTokensFromText(*msg.Content.ContentStr)) and skip
iterating ContentBlocks; otherwise (ContentStr == nil) iterate
msg.Content.ContentBlocks and add
estimate.add(estimateCountTokensFromContentBlock(block)) as before so
ContentBlocks are only counted when ContentStr is nil.

🤖 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/integrations/python/config.yml`:
- Line 324: provider_scenarios.azure has count_tokens enabled but
providers.azure lacks a corresponding model mapping; add a count_tokens entry
under the providers.azure models mapping so the scenario resolves to a concrete
Azure model. Locate the providers.azure configuration block and add a key named
count_tokens mapped to the appropriate Azure model identifier (the same model
type used for token counting elsewhere in the repo or tests), ensuring the
mapping name matches provider_scenarios.azure.count_tokens so runtime resolution
will succeed.

---

Nitpick comments:
In `@core/count_tokens_fallback.go`:
- Around line 137-143: The token estimator is double-counting when both
ResponsesMessageContent.ContentStr and ContentBlocks are present; update the
logic in core/count_tokens_fallback.go (the branch that inspects msg.Content) to
treat ResponsesMessageContent as exclusive: if msg.Content.ContentStr != nil,
only call estimate.withText(estimateTokensFromText(*msg.Content.ContentStr)) and
skip iterating ContentBlocks; otherwise (ContentStr == nil) iterate
msg.Content.ContentBlocks and add
estimate.add(estimateCountTokensFromContentBlock(block)) as before so
ContentBlocks are only counted when ContentStr is nil.

🪄 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: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 13d5a499-8b38-44bb-9b66-b8f77b81385a

📥 Commits

Reviewing files that changed from the base of the PR and between 96dbf26 and 3f4c861.

📒 Files selected for processing (4)

• core/bifrost.go
• core/bifrost_test.go
• core/count_tokens_fallback.go
• tests/integrations/python/config.yml

[greptile-apps]

Confidence Score: 5/5

Safe to merge — the fallback is strictly opt-in, the estimation path is isolated to a new file, and the existing error-passthrough behavior is fully preserved when the feature is disabled.

All changed paths are additive: a new DB column with a safe false default, a new estimation function that only runs when explicitly enabled, and hooks that short-circuit immediately when the state key is absent. The DB migration is idempotent, the fallback only fires after all configured provider fallbacks are exhausted, and the test suite covers the four critical behavioral contracts.

No files require special attention.

Important Files Changed

Filename	Overview
plugins/compat/counttokensfallback.go	New file implementing modality-aware token estimation logic for the graceful fallback path; estimation helpers are well-structured with clear value-object chaining.
plugins/compat/main.go	Adds CountTokensFallback flag to Config, threads it through Init/PreLLMHook/PostLLMHook, and exposes account for per-provider allowlist checks; logic is consistent with existing compat feature patterns.
plugins/compat/counttokensfallback_test.go	New tests covering unsupported-provider recovery, fallback-preservation ordering, explicit disallow respect, and mixed-modality token detail tracking; good coverage of the critical paths.
framework/configstore/migrations.go	Adds migrationAddCompatCountTokensFallbackColumn with idempotent HasColumn guard and correct false-default UPDATE; rollback drops the column safely.
transports/bifrost-http/lib/ctx.go	Adds count_tokens_fallback to all three x-bf-compat header paths (true, [*], named feature array); consistent with existing compat feature handling.

_{Reviews (5): Last reviewed commit: "fix: move count_tokens fallback to compa..." | Re-trigger Greptile}

[coderabbitai]

🧹 Nitpick comments (1)

core/count_tokens_fallback.go (1)

1-1: ⚡ Quick win

Rename this file to match the Go filename convention.

core/count_tokens_fallback.go uses underscores; the repository rule for non-test Go files requires concatenated lowercase words.

As per coding guidelines, "No underscores in Go filenames except for _test.go suffix; concatenate words in lowercase for multi-word filenames."

🤖 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/count_tokens_fallback.go` at line 1, The file name uses underscores
which violates the repo Go filename convention; rename
core/count_tokens_fallback.go to a concatenated lowercase name like
core/counttokensfallback.go and update any references or build/import paths if
the filename is referenced elsewhere (ensure package bifrost and the
function/identifier names such as CountTokensFallback remain unchanged); run go
build or go test to verify the package compiles after the rename.

🤖 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 `@core/count_tokens_fallback.go`:
- Line 1: The file name uses underscores which violates the repo Go filename
convention; rename core/count_tokens_fallback.go to a concatenated lowercase
name like core/counttokensfallback.go and update any references or build/import
paths if the filename is referenced elsewhere (ensure package bifrost and the
function/identifier names such as CountTokensFallback remain unchanged); run go
build or go test to verify the package compiles after the rename.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 2d961abf-cee5-4c3f-b826-048f88a97777

📥 Commits

Reviewing files that changed from the base of the PR and between 3f4c861 and 2d5f76e.

📒 Files selected for processing (4)

• core/bifrost.go
• core/bifrost_test.go
• core/count_tokens_fallback.go
• tests/integrations/python/config.yml

[coderabbitai]

♻️ Duplicate comments (1)

tests/integrations/python/config.yml (1)

324-324: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Azure count_tokens enablement is incomplete without a provider model mapping.

provider_scenarios.azure.count_tokens is now enabled, but providers.azure still has no count_tokens entry. This can break capability resolution or route incorrectly at runtime.

Suggested patch

   azure:
     chat: "gpt-4o"
     vision: "gpt-4o"
     tools: "gpt-4o-mini"
     streaming: "gpt-4o-mini"
     speech: "gpt-4o-mini-tts"
     transcription: "whisper"
     embeddings: "text-embedding-3-small"
     image_generation: "gpt-image-1"
     thinking: "o1"
+    count_tokens: "gpt-4o-mini"
     batch_file_upload: "gpt-4o-2"

🤖 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/integrations/python/config.yml` at line 324,
provider_scenarios.azure.count_tokens was enabled but providers.azure lacks a
corresponding count_tokens mapping; update the configuration by adding a
count_tokens entry under providers.azure that maps the Azure model(s) used for
token counting (matching the keys referenced by provider_scenarios.azure) so
capability resolution can find the provider implementation (look for
provider_scenarios.azure.count_tokens and providers.azure to add the appropriate
count_tokens mapping).

🤖 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.

Duplicate comments:
In `@tests/integrations/python/config.yml`:
- Line 324: provider_scenarios.azure.count_tokens was enabled but
providers.azure lacks a corresponding count_tokens mapping; update the
configuration by adding a count_tokens entry under providers.azure that maps the
Azure model(s) used for token counting (matching the keys referenced by
provider_scenarios.azure) so capability resolution can find the provider
implementation (look for provider_scenarios.azure.count_tokens and
providers.azure to add the appropriate count_tokens mapping).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 5bc1b251-4163-489d-bdb5-704cafc2e9ae

📥 Commits

Reviewing files that changed from the base of the PR and between 2d5f76e and cb1c9aa.

📒 Files selected for processing (4)

• core/bifrost.go
• core/bifrost_test.go
• core/counttokensfallback.go
• tests/integrations/python/config.yml

[0xPixelNinja]

Pushed the cleanup and follow up fixes, please merge if this looks good

The merge-base changed after approval.

[Pratham-Mishra04]

Hey @0xPixelNinja thanks for the PR! This feature is useful but I think we should move it to compat plugin under a toggle cause not everyone would want this behavior

[0xPixelNinja]

Hey @0xPixelNinja thanks for the PR! This feature is useful but I think we should move it to compat plugin under a toggle cause not everyone would want this behavior

Sure, I will move it to the compat plugin behind a toggle :)

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

plugins/compat/counttokensfallback_test.go (1)

30-38: ⚡ Quick win

Add explicit tests for disabled config and request-level override.

Line 33 hard-codes CountTokensFallback: true, so this suite never validates the off-by-default path or request override path in this stack. Please add one test with CountTokensFallback: false (expect original unsupported_operation to remain) and one with per-request override enabled (expect synthesized response).

Proposed minimal refactor to enable both branches in tests

-func newCompatPluginForCountTokensFallback(t *testing.T, account schemas.Account) *CompatPlugin {
+func newCompatPluginForCountTokensFallback(t *testing.T, account schemas.Account, enabled bool) *CompatPlugin {
 	t.Helper()

-	plugin, err := Init(Config{CountTokensFallback: true}, bifrost.NewNoOpLogger(), nil, account)
+	plugin, err := Init(Config{CountTokensFallback: enabled}, bifrost.NewNoOpLogger(), nil, account)
 	if err != nil {
 		t.Fatalf("init compat plugin: %v", err)
 	}
 	return plugin
}

- plugin := newCompatPluginForCountTokensFallback(t, nil)
+ plugin := newCompatPluginForCountTokensFallback(t, nil, true)

🤖 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 `@plugins/compat/counttokensfallback_test.go` around lines 30 - 38, The test
helper newCompatPluginForCountTokensFallback always sets
Config.CountTokensFallback=true so tests never exercise the disabled-default and
request-level override paths; update tests by (A) adding a test that constructs
the CompatPlugin via Init with Config{CountTokensFallback:false} and asserts the
original unsupported_operation behavior remains, and (B) adding a test that
constructs the plugin with CountTokensFallback:false but sends a request with
the per-request override enabled (use whatever request field/flag the code
inspects for request-level fallback) and asserts the synthesized response is
returned; to implement this you can either add a new helper that accepts a
countTokensFallback bool or overload newCompatPluginForCountTokensFallback to
accept that flag, and locate initialization logic in Init/Config and request
handling code that checks the per-request override to craft the assertions.

🤖 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 `@plugins/compat/counttokensfallback.go`:
- Around line 170-172: The current conditional only counts custom tool input
when both msg.ResponsesToolMessage and msg.ResponsesCustomToolCall are non-nil,
which can undercount; change the check to simply if msg.ResponsesCustomToolCall
!= nil { estimate =
estimate.withText(estimateTokensFromText(msg.ResponsesCustomToolCall.Input)) }
so that estimateTokensFromText is invoked whenever a ResponsesCustomToolCall
exists (retain the existing variables: msg.ResponsesCustomToolCall, estimate,
estimate.withText, and estimateTokensFromText).

---

Nitpick comments:
In `@plugins/compat/counttokensfallback_test.go`:
- Around line 30-38: The test helper newCompatPluginForCountTokensFallback
always sets Config.CountTokensFallback=true so tests never exercise the
disabled-default and request-level override paths; update tests by (A) adding a
test that constructs the CompatPlugin via Init with
Config{CountTokensFallback:false} and asserts the original unsupported_operation
behavior remains, and (B) adding a test that constructs the plugin with
CountTokensFallback:false but sends a request with the per-request override
enabled (use whatever request field/flag the code inspects for request-level
fallback) and asserts the synthesized response is returned; to implement this
you can either add a new helper that accepts a countTokensFallback bool or
overload newCompatPluginForCountTokensFallback to accept that flag, and locate
initialization logic in Init/Config and request handling code that checks the
per-request override to craft the assertions.

🪄 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: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 2eb47717-cded-47b2-84ec-bd01c635c2cf

📥 Commits

Reviewing files that changed from the base of the PR and between cb1c9aa and a99ce71.

📒 Files selected for processing (12)

• core/schemas/bifrost.go
• framework/configstore/clientconfig.go
• framework/configstore/migrations.go
• framework/configstore/rdb.go
• framework/configstore/tables/clientconfig.go
• plugins/compat/counttokensfallback.go
• plugins/compat/counttokensfallback_test.go
• plugins/compat/main.go
• transports/bifrost-http/handlers/config.go
• transports/bifrost-http/lib/ctx.go
• transports/bifrost-http/server/plugins.go
• transports/config.schema.json

✅ Files skipped from review due to trivial changes (1)

• core/schemas/bifrost.go

[0xPixelNinja]

@Pratham-Mishra04, moved it to the compat plugin please review when you get a chance

The merge-base changed after approval.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Actionable comments posted: 1

🤖 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 `@plugins/compat/main.go`:
- Around line 228-229: The shared typed context key is missing, causing the
compat override path to fail type checking. Add
BifrostContextKeyCompatCountTokensFallback to the typed context-key declarations
in core/schemas/bifrost.go, then ensure plugins/compat/main.go lines 228-229
reads that declared key and transports/bifrost-http/lib/ctx.go lines 582-611
clears and sets the same key.

🪄 Autofix

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: 95b94831-4cbd-47b4-bf98-3c6fb2d94a3b

📥 Commits

Reviewing files that changed from the base of the PR and between d704af4 and 0c4830b.

📒 Files selected for processing (15)

• core/schemas/bifrost.go
• framework/configstore/clientconfig.go
• framework/configstore/migrations.go
• framework/configstore/rdb.go
• framework/configstore/tables/clientconfig.go
• plugins/compat/counttokensfallback.go
• plugins/compat/counttokensfallback_test.go
• plugins/compat/hooks_test.go
• plugins/compat/main.go
• transports/bifrost-http/handlers/config.go
• transports/bifrost-http/lib/ctx.go
• transports/bifrost-http/server/plugins.go
• transports/config.schema.json
• ui/app/workspace/config/views/compatibilityView.tsx
• ui/lib/types/config.ts

🚧 Files skipped from review as they are similar to previous changes (11)

• transports/bifrost-http/server/plugins.go
• core/schemas/bifrost.go
• framework/configstore/rdb.go
• transports/config.schema.json
• ui/lib/types/config.ts
• ui/app/workspace/config/views/compatibilityView.tsx
• framework/configstore/tables/clientconfig.go
• framework/configstore/clientconfig.go
• transports/bifrost-http/handlers/config.go
• framework/configstore/migrations.go
• plugins/compat/counttokensfallback_test.go

[0xPixelNinja]

hi @akshaydeo, rebased onto latest dev and resolved the conflicts, CLA signed now too. Should be good to merge whenever you get a chance

The merge-base changed after approval.