feat: update routing engine tracking to support multiple engines#1635
feat: update routing engine tracking to support multiple engines#1635Pratham-Mishra04 wants to merge 1 commit intomainfrom
Conversation
🧪 Test Suite AvailableThis PR can be tested by a repository admin. |
📝 WalkthroughSummary by CodeRabbitRelease Notes
WalkthroughThe pull request converts single routing engine tracking to multi-engine support throughout the system. Context keys are renamed, type changed from string to array, database migrations add new columns, and UI/API layer updated to handle multiple routing engines. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/governance/main.go (1)
516-520:⚠️ Potential issue | 🟠 MajorReplace raw context key
"model"with a typed key.
ctx.SetValue("model", ...)at lines 519 and 655 uses a string literal key, violating the context-key rule. A typed keyBifrostContextKeyModeldoes not yet exist incore/schemas/bifrost.go—it must be added first as aBifrostContextKeyconstant, then used consistently.Update required at:
plugins/governance/main.golines 519 and 655 (SetValue calls)transports/bifrost-http/integrations/genai.golines 413 and 487 (Value retrieval calls)Define the key in
core/schemas/bifrost.go:BifrostContextKeyModel BifrostContextKey = "bifrost-model"Then replace all raw
"model"usages withschemas.BifrostContextKeyModel.
🤖 Fix all issues with AI agents
In `@framework/logstore/migrations.go`:
- Around line 1091-1126: The AddColumn call used for the old name
("routing_engine_used") fails on fresh installs because GORM requires a
corresponding struct field; update the migration to avoid relying on a
non-existent struct field: either add a backward-compatible field to the Log
struct (e.g., a string field tagged `gorm:"column:routing_engine_used"`
alongside the existing RoutingEnginesUsedStr) so AddColumn(&Log{},
"routing_engine_used") can succeed, or change the migration that currently calls
migrator.AddColumn(&Log{}, "routing_engine_used") to execute raw SQL (e.g.,
tx.Exec("ALTER TABLE logs ADD COLUMN routing_engine_used VARCHAR(...)")) instead
of using the GORM struct-based AddColumn; update the
migrationAddRoutingEnginesUsedJSONColumn flow to match whichever approach you
choose.
In `@framework/logstore/rdb.go`:
- Around line 68-78: Skip empty strings in filters.RoutingEngineUsed and change
the LIKE predicate to match whole comma-separated tokens by wrapping the stored
column with commas; e.g. for each non-empty engine push "CONCAT(',',
routing_engines_used, ',') LIKE ?" into engineConditions and "%,"+engine+",%"
into engineArgs before calling baseQuery.Where (retain
engineConditions/engineArgs variables and the baseQuery.Where call). This avoids
matching partial substrings and prevents empty-filter behavior.
In `@plugins/governance/main.go`:
- Around line 533-534: The build fails because bifrost.AppendToContextList is
not defined; locate where AppendToContextList actually resides (e.g., in the
core package or another package) and either export and add it to the bifrost
package or change the import/receiver to the correct package that defines
AppendToContextList; update the usages in main.go that reference
bifrost.AppendToContextList (and the duplicate calls around the other
occurrence) to call the exported function from the correct package, ensuring the
package is imported so the symbol resolves at compile time (also confirm
schemas.BifrostContextKeyRoutingEnginesUsed is reachable from the same package).
In `@ui/app/workspace/logs/views/logDetailsSheet.tsx`:
- Around line 308-319: The mapped Badge elements use key={engine}, which can
collide if log.routing_engines_used contains duplicates; update the mapping in
the LogEntryDetailsView render to use a stable composite key or dedupe the list
first: either map with both value and index (e.g., use `${engine}-${idx}`) when
iterating log.routing_engines_used in the Badge rendering, or replace the source
with a deduped array (e.g., Array.from(new Set(log.routing_engines_used)))
before mapping; ensure you update the map call that references
RoutingEngineUsedColors, RoutingEngineUsedIcons, and RoutingEngineUsedLabels so
the keys remain unique and stable.
🧹 Nitpick comments (4)
transports/changelog.md (1)
1-1: Minor grammar fix: hyphenate "multi-level".The compound adjective "multi level" should be hyphenated when used as a modifier.
📝 Suggested fix
-- feat: added multi level routing support for routing rules + vk based provider routing +- feat: added multi-level routing support for routing rules + vk-based provider routingplugins/governance/changelog.md (1)
1-1: Minor grammar fix: hyphenate "multi-level".Same as transports/changelog.md — compound adjectives should be hyphenated.
📝 Suggested fix
-- feat: added multi level routing support for routing rules + vk based provider routing +- feat: added multi-level routing support for routing rules + vk-based provider routingframework/logstore/migrations.go (1)
1095-1095: Consider renaming the migration ID for clarity.The migration ID
logs_add_routing_engines_used_json_columnsuggests adding a JSON column, but the migration actually renames an existing column. A more accurate ID likelogs_rename_routing_engine_to_engines_usedwould better reflect the operation.plugins/telemetry/main.go (1)
380-385: Consider de‑duplicating routing engines before labeling.If the context list can contain repeats, label values like
"governance,governance"add noise. A small de‑dupe keeps labels cleaner while preserving order.♻️ Suggested de-duplication before join
- // Get routing engines array and join into comma-separated string - routingEngines := []string{} - if engines, ok := ctx.Value(schemas.BifrostContextKeyRoutingEnginesUsed).([]string); ok { - routingEngines = engines - } - routingEngineUsed := strings.Join(routingEngines, ",") + // Get routing engines array and join into comma-separated string + routingEngines := []string{} + if engines, ok := ctx.Value(schemas.BifrostContextKeyRoutingEnginesUsed).([]string); ok { + seen := make(map[string]struct{}, len(engines)) + for _, engine := range engines { + if engine == "" { + continue + } + if _, ok := seen[engine]; ok { + continue + } + seen[engine] = struct{}{} + routingEngines = append(routingEngines, engine) + } + } + routingEngineUsed := strings.Join(routingEngines, ",")
Pratham-Mishra04
force-pushed
the
02-13-feat_added_multi_routing_engine_support
branch
from
65a466d to
6806b74
Compare
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/governance/main.go (1)
516-520:⚠️ Potential issue | 🟠 MajorUse a typed context key instead of raw "model" strings.
Lines 519 and 655 use raw string keys
ctx.SetValue("model", ...), which violates the no-raw-context-keys rule. However, the suggestedschemas.BifrostContextKeyModeldoes not exist.Define a typed key following the existing pattern in this file (e.g.,
modelContextKey schemas.BifrostContextKey = "bifrost-model"), or addBifrostContextKeyModeltocore/schemas/bifrost.goif it should be shared across plugins. Update both lines to use the typed key.
🤖 Fix all issues with AI agents
In `@framework/logstore/rdb.go`:
- Around line 68-85: The routing engine filter currently builds SQL using
CONCAT(',', routing_engines_used, ',') which fails on SQLite; make the filter
dialect-aware by selecting the correct concatenation expression based on the DB
dialect (use "CONCAT(',', routing_engines_used, ',')" for MySQL/Postgres and
"','" || routing_engines_used || ',' (SQLite) or equivalent), then use that
expression when constructing engineConditions and passing to baseQuery.Where;
detect the dialect the same way other code in this package does (e.g., via the
DB dialector/name helper used elsewhere) and update the logic that builds
engineConditions, engineArgs, and the call to baseQuery.Where to use the chosen
concat expression for routing_engines_used and keep the existing LIKE
"%,engine,%" pattern.
In `@plugins/governance/changelog.md`:
- Line 1: Update the changelog line "feat: added multi level routing support for
routing rules + vk based provider routing" to hyphenate compound modifiers:
change "multi level routing" to "multi-level routing" and "vk based provider
routing" to "vk-based provider routing" so the final line reads e.g. "feat:
added multi-level routing support for routing rules + vk-based provider
routing".
In `@plugins/logging/operations.go`:
- Around line 471-492: The query filter in GetAvailableRoutingEngines
incorrectly checks for a JSON empty array; update the call to p.store.FindAll to
remove the "!= '[]'" condition so it only filters "routing_engines_used IS NOT
NULL" (since routing_engines_used is stored as comma-separated values or NULL).
Keep the rest of GetAvailableRoutingEngines unchanged (parsing
log.RoutingEnginesUsed and deduplicating) and ensure p.store.FindAll is invoked
with the simplified WHERE clause referencing routing_engines_used.
In `@transports/changelog.md`:
- Line 1: Update the changelog line to hyphenate compound modifiers: change
"multi level routing support" to "multi-level routing support" and "vk based
provider routing" to "vk-based provider routing" in the existing changelog entry
(the line containing "feat: added multi level routing support for routing rules
+ vk based provider routing").
🧹 Nitpick comments (1)
plugins/governance/main.go (1)
487-490: Consider tracking the TODO as a follow‑up task.
Right now the code silently continues when all providers are filtered out; if that’s intentional, a small follow‑up issue to return a clear error/telemetry signal would help avoid blind spots.Would you like me to draft a follow‑up issue or a minimal error response patch?
Pratham-Mishra04
force-pushed
the
02-13-feat_added_multi_routing_engine_support
branch
from
6806b74 to
e7890aa
Compare
core/utils.go
Outdated
| if !ok { | ||
| existingValues = []T{} | ||
| } | ||
| ctx.SetValue(key, append(existingValues, value)) |
There was a problem hiding this comment.
@Pratham-Mishra04 instead of this we can just introduce this method in BifrostContext itself.
ctx.AppendValue() which internally can ensure all this
There was a problem hiding this comment.
yes can do
There was a problem hiding this comment.
Cannot set generics to an interface so will have to do datatype specific function like ctx.AppendStringValue() - is that okay?
| p.logger.Debug("[Governance] Allowed providers after filtering: %v", allowedProviders) | ||
|
|
||
| if len(allowedProviderConfigs) == 0 { | ||
| // TODO: Send proper error if (overall VK budget/rate limit) or (all provider budgets/rate limits) are violated |
There was a problem hiding this comment.
we can start sending already right?
There was a problem hiding this comment.
yes but will need to brainstorming first cause early returns directly (from HTTPPreHook) will present request from entering plugin flow, affecting error logging and loadbalancing
Pratham-Mishra04
force-pushed
the
02-13-feat_added_multi_routing_engine_support
branch
from
e7890aa to
9360f73
Compare
Pratham-Mishra04
force-pushed
the
02-13-feat_added_multi_routing_engine_support
branch
from
9360f73 to
e5acb8b
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@plugins/governance/main.go`:
- Around line 650-666: Replace the raw string key "model" used in ctx.SetValue
with a typed constant: define BifrostContextKeyModel of type BifrostContextKey
with the value "model" in the schemas/bifrost package (symbol:
BifrostContextKeyModel), then update both ctx.SetValue("model", newModel) calls
in the governance plugin to ctx.SetValue(schemas.BifrostContextKeyModel,
newModel); ensure you reference the schemas package
(schemas.BifrostContextKeyModel) and adjust any imports/types so the typed key
compiles with the existing ctx.SetValue signature.
🧹 Nitpick comments (2)
framework/logstore/migrations.go (1)
1138-1140: Wrap the migration error for consistency.Other migration functions in this file wrap errors with context (e.g.,
fmt.Errorf("error while adding X: %s", err.Error())), but this function returnsm.Migrate()directly.♻️ Suggested fix
- return m.Migrate() + err := m.Migrate() + if err != nil { + return fmt.Errorf("error while adding routing engines used column: %s", err.Error()) + } + return nilplugins/governance/main.go (1)
487-490: TODO for provider/budget violation handling.If you want, I can help turn this into a concrete error response and add tests.
| newModel := decision.Model + genaiRequestSuffix | ||
| // Add provider prefix if present (because there can be other routing rules down stream that can add the provider) | ||
| if decision.Provider != "" { | ||
| newModel = decision.Provider + "/" + newModel | ||
| } | ||
| ctx.SetValue("model", newModel) | ||
| } else { | ||
| // For regular requests, update in body | ||
| body["model"] = decision.Provider + "/" + decision.Model | ||
| newModel := decision.Model | ||
| // Add provider prefix if present (because there can be other routing rules down stream that can add the provider) | ||
| if decision.Provider != "" { | ||
| newModel = decision.Provider + "/" + newModel | ||
| } | ||
| body["model"] = newModel | ||
| } | ||
| ctx.SetValue(schemas.BifrostContextKeyRoutingEngineUsed, "routing-rule") | ||
| // Append routing-rule to routing engines used | ||
| schemas.AppendToContextList(ctx, schemas.BifrostContextKeyRoutingEnginesUsed, schemas.RoutingEngineRoutingRule) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, find and check the file structure
find . -name "main.go" -path "*/governance/*" | head -5
find . -name "bifrost.go" -path "*/schemas/*" | head -5Repository: maximhq/bifrost
Length of output: 113
🏁 Script executed:
# Check the content at the specified lines in plugins/governance/main.go
if [ -f "plugins/governance/main.go" ]; then
echo "=== Lines 650-666 in plugins/governance/main.go ==="
sed -n '650,666p' plugins/governance/main.go
fiRepository: maximhq/bifrost
Length of output: 885
🏁 Script executed:
# Check line 519 for the reference to loadBalanceProvider
if [ -f "plugins/governance/main.go" ]; then
echo "=== Lines 515-525 around line 519 ==="
sed -n '515,525p' plugins/governance/main.go
fiRepository: maximhq/bifrost
Length of output: 588
🏁 Script executed:
# Search for BifrostContextKeyModel in the schemas
fd bifrost.go | xargs grep -n "BifrostContextKeyModel" 2>/dev/null || echo "Key not found"Repository: maximhq/bifrost
Length of output: 72
🏁 Script executed:
# Run the suggested verification script to find all raw "model" context key usages
rg -n 'SetValue\("model"\)|Value\("model"\)' --type goRepository: maximhq/bifrost
Length of output: 224
🏁 Script executed:
# View core/schemas/bifrost.go to see existing BifrostContextKey patterns
head -100 core/schemas/bifrost.goRepository: maximhq/bifrost
Length of output: 3167
🏁 Script executed:
# Search for all BifrostContextKey definitions to understand the pattern
rg "BifrostContextKey" core/schemas/bifrost.go -A 1Repository: maximhq/bifrost
Length of output: 9364
🏁 Script executed:
# More comprehensive search for "model" key usage with different patterns
rg '(SetValue|Value|UserValue)\s*\(\s*["\x27]?model["\x27]?' --type go -nRepository: maximhq/bifrost
Length of output: 565
🏁 Script executed:
# Check if there are any type definitions for context keys
rg "type.*[Kk]ey" core/schemas/bifrost.go -B 2 -A 2Repository: maximhq/bifrost
Length of output: 427
🏁 Script executed:
# Check the genai.go context usage to understand SetUserValue vs SetValue
sed -n '410,420p' transports/bifrost-http/integrations/genai.go
sed -n '484,494p' transports/bifrost-http/integrations/genai.goRepository: maximhq/bifrost
Length of output: 839
🏁 Script executed:
# Check the BifrostContext implementation to understand SetValue vs SetUserValue
fd -e go -x grep -l "type BifrostContext" {} \;Repository: maximhq/bifrost
Length of output: 110
🏁 Script executed:
# View the BifrostContext definition
rg "type BifrostContext" -A 20Repository: maximhq/bifrost
Length of output: 5234
🏁 Script executed:
# Check if model key set in governance is read in genai by searching context flow
rg -B 3 -A 3 'ctx.UserValue\("model"\)' transports/bifrost-http/integrations/genai.goRepository: maximhq/bifrost
Length of output: 658
🏁 Script executed:
# Find the SetValue and Value method implementations in BifrostContext
rg -A 10 "func.*BifrostContext.*SetValue" core/schemas/context.go
rg -A 10 "func.*BifrostContext.*Value\(" core/schemas/context.goRepository: maximhq/bifrost
Length of output: 2677
🏁 Script executed:
# Check if BifrostContext uses standard context.Context Value/WithValue
rg "context.WithValue|context.Value" core/schemas/context.goRepository: maximhq/bifrost
Length of output: 107
🏁 Script executed:
# Check if genai functions receive BifrostContext and use it
rg -B 5 'ctx.UserValue\("model"\)' transports/bifrost-http/integrations/genai.go | head -40Repository: maximhq/bifrost
Length of output: 628
🏁 Script executed:
# Check the function signatures that use ctx.UserValue("model")
rg 'func.*\(ctx \*fasthttp.RequestCtx' transports/bifrost-http/integrations/genai.go | head -5Repository: maximhq/bifrost
Length of output: 433
🏁 Script executed:
# Check where the model key is set in governance main.go to understand context flow
rg -B 10 'ctx.SetValue\("model"' plugins/governance/main.go | head -30Repository: maximhq/bifrost
Length of output: 1134
🏁 Script executed:
# Search for reads of the "model" key in BifrostContext context (Value calls)
rg 'Value\("model"\)' --type goRepository: maximhq/bifrost
Length of output: 216
🏁 Script executed:
# Check if there are any references to the model in governance that read it back
rg -A 5 'SetValue\("model"' plugins/governance/main.go | grep -i "value\|model"Repository: maximhq/bifrost
Length of output: 278
🏁 Script executed:
# Check the full context of the loadBalanceProvider function around line 519
sed -n '510,540p' plugins/governance/main.goRepository: maximhq/bifrost
Length of output: 1350
Replace raw "model" context key with a typed key.
ctx.SetValue("model", ...) at lines 519 and 655 uses a string literal key, which violates the context-key rule. Switch to a typed BifrostContextKey constant.
Suggested fix
First, add the typed key constant to core/schemas/bifrost.go:
const (
BifrostContextKeyModel BifrostContextKey = "model"
)Then replace both usages in plugins/governance/main.go:
- ctx.SetValue("model", newModelWithRequestSuffix)
+ ctx.SetValue(schemas.BifrostContextKeyModel, newModelWithRequestSuffix)- ctx.SetValue("model", newModel)
+ ctx.SetValue(schemas.BifrostContextKeyModel, newModel)🤖 Prompt for AI Agents
In `@plugins/governance/main.go` around lines 650 - 666, Replace the raw string
key "model" used in ctx.SetValue with a typed constant: define
BifrostContextKeyModel of type BifrostContextKey with the value "model" in the
schemas/bifrost package (symbol: BifrostContextKeyModel), then update both
ctx.SetValue("model", newModel) calls in the governance plugin to
ctx.SetValue(schemas.BifrostContextKeyModel, newModel); ensure you reference the
schemas package (schemas.BifrostContextKeyModel) and adjust any imports/types so
the typed key compiles with the existing ctx.SetValue signature.
2 participants
Summary
Enhance routing engine tracking by replacing the single routing engine field with an array-based approach to support multiple routing engines in a single request. This allows for more accurate tracking when requests pass through multiple routing components.
Changes
BifrostContextKeyRoutingEngineUsedtoBifrostContextKeyRoutingEnginesUsedand changed its type from string to []stringAppendToContextListutility function to easily add routing engines to the contextType of change
Affected areas
How to test
Breaking changes
This change renames the context key from
routing_engine_usedtorouting_engines_usedand changes its type from string to []string. Any custom plugins that directly access this field will need to be updated.Database migration is included to rename the column and maintain backward compatibility.
Security considerations
No security implications. This change only affects telemetry and logging data structure.
Checklist
docs/contributing/README.mdand followed the guidelines