[DO NOT MERGE] refactor: replace key selector with key sorter for improved retry handling #1163
Conversation
📝 WalkthroughWalkthroughReplaces single-key selection with a pluggable key sorter returning multiple keys, threads retry attempt counts through request handlers, records tried-key IDs/names in context and DB, and propagates tried-key data into logging, telemetry, and OTEL spans. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant Bifrost
participant Provider
participant LogStore
participant Telemetry
Client->>Bifrost: Send request
Note right of Bifrost: keySorter(ctx, keys, provider, model) → []Key
Bifrost->>Provider: Send request using selected Key(s) (attempt 0)
alt provider succeeds
Provider-->>Bifrost: Response
Bifrost->>LogStore: update log (selectedKey, triedKeyIDs/names)
Bifrost->>Telemetry: emit span/labels (tried_key_ids/names)
Bifrost-->>Client: Return response
else provider error
Provider-->>Bifrost: Error
Bifrost->>Bifrost: executeRequestWithRetries → requestHandler(attempts+1)
Bifrost->>Provider: Retry with next Key(s)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Suggested reviewers
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
This stack of pull requests is managed by Graphite. Learn more about stacking. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/logging/pool.go (1)
13-37: Remove duplicate nil assignments.Lines 22-23 set
TriedKeyIDsandTriedKeyNamestonil, but lines 28-29 repeat the same assignments. This duplication is unnecessary and should be removed.🔎 Proposed fix to remove duplication
func (p *LoggerPlugin) putLogMessage(msg *LogMessage) { // Reset the message fields to avoid memory leaks msg.Operation = "" msg.RequestID = "" msg.ParentRequestID = "" msg.NumberOfRetries = 0 msg.FallbackIndex = 0 msg.SelectedKeyID = "" msg.SelectedKeyName = "" msg.TriedKeyIDs = nil msg.TriedKeyNames = nil msg.VirtualKeyID = "" msg.VirtualKeyName = "" msg.Timestamp = time.Time{} msg.InitialData = nil - msg.TriedKeyIDs = nil - msg.TriedKeyNames = nil // Don't reset UpdateData and StreamResponse here since they're returned // to their own pools in the defer function - just clear the pointers msg.UpdateData = nil msg.StreamResponse = nil p.logMsgPool.Put(msg) }
🧹 Nitpick comments (1)
core/schemas/bifrost.go (1)
16-22:KeySorterFuncuses*context.Context- consider usingcontext.Contextby value.The function signature uses
*context.Contextwhile similar functions in the codebase (e.g.,FetchNewRequestIDFuncincore/schemas/mcp.go) usecontext.Contextby value. Go's context package recommends passing context by value, not by pointer.If the sorter needs to modify the context, the standard pattern is to return a new context. If it doesn't need modification, pass by value for consistency.
🔎 Consider aligning with Go conventions
-type KeySorterFunc func(ctx *context.Context, keys []Key, providerKey ModelProvider, model string) ([]Key, error) +type KeySorterFunc func(ctx context.Context, keys []Key, providerKey ModelProvider, model string) ([]Key, error)
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (18)
core/bifrost.gocore/bifrost_test.gocore/key-sorting-algos/weighted-random.gocore/schemas/bifrost.gocore/schemas/context.gocore/utils.goframework/logstore/migrations.goframework/logstore/tables.goplugins/logging/main.goplugins/logging/operations.goplugins/logging/pool.goplugins/logging/utils.goplugins/otel/converter.goplugins/otel/main.goplugins/telemetry/main.goplugins/telemetry/utils.goui/app/workspace/logs/views/logDetailsSheet.tsxui/lib/types/logs.ts
💤 Files with no reviewable changes (2)
- plugins/telemetry/utils.go
- plugins/logging/utils.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/otel/main.gocore/key-sorting-algos/weighted-random.gocore/bifrost_test.goframework/logstore/tables.goplugins/logging/operations.gocore/schemas/context.goui/app/workspace/logs/views/logDetailsSheet.tsxframework/logstore/migrations.goplugins/otel/converter.goui/lib/types/logs.tsplugins/telemetry/main.gocore/utils.goplugins/logging/main.gocore/bifrost.gocore/schemas/bifrost.goplugins/logging/pool.go
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/otel/main.gocore/key-sorting-algos/weighted-random.gocore/bifrost_test.goframework/logstore/tables.goplugins/logging/operations.gocore/schemas/context.goframework/logstore/migrations.goplugins/otel/converter.goplugins/telemetry/main.gocore/utils.goplugins/logging/main.gocore/bifrost.gocore/schemas/bifrost.goplugins/logging/pool.go
🧬 Code graph analysis (7)
plugins/otel/main.go (2)
core/utils.go (1)
GetStringSliceFromContext(303-310)core/schemas/bifrost.go (2)
BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)
core/key-sorting-algos/weighted-random.go (2)
core/schemas/account.go (1)
Key(8-20)core/schemas/bifrost.go (1)
ModelProvider(46-46)
core/schemas/context.go (1)
core/schemas/bifrost.go (2)
BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)
ui/app/workspace/logs/views/logDetailsSheet.tsx (1)
ui/app/workspace/logs/views/logEntryDetailsView.tsx (1)
LogEntryDetailsView(15-49)
framework/logstore/migrations.go (2)
framework/migrator/migrator.go (3)
DefaultOptions(100-106)New(131-149)Migration(62-69)framework/logstore/tables.go (2)
Log(75-140)Log(143-145)
core/bifrost.go (5)
core/schemas/bifrost.go (9)
KeySorterFunc(22-22)KeySorterWeightedRandom(27-27)KeySorterCustom(28-28)RequestType(102-102)BifrostContextKeySelectedKeyID(140-140)BifrostContextKeySelectedKeyName(141-141)BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)ModelProvider(46-46)core/key-sorting-algos/weighted-random.go (1)
WeightedRandomKeySorter(35-105)ui/lib/types/config.ts (2)
RequestType(138-160)ModelProvider(198-201)core/utils.go (1)
IsStreamRequestType(197-199)core/schemas/account.go (1)
Key(8-20)
core/schemas/bifrost.go (2)
core/schemas/account.go (1)
Key(8-20)core/schemas/mcp.go (1)
MCPConfig(14-23)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (30)
core/utils.go (2)
302-310: LGTM: Well-structured context helper.The implementation follows the established pattern of
GetStringFromContextandGetIntFromContext, safely extracting a[]stringfrom context with proper type assertion and returning an empty slice (not nil) when the value is absent or of the wrong type.
410-418: LGTM: Standard clamping utility.The
clampfunction correctly constrains an integer to the inclusive range[min, max]. The logic is straightforward and correct.ui/lib/types/logs.ts (1)
267-268: LGTM: Type definitions align with backend schema.The optional
tried_key_idsandtried_key_namesfields correctly match the backend's[]stringtypes and are appropriately positioned alongside other key-related fields in theLogEntryinterface.core/schemas/context.go (1)
20-21: LGTM: Correctly protects internal context keys.Adding
BifrostContextKeyTriedKeyIDsandBifrostContextKeyTriedKeyNamesto thereservedKeysslice ensures these system-managed fields cannot be overwritten by user code viaSetValue, maintaining the integrity of retry tracking data.plugins/otel/main.go (2)
252-253: LGTM: Context extraction follows established patterns.The extraction of
triedKeyIDsandtriedKeyNamesfrom context usingGetStringSliceFromContextis consistent with the approach used for other key-related fields and correctly integrates tried-key tracking into the OTEL telemetry pipeline.
280-303: LGTM: Tried-key parameters properly threaded through resource span completion.Both streaming-final and non-streaming code paths correctly pass
triedKeyIDsandtriedKeyNamestocompleteResourceSpan, ensuring observability of retry attempts across all request types.Also applies to: 307-325
ui/app/workspace/logs/views/logDetailsSheet.tsx (1)
272-277: LGTM: UI correctly renders tried-key information.The conditional rendering of "Tried Keys" and "Tried Key Names" appropriately checks for non-empty arrays before displaying them, and using
.join(", ")to format the arrays as comma-separated strings is a suitable UI presentation for this metadata.plugins/logging/operations.go (2)
49-71: LGTM: Non-streaming update path correctly handles tried-key tracking.The
updateLogEntryfunction properly acceptstriedKeyIDsandtriedKeyNamesparameters and conditionally persists them only when non-empty, avoiding unnecessary database writes for requests that don't involve retries.
193-216: LGTM: Streaming update path mirrors non-streaming implementation.The
updateStreamingLogEntryfunction applies the same conditional persistence logic fortriedKeyIDsandtriedKeyNamesas the non-streaming path, ensuring consistent tried-key tracking across both request types.framework/logstore/migrations.go (2)
46-48: LGTM: Migration properly sequenced.The call to
migrationAddTriedKeyIDsAndTriedKeyNamesColumnsis correctly positioned after existing migrations and follows the established error-handling pattern.
681-722: LGTM: Migration correctly adds tried-key columns with proper rollback.The migration function properly:
- Uses transactions for atomicity
- Checks for column existence before adding (idempotent)
- Adds both
tried_key_idsandtried_key_namescolumns in a single migration- Provides a matching rollback that drops the columns if they exist
- Wraps errors with descriptive messages
core/bifrost_test.go (1)
59-62: Handler signature updated correctly for attempts tracking.The test handlers are properly updated to accept the new
attempts intparameter, aligning with the core retry logic changes. The parameter is intentionally unused in most tests since they focus on verifying retry counts viacallCountrather than the passed attempts value.core/key-sorting-algos/weighted-random.go (3)
33-40: Well-implemented weighted random key sorter with good performance characteristics.The algorithm using
log(u)/wfor weighted random selection is mathematically correct (Gumbel-max trick variant). The use ofsync.Poolfor both RNG instances and score buffers effectively minimizes allocations after warmup.Note that this function sorts the input slice in-place, which is efficient but callers should be aware they cannot reuse the original ordering. The function comment mentions "shuffles keys in-place" which is good documentation.
99-102: Fisher-Yates shuffle on zero-weight keys looks correct.The loop bounds correctly shuffle only the
keys[z:]tail containing zero-weight keys. The index calculationj := z + r.Intn(i-z+1)ensuresjis in range[z, i]which is correct for the unshuffled portion.
16-23: Thread-safe RNG pooling with atomic seed increment.Good approach using
atomic.AddUint64to ensure each pooled RNG gets a unique seed. This avoids contention while maintaining randomness across goroutines.plugins/logging/main.go (2)
370-371: Context retrieval for tried keys is consistent with existing patterns.Using
bifrost.GetStringSliceFromContextaligns with the new utility function and follows the same pattern as other context value retrievals in this function.
66-67: NewTriedKeyIDsandTriedKeyNamesfields properly integrated intoLogMessagewith complete pool reset handling.The fields are correctly defined, retrieved from context, assigned to the LogMessage struct, and passed through all updateLogEntry and updateStreamingLogEntry calls. The pool reset logic in putLogMessage (pool.go, lines 22–23) properly clears these slice fields to nil, preventing data leakage between pooled objects.
plugins/otel/converter.go (2)
650-651: Function signature properly extended for tried key tracking.The new parameters integrate well with the existing function signature pattern. The placement after
selectedKeyNamemaintains logical grouping of key-related parameters.
1068-1073: Tried keys correctly added to both span and resource attributes.The implementation properly:
- Guards against empty slices with
len() > 0checks- Uses comma-joined format consistent with other multi-value attributes in this file
- Adds to both span attributes (
gen_ai.*prefix) and resource attributes (without prefix)Also applies to: 1096-1101
plugins/telemetry/main.go (2)
101-102: New Prometheus labels for tried keys added to default labels.These labels integrate properly with the existing telemetry structure. However, be aware this introduces new dimensions to all Bifrost metrics, which:
- Is a breaking change for existing Prometheus queries/alerts
- May increase label cardinality if requests frequently try multiple keys
If cardinality becomes an issue, consider making these labels opt-in via config or limiting the number of keys recorded.
333-334: Comma-joined format for tried keys in Prometheus labels.The implementation correctly joins the string slices with commas, maintaining consistency with how the OTEL plugin handles these values.
core/schemas/bifrost.go (2)
24-29: Clean extensibility design for key sorting algorithms.The
KeySortertype with constants provides a good extension point for future sorting algorithms while maintaining type safety. The distinction betweenKeySorterWeightedRandomandKeySorterCustomallows both built-in and user-defined sorters.
142-143: New context keys for tried key tracking follow existing conventions.The naming pattern (
bifrost-tried-key-*) is consistent with other Bifrost context keys, and the comments clearly document the expected types and that these are set internally.framework/logstore/tables.go (1)
7-7: Consider whether sonic's default configuration is appropriate for this use case, or useConfigStdfor strictencoding/jsoncompatibility.
sonic.Marshal()without configuration differs fromencoding/jsonin several ways:EscapeHTMLdefaults tofalse(vs.truein the standard library),SortKeysis off, and there are subtle differences in number handling. While sonic can be a high-performance alternative, usingsonic.ConfigStdensures exact parity withencoding/jsonif that compatibility is required. If the behavioral differences are acceptable for serializing these fields, the current usage is fine; otherwise, consider passingsonic.ConfigStdto the Marshal calls.core/bifrost.go (6)
17-17: LGTM!The new import and
keySorterfield are properly defined. The field type correctly matches theKeySorterFuncsignature fromcore/schemas/bifrost.go.Also applies to: 73-73
124-139: LGTM!The key sorter initialization logic is well-structured:
- Sensible default to
WeightedRandomalgorithm- Proper validation requiring
CustomKeySorterwhen using custom algorithm- Fail-fast behavior for unsupported algorithms
331-333: LGTM!The
ListModelsRequestcorrectly ignores theattemptsparameter since it already passes all keys to the provider and doesn't need key rotation on retry.
2868-2886: LGTM!The call site correctly handles the new
[]schemas.Keyreturn type fromselectKeyFromProviderForModel.
2737-2773: LGTM!The
executeRequestWithRetriesfunction signature correctly adds theattemptsparameter to enable key rotation on retries. The implementation properly passes the attempt count to the handler.
3437-3551: LGTM!The
selectKeyFromProviderForModelfunction correctly returns a sorted slice of keys via thekeySorter, enabling deterministic key rotation on retries. All return paths are consistent with the new[]schemas.Keyreturn type.
Pratham-Mishra04
force-pushed
the
12-23-feat_added_responses_mcp_tool_execute_endpoint
branch
from
06d67a3 to
4a1a08e
Compare
Pratham-Mishra04
force-pushed
the
12-23-feat_added_retries_across_multiple_keys
branch
from
d68a8fc to
6e9445e
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
core/key-sorting-algos/weighted-random.go (1)
76-80: Consider a bounded retry for zero-check.While
Float64()returning exactly 0.0 is extremely rare, a defensive bounded loop prevents theoretical infinite loops.🔎 Proposed bounded retry
// u in (0,1]; ensure non-zero so log(u) is finite. u := r.Float64() - for u == 0 { + for i := 0; u == 0 && i < 10; i++ { u = r.Float64() } + if u == 0 { + u = 1e-15 // fallback to tiny positive value + } scores[i] = math.Log(u) / wplugins/telemetry/main.go (1)
101-102: Consider cardinality impact of tried_key labels.Adding
tried_key_idsandtried_key_namesas Prometheus labels could significantly increase metric cardinality, especially if there are many keys and retry patterns vary. Each unique combination creates a new time series.If cardinality becomes problematic, consider:
- Moving these to span attributes only (OTEL) rather than metric labels
- Using a count of tried keys instead of the full list
- Making these labels opt-in via configuration
Also applies to: 333-334
plugins/logging/operations.go (1)
66-83: Consider extracting tried-key serialization to a helper function.The serialization logic for
triedKeyIDsandtriedKeyNamesis duplicated in bothupdateLogEntryandupdateStreamingLogEntry. Extracting this to a helper would improve maintainability.🔎 Proposed helper function
// serializeTriedKeys serializes tried key IDs and names into the updates map func (p *LoggerPlugin) serializeTriedKeys(updates map[string]interface{}, triedKeyIDs, triedKeyNames []string) { if len(triedKeyIDs) > 0 { tempEntry := &logstore.Log{TriedKeyIDsParsed: triedKeyIDs} if err := tempEntry.SerializeFields(); err != nil { p.logger.Error("failed to serialize tried key IDs: %v", err) } else { updates["tried_key_ids"] = tempEntry.TriedKeyIDs } } if len(triedKeyNames) > 0 { tempEntry := &logstore.Log{TriedKeyNamesParsed: triedKeyNames} if err := tempEntry.SerializeFields(); err != nil { p.logger.Error("failed to serialize tried key names: %v", err) } else { updates["tried_key_names"] = tempEntry.TriedKeyNames } } }core/schemas/bifrost.go (2)
16-22: Consider usingcontext.Contextinstead of*context.Context.The
KeySorterFuncsignature uses*context.Context, but Go conventions typically passcontext.Contextby value (it's already a reference type internally). Using a pointer to Context is unusual and may cause confusion.🔎 Proposed signature change
-type KeySorterFunc func(ctx *context.Context, keys []Key, providerKey ModelProvider, model string) ([]Key, error) +type KeySorterFunc func(ctx context.Context, keys []Key, providerKey ModelProvider, model string) ([]Key, error)Note: This would require updating the
WeightedRandomKeySorterimplementation incore/key-sorting-algos/weighted-random.goas well.
41-42: Document the precedence between KeySorterAlgorithm and CustomKeySorter.It would be helpful to document in comments what happens when:
- Neither
KeySorterAlgorithmnorCustomKeySorteris set (default behavior?)- Both are set (which takes precedence?)
🔎 Suggested documentation
MCPConfig *MCPConfig // MCP (Model Context Protocol) configuration for tool integration - KeySorterAlgorithm KeySorter // Name of the key sorter algorithm to use - CustomKeySorter KeySorterFunc // Custom key sorter function to use + KeySorterAlgorithm KeySorter // Name of the key sorter algorithm to use (default: weighted_random) + CustomKeySorter KeySorterFunc // Custom key sorter function (takes precedence over KeySorterAlgorithm if set)
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (11)
core/go.sumis excluded by!**/*.sumframework/go.sumis excluded by!**/*.sumplugins/governance/go.sumis excluded by!**/*.sumplugins/jsonparser/go.sumis excluded by!**/*.sumplugins/logging/go.sumis excluded by!**/*.sumplugins/maxim/go.sumis excluded by!**/*.sumplugins/mocker/go.sumis excluded by!**/*.sumplugins/otel/go.sumis excluded by!**/*.sumplugins/semanticcache/go.sumis excluded by!**/*.sumplugins/telemetry/go.sumis excluded by!**/*.sumtransports/go.sumis excluded by!**/*.sum
📒 Files selected for processing (29)
core/bifrost.gocore/bifrost_test.gocore/go.modcore/key-sorting-algos/weighted-random.gocore/schemas/bifrost.gocore/schemas/context.gocore/utils.goframework/go.modframework/logstore/migrations.goframework/logstore/tables.goplugins/governance/go.modplugins/jsonparser/go.modplugins/logging/go.modplugins/logging/main.goplugins/logging/operations.goplugins/logging/pool.goplugins/logging/utils.goplugins/maxim/go.modplugins/mocker/go.modplugins/otel/converter.goplugins/otel/go.modplugins/otel/main.goplugins/semanticcache/go.modplugins/telemetry/go.modplugins/telemetry/main.goplugins/telemetry/utils.gotransports/go.modui/app/workspace/logs/views/logDetailsSheet.tsxui/lib/types/logs.ts
💤 Files with no reviewable changes (2)
- plugins/logging/utils.go
- plugins/telemetry/utils.go
✅ Files skipped from review due to trivial changes (1)
- plugins/maxim/go.mod
🚧 Files skipped from review as they are similar to previous changes (7)
- framework/logstore/migrations.go
- plugins/logging/pool.go
- core/bifrost_test.go
- ui/app/workspace/logs/views/logDetailsSheet.tsx
- ui/lib/types/logs.ts
- core/schemas/context.go
- core/utils.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/telemetry/main.goplugins/semanticcache/go.modplugins/otel/main.gocore/key-sorting-algos/weighted-random.goplugins/governance/go.modcore/go.modplugins/otel/converter.goplugins/logging/main.goplugins/otel/go.modtransports/go.modcore/schemas/bifrost.goplugins/logging/operations.goplugins/mocker/go.modframework/logstore/tables.goframework/go.modplugins/telemetry/go.modplugins/logging/go.modplugins/jsonparser/go.modcore/bifrost.go
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/telemetry/main.goplugins/otel/main.gocore/key-sorting-algos/weighted-random.goplugins/otel/converter.goplugins/logging/main.gocore/schemas/bifrost.goplugins/logging/operations.goframework/logstore/tables.gocore/bifrost.go
🧬 Code graph analysis (6)
plugins/telemetry/main.go (2)
core/utils.go (3)
GetStringFromContext(293-300)GetStringSliceFromContext(303-310)GetIntFromContext(313-320)core/schemas/bifrost.go (7)
BifrostContextKey(131-131)BifrostContextKeySelectedKeyID(140-140)BifrostContextKeySelectedKeyName(141-141)BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)BifrostContextKeyNumberOfRetries(144-144)BifrostContextKeyFallbackIndex(145-145)
plugins/otel/main.go (2)
core/utils.go (1)
GetStringSliceFromContext(303-310)core/schemas/bifrost.go (2)
BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)
core/key-sorting-algos/weighted-random.go (2)
core/schemas/account.go (1)
Key(8-20)core/schemas/bifrost.go (1)
ModelProvider(46-46)
plugins/logging/main.go (2)
core/schemas/bifrost.go (7)
BifrostContextKeyFallbackIndex(145-145)BifrostContextKeySelectedKeyID(140-140)BifrostContextKeySelectedKeyName(141-141)BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)BifrostContextKey(131-131)BifrostContextKeyNumberOfRetries(144-144)core/utils.go (3)
GetIntFromContext(313-320)GetStringFromContext(293-300)GetStringSliceFromContext(303-310)
core/schemas/bifrost.go (3)
core/schemas/account.go (1)
Key(8-20)ui/lib/types/config.ts (1)
ModelProvider(198-201)core/schemas/mcp.go (1)
MCPConfig(14-23)
plugins/logging/operations.go (1)
framework/logstore/tables.go (2)
Log(75-142)Log(145-147)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (12)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (34)
plugins/semanticcache/go.mod (3)
1-114: Review this with actual stack context.The referenced PRs (#1161, #1160) cannot be verified in the current repository. Before proceeding with the prescriptive guidance about downstack PR merge order and stack-wide verification, confirm these PRs exist and are related to this change.
If this change is part of a dependency stack:
- Verify direct dependency versions (bifrost/core v1.2.40, bifrost/framework v1.1.50) align with downstack modules
- Ensure all stack PRs are merged before deploying to avoid version mismatches
- Test the full stack together after all changes are integrated
If this is a standalone change, standard go mod tidying on the root module and local build/test validation are sufficient.
38-39: Sonic version alignment across main modules is correct, but test/example modules lag.The sonic v1.14.2 + loader v0.4.0 versions are correctly aligned across core, framework, and all production plugins including semanticcache. However, test and example modules (examples/plugins/hello-world, tests/scripts/1millogs, tests/core-mcp) remain on sonic v1.14.1 + loader v0.3.0. Verify whether this version lag in test/example modules is intentional or if they should be updated to match the production stack.
44-64: No action required; go-openapi updates are intentional transitive dependencies.The 21 go-openapi packages in lines 44–64 are indirect dependencies pulled in by Weaviate v1.34.5 (and potentially invopop/jsonschema v0.13.0 and mark3labs/mcp-go v0.43.2) in the framework module. Semanticcache inherits these through its direct dependency on framework v1.1.50.
Version consistency verified: All 21 packages have identical versions across framework, semanticcache, and other affected plugins (governance, logging, maxim, otel, telemetry). Core has no go-openapi dependencies, creating no conflicts. The updates are part of a single coherent commit ("feat: added retries across multiple keys") applied uniformly across the module stack.
plugins/jsonparser/go.mod (3)
32-33: No compatibility concerns identified with bytedance/sonic v1.14.2.The v1.14.2 release contains bug fixes and minor features with no breaking changes. Notably, the decoder fix for interface map keys and the improvement to validate skip behavior (to align with encoding/json) support the JSON serialization migration and enhance compatibility with existing code.
Likely an incorrect or invalid review comment.
50-54: No compatibility concerns identified.The golang.org/x package versions listed (golang.org/x/arch v0.23.0, golang.org/x/net v0.48.0, golang.org/x/oauth2 v0.34.0, golang.org/x/sys v0.39.0, golang.org/x/text v0.32.0) are compatible with Go 1.25.5. These modules are routinely maintained and used with current Go releases, including Go 1.25.x, with no reported breaking compatibility issues.
12-13: AWS SDK versions are compatible and stable. Verification confirms:
- config v1.32.6, credentials v1.19.6, s3 v1.94.0, and sso v1.30.8 are all compatible with the base AWS SDK v1.41.0
- No breaking changes in credentials or S3 APIs that would affect retry logic
- All versions are stable releases (v1.41.0 released Nov 19, 2025)
No further action required.
plugins/governance/go.mod (1)
17-18: Dependency updates are consistent with the broader module refresh.The updates include AWS SDK v2, Sonic, OpenAPI tooling, OTEL, and Go standard library packages. These align with the ecosystem-wide dependency refresh across the codebase.
Also applies to: 28-28, 30-30, 37-38, 44-64, 74-74, 77-77, 83-84, 89-91, 94-94, 96-109
plugins/otel/go.mod (1)
8-9: OTEL plugin dependencies updated consistently.The direct upgrades to grpc (1.77.0) and protobuf (1.36.11), along with OTEL 1.39.0, maintain compatibility across the telemetry stack.
Also applies to: 17-18, 28-28, 30-30, 41-61, 72-72, 74-74, 80-81, 85-87, 90-90, 92-94, 96-102, 110-111, 115-117
core/go.mod (1)
8-10: Core module dependencies updated appropriately.The updates maintain consistency with downstream plugin modules. Note the addition of
github.com/Masterminds/semver/v3 v3.3.1as a new indirect dependency.Also applies to: 12-12, 17-17, 20-22, 27-27, 39-39, 45-45, 50-50, 52-52, 64-66
framework/go.mod (1)
8-13: Framework dependencies comprehensively updated.The updates span vector databases (Qdrant, Weaviate, Redis), OTEL components, and OpenAPI tooling, maintaining consistency with the broader dependency refresh.
Also applies to: 23-23, 28-38, 45-49, 56-57, 64-64, 70-71, 76-85, 89-90, 92-92, 103-103, 106-114
plugins/logging/go.mod (1)
6-6: Logging plugin dependencies updated consistently.The direct Sonic dependency (1.14.2) aligns with the PR's shift to bytedance/sonic for JSON serialization, and all indirect dependencies follow the coordinated update pattern.
Also applies to: 16-17, 27-27, 29-29, 36-36, 42-62, 72-72, 75-75, 81-82, 87-89, 92-92, 94-107
plugins/telemetry/go.mod (1)
9-9: Telemetry plugin dependencies updated as expected.The direct fasthttp dependency (1.68.0) and all indirect dependencies follow the coordinated update pattern across the codebase.
Also applies to: 17-18, 28-28, 30-30, 38-39, 45-65, 75-75, 78-78, 88-89, 94-95, 98-98, 100-113
plugins/otel/main.go (2)
252-253: LGTM: Tried keys context extraction added.The implementation correctly uses
GetStringSliceFromContextto retrieve the tried key IDs and names from the context, consistent with the pattern used for other context values in this function.
292-293: LGTM: Tried keys propagated to resource spans.Both streaming (lines 292-293) and non-streaming (lines 317-318) paths correctly pass
triedKeyIDsandtriedKeyNamestocompleteResourceSpan, ensuring comprehensive telemetry coverage of retry behavior.Also applies to: 317-318
plugins/mocker/go.mod (1)
15-16: Routine dependency updates approved.Indirect dependency version bumps are minor/patch updates that align with standard maintenance. No concerns with the specified versions.
core/key-sorting-algos/weighted-random.go (2)
33-45: LGTM! Well-designed weighted random key sorter.The implementation correctly uses the exponential sampling algorithm (
log(u)/w) for weighted random selection, which provides proper probability distribution. The use ofsync.Poolfor both RNG and score buffers minimizes allocations in hot paths.
98-102: LGTM! Fisher-Yates shuffle is correctly implemented.The shuffle on zero-weight keys ensures varied ordering in the tail, preventing deterministic fallback patterns for zero-weight keys.
plugins/logging/main.go (2)
66-67: LGTM! Tried-key fields properly added to LogMessage.The new
TriedKeyIDsandTriedKeyNamesfields are correctly integrated into the logging flow, enabling observability of key selection across retries.
367-388: LGTM! Context extraction follows established patterns.The use of
bifrost.GetStringFromContextandbifrost.GetStringSliceFromContextfor extracting context values is consistent with the rest of the codebase and aligns with the new helper functions incore/utils.go.framework/logstore/tables.go (3)
86-87: LGTM! Tried-key fields correctly structured for GORM persistence.The approach of using
stringfields for database storage with separateParsedvariants (taggedgorm:"-") follows the established pattern for other JSON-serialized fields in this struct. This addresses the previous review concern about GORM slice serialization.Also applies to: 136-137
293-307: LGTM! Serialization follows established patterns.The serialization of
TriedKeyIDsParsedandTriedKeyNamesParsedcorrectly mirrors the pattern used for other JSON-serialized fields, with proper nil checks before marshaling.
7-7: No action needed. Sonic is explicitly designed to be compatible with customMarshalJSON/UnmarshalJSONimplementations—it natively supports thejson.Marshalerandjson.Unmarshalerinterfaces and provides maximum stdlib compatibility through itsConfigStdprofile. The migration is safe.transports/go.mod (2)
7-7: LGTM! Dependency updates are reasonable patch/minor bumps.The updates to
bytedance/sonic(1.14.2),valyala/fasthttp(1.68.0), and indirect dependencies appear to be standard maintenance updates.Also applies to: 22-22
11-11: Confirm mcp-go version and test compatibility with v0.43.0 features.The version
v0.43.2does not exist as an official release (latest v0.43.* is v0.43.0). No breaking changes are documented between v0.41.1 and v0.43.0, but verify that the code intransports/bifrost-http/handlers/mcp_server.gois compatible with v0.43.0's new features:SessionIdManagerResolver,SessionWithResourceTemplates, HTTP/StdioRoots, and theTitlefield added to theImplementationstruct. Run tests to confirm compatibility.plugins/telemetry/main.go (1)
307-321: LGTM! Context extraction is clean and consistent.The use of
bifrost.GetStringFromContext,bifrost.GetStringSliceFromContext, andbifrost.GetIntFromContextprovides a clean pattern for extracting context values with safe defaults.plugins/logging/operations.go (1)
69-73: LGTM! Error handling follows graceful degradation pattern.Logging serialization errors without failing the operation is appropriate for a logging plugin - it ensures that a serialization issue doesn't break the main request flow.
Also applies to: 78-82
plugins/otel/converter.go (2)
650-651: LGTM! Signature extension for tried-key tracking.The addition of
triedKeyIDsandtriedKeyNamesparameters aligns with the broader PR objective of tracking tried keys across retries for observability.
1068-1073: LGTM! Tried-key attributes correctly added to both span and resource.Following the established pattern for
selected_key_id/virtual_key_id, the tried-key attributes are added to both span attributes (withgen_ai.prefix) and resource attributes, ensuring visibility across different OTEL tooling views.Also applies to: 1096-1101
core/schemas/bifrost.go (1)
142-143: LGTM! Context keys follow established patterns.The new
BifrostContextKeyTriedKeyIDsandBifrostContextKeyTriedKeyNamescontext keys follow the existing naming convention and include appropriate documentation noting they are set by Bifrost internally.core/bifrost.go (5)
124-139: LGTM! Solid initialization logic for the key sorter.The initialization properly handles the default case (WeightedRandom), validates custom sorters, and fails fast on unsupported algorithms. The error checking for missing custom sorter functions prevents runtime nil pointer dereferences.
Based on coding guidelines, this change is part of the PR stack (#1163, #1161, #1160) that introduces deterministic retry behavior. The pluggable sorter design aligns well with the stack's objectives.
331-331: Correct adaptation to new executeRequestWithRetries signature.The underscore correctly indicates that ListModels doesn't need the attempts parameter, which makes sense since list operations don't benefit from key-based retry logic.
2734-2762: Well-structured fix that addresses the critical bug from the previous review.The function correctly accumulates tried keys across retry attempts by using pointer parameters (
*triedKeyIDs = append(*triedKeyIDs, prevKey.ID)). The semantic is sound: "tried keys" represent keys attempted before the current one, which is appropriate for observability and debugging.Note on retry behavior beyond available keys:
When retry attempts exceed the number of available keys, theclampfunction causes the last key to be reused. For example, with 3 keys and 5 retries:
- Attempts 0, 1, 2 use keys[0], keys[1], keys[2]
- Attempts 3, 4, 5 reuse keys[2]
This is likely acceptable given that retries only occur for rate limits and transient network errors where backoff + key reuse may succeed. However, confirm this aligns with the intended retry behavior across the stack.
Could you verify that reusing the last key when attempts exceed the key count aligns with the expected behavior for the stack's deterministic retry feature?
2933-2944: Excellent! The tried keys are now properly accumulated across retry attempts.This implementation correctly addresses the critical issue flagged in the previous review. The slices are declared outside the retry closures and modified via pointer parameters, ensuring that each retry adds to the cumulative list of tried keys rather than resetting them.
The approach is applied consistently to both streaming and non-streaming request paths, maintaining correctness across all code paths.
✅ Confirming that the bug identified in the past review comments has been fully resolved.
3445-3559: Clean refactor to support pluggable key sorting.The signature change from returning a single
Keyto[]schemas.Keyaligns with the sorter-based architecture introduced in this stack. The integration point at line 3553 (sortedKeys, err := bifrost.keySorter(...)) properly delegates sorting to the configured algorithm, enabling the deterministic retry behavior described in the PR objectives.All return paths have been correctly updated to return key slices, and error handling remains robust.
Based on the stack context, this change enables the sorted-list consumption strategy for retries described in PR #1163.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (5)
core/schemas/bifrost.go (1)
16-29: Key sorter types and config extension look correct; minor doc nit only.The
KeySorterFunc/KeySorterabstractions andBifrostConfigfields (KeySorterAlgorithm,CustomKeySorter) are wired in a way that matches how core/bifrost.go uses them (defaulting to weighted random, with a validated custom override). Only small nit: the comment on Line 16 could say “all the available keys for the provider and model” for clarity.Also applies to: 35-43
plugins/logging/operations.go (1)
54-83: Tried-key logging is integrated correctly; consider a tiny helper to reduce duplication.The new
triedKeyIDs/triedKeyNamesparameters are propagated and serialized in the same pattern as other JSON fields, and failures are non-fatal (just log and continue), which is appropriate for observability-only data. You duplicate the templogstore.Log+SerializeFieldspattern four times; if this evolves further, a small helper (e.g.,serializeTriedKeys) could centralize that logic, but it’s not urgent.Also applies to: 210-240
core/bifrost.go (1)
3445-3559: Key sorter integration inselectKeyFromProviderForModelis mostly robust; add a guard for empty outputs.The refactor to return
[]schemas.Keyand then callbifrost.keySorter(ctx, supportedKeys, providerKey, model)is clean: you still honor direct keys, skip-selection flags, batch/file constraints, model/deployment filtering, and explicit key-name selection before sorting.One defensive improvement: if
supportedKeysis non-empty butkeySorterreturns an empty slice without error, you’ll fall back to a zero-valued key downstream instead of failing fast. Consider treating that as an error, e.g.:Proposed guard for empty sorter output
- sortedKeys, err := bifrost.keySorter(ctx, supportedKeys, providerKey, model) - if err != nil { - return nil, err - } - - return sortedKeys, nil + sortedKeys, err := bifrost.keySorter(ctx, supportedKeys, providerKey, model) + if err != nil { + return nil, err + } + if len(sortedKeys) == 0 { + return nil, fmt.Errorf("key sorter returned no keys for provider %v and model %s", providerKey, model) + } + return sortedKeys, nilThis makes misbehaving custom sorters easier to diagnose.
ui/app/workspace/logs/views/logDetailsSheet.tsx (1)
272-277: Tried-key display wiring is correct; you may later improve UX for long lists.The new
LogEntryDetailsViewrows forlog.tried_key_idsandlog.tried_key_namesare correctly guarded on non-empty arrays and render a simple comma‑separated list, which is fine for now. If you expect many keys per request, a future enhancement could be truncation with a tooltip or popover, but it’s not required for functionality.plugins/otel/converter.go (1)
650-651: OTEL span/resource attributes for tried keys are wired consistently with existing metadata.
completeResourceSpannow:
- Accepts
triedKeyIDs/triedKeyNames.- Emits them on the span as
gen_ai.tried_key_ids/gen_ai.tried_key_names(comma‑joined).- Mirrors them on the resource as
tried_key_ids/tried_key_names, alongside the existing virtual/selected/team/customer attributes.This matches your current pattern of flattening slices into string attributes and should play nicely with most OTEL backends; just confirm your collectors/exporters see and index these new attributes as expected.
Also applies to: 1068-1073, 1096-1101
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (11)
core/go.sumis excluded by!**/*.sumframework/go.sumis excluded by!**/*.sumplugins/governance/go.sumis excluded by!**/*.sumplugins/jsonparser/go.sumis excluded by!**/*.sumplugins/logging/go.sumis excluded by!**/*.sumplugins/maxim/go.sumis excluded by!**/*.sumplugins/mocker/go.sumis excluded by!**/*.sumplugins/otel/go.sumis excluded by!**/*.sumplugins/semanticcache/go.sumis excluded by!**/*.sumplugins/telemetry/go.sumis excluded by!**/*.sumtransports/go.sumis excluded by!**/*.sum
📒 Files selected for processing (29)
core/bifrost.gocore/bifrost_test.gocore/go.modcore/key-sorting-algos/weighted-random.gocore/schemas/bifrost.gocore/schemas/context.gocore/utils.goframework/go.modframework/logstore/migrations.goframework/logstore/tables.goplugins/governance/go.modplugins/jsonparser/go.modplugins/logging/go.modplugins/logging/main.goplugins/logging/operations.goplugins/logging/pool.goplugins/logging/utils.goplugins/maxim/go.modplugins/mocker/go.modplugins/otel/converter.goplugins/otel/go.modplugins/otel/main.goplugins/semanticcache/go.modplugins/telemetry/go.modplugins/telemetry/main.goplugins/telemetry/utils.gotransports/go.modui/app/workspace/logs/views/logDetailsSheet.tsxui/lib/types/logs.ts
💤 Files with no reviewable changes (2)
- plugins/logging/utils.go
- plugins/telemetry/utils.go
🚧 Files skipped from review as they are similar to previous changes (7)
- plugins/otel/main.go
- core/bifrost_test.go
- core/key-sorting-algos/weighted-random.go
- plugins/logging/pool.go
- plugins/logging/main.go
- framework/logstore/tables.go
- core/utils.go
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/semanticcache/go.modui/app/workspace/logs/views/logDetailsSheet.tsxframework/logstore/migrations.goplugins/logging/go.modplugins/telemetry/main.goui/lib/types/logs.tsplugins/jsonparser/go.modcore/schemas/bifrost.gocore/go.modplugins/mocker/go.modplugins/telemetry/go.modcore/schemas/context.goplugins/governance/go.modframework/go.modplugins/maxim/go.modplugins/logging/operations.goplugins/otel/go.modplugins/otel/converter.gocore/bifrost.gotransports/go.mod
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
framework/logstore/migrations.goplugins/telemetry/main.gocore/schemas/bifrost.gocore/schemas/context.goplugins/logging/operations.goplugins/otel/converter.gocore/bifrost.go
🧬 Code graph analysis (7)
ui/app/workspace/logs/views/logDetailsSheet.tsx (1)
ui/app/workspace/logs/views/logEntryDetailsView.tsx (1)
LogEntryDetailsView(15-49)
framework/logstore/migrations.go (2)
framework/migrator/migrator.go (3)
DefaultOptions(100-106)New(131-149)Migration(62-69)framework/logstore/tables.go (2)
Log(75-142)Log(145-147)
plugins/telemetry/main.go (2)
core/utils.go (3)
GetStringFromContext(293-300)GetStringSliceFromContext(303-310)GetIntFromContext(313-320)core/schemas/bifrost.go (7)
BifrostContextKey(131-131)BifrostContextKeySelectedKeyID(140-140)BifrostContextKeySelectedKeyName(141-141)BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)BifrostContextKeyNumberOfRetries(144-144)BifrostContextKeyFallbackIndex(145-145)
core/schemas/bifrost.go (2)
core/schemas/account.go (1)
Key(8-20)core/schemas/mcp.go (1)
MCPConfig(14-23)
core/schemas/context.go (1)
core/schemas/bifrost.go (2)
BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)
plugins/logging/operations.go (1)
framework/logstore/tables.go (2)
Log(75-142)Log(145-147)
core/bifrost.go (2)
core/schemas/bifrost.go (8)
KeySorterFunc(22-22)KeySorterWeightedRandom(27-27)KeySorterCustom(28-28)BifrostContextKeySelectedKeyID(140-140)BifrostContextKeySelectedKeyName(141-141)BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)ModelProvider(46-46)core/key-sorting-algos/weighted-random.go (1)
WeightedRandomKeySorter(35-105)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (17)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (30)
plugins/mocker/go.mod (3)
7-7: Verify that bifrost/core v1.2.40 includes all key sorter refactoring changes.This PR introduces a significant refactoring around key sorters and retry handling. Ensure that the core module version bump (to v1.2.40) reflects all changes from the broader PR stack and is fully compatible with this plugin.
Please confirm:
- Does v1.2.40 include the pluggable KeySorterFunc and multi-key handling changes?
- Are there any breaking changes in the core API that require updates in this plugin?
- Is this the correct version from the downstack PR?
35-36: Verify bytedance/sonic versions support the JSON serialization switch.The updates to bytedance/sonic (v1.14.2) and sonic/loader (v0.4.0) align with the PR objective to switch from encoding/json to bytedance/sonic for improved serialization performance. Verify that these specific versions are stable and required for the serialization changes across the codebase.
Is the bump from sonic/loader v0.3.0 → v0.4.0 a significant enough change that warrants validation? The loader module is typically infrastructure-level and minor version bumps may have implications.
15-16: Clarify the necessity and rationale for broad dependency updates.The go.mod file shows systematic version bumps across AWS SDK, stdlib modules (arch, net, oauth2, sys, text), and other ecosystem packages. While these appear to be minor/patch-level updates, it's unclear whether:
- These updates are required for the key sorter refactoring and retry handling changes, or
- They were pulled in transitively by updating bifrost/core or other direct dependencies.
For a PR focused on key sorting and retry logic, the breadth of indirect dependency changes is noteworthy. Ensure that:
- Each update is intentional and necessary
- No security advisories affect these versions
- Go's minimal version selection has not inadvertently pulled in unexpected versions
If these are transitive updates from bifrost/core v1.2.40, that's expected and can be documented. If they were manually bumped, please clarify the rationale.
Also applies to: 26-26, 28-28, 40-40, 43-43, 50-50, 53-57
plugins/jsonparser/go.mod (2)
12-13: Verify necessity and alignment of dependency updates across the stack.This file shows updates to 10+ indirect dependencies. Given this is part of a multi-PR stack, confirm whether these updates are:
- Directly required for the key sorter/retry feature implementation
- Automatically pulled in as transitive dependency updates
- Bundled modernization updates
Particularly for context: ensure the core module's go.mod (if different) aligns with these versions to prevent version inconsistencies across the monorepo.
32-33: No issues found. The bytedance/sonic v1.14.2 and sonic/loader v0.4.0 versions are correctly paired and consistently aligned across all modules (core/go.mod and plugins/jsonparser/go.mod). The versions are compatible with each other and the rest of the codebase.core/schemas/context.go (1)
20-21: LGTM! Reserved keys extended correctly.The addition of
BifrostContextKeyTriedKeyIDsandBifrostContextKeyTriedKeyNamesto the reserved keys list is consistent with the broader tried-key tracking feature introduced in this PR. The implementation correctly prevents user code from overwriting these system-managed context values.framework/logstore/migrations.go (2)
46-48: LGTM! Migration invocation correctly placed.The new migration is invoked in the proper sequence after
migrationAddRawRequestColumn, ensuring the logs table is progressively updated.
681-722: LGTM! Migration follows established patterns.The migration implementation correctly:
- Uses transactions for atomicity
- Guards against re-execution with
HasColumnchecks- Provides symmetric rollback logic
- Uses descriptive error messages
- Matches the Log struct field definitions (type:text for JSON-serialized arrays)
core/go.mod (2)
8-22: LGTM! Dependency updates are incremental.The version bumps for AWS SDK v2 (config, credentials, S3), Bytedance Sonic, MCP library, fasthttp, oauth2, and golang.org/x/text are all patch or minor updates, consistent with the PR's goals of improving serialization performance and maintaining up-to-date dependencies.
27-66: LGTM! Indirect dependency updates are consistent.The indirect dependency version bumps are synchronized across the dependency tree, including the addition of
Masterminds/semver/v3 v3.3.1and updates to various transitive dependencies. All changes are incremental patch/minor versions.plugins/semanticcache/go.mod (1)
18-109: LGTM! Dependency updates synchronized with core module.The version bumps across AWS SDK v2, Bytedance Sonic, OpenAPI tooling, OpenTelemetry, client libraries (qdrant, redis, weaviate, mongo-driver), and golang.org packages are consistent with the broader PR dependency refresh. All updates are incremental patch/minor versions.
plugins/maxim/go.mod (1)
18-109: LGTM! Dependency updates synchronized with core module.The version bumps mirror the dependency refresh pattern across the repository, with incremental updates to AWS SDK v2, Bytedance Sonic, OpenAPI components, OpenTelemetry, and various golang.org/x/* packages. All changes are patch/minor versions.
framework/go.mod (2)
8-16: LGTM! Direct dependency updates are appropriate.The updates to qdrant/go-client, redis/go-redis, weaviate ecosystem, and golang.org/x/crypto are incremental and align with the broader dependency refresh across the repository.
18-116: LGTM! Indirect dependency updates are synchronized.The indirect dependency version bumps across AWS SDK v2, Bytedance Sonic, OpenAPI tooling, OpenTelemetry, and golang.org packages are consistent with updates in other modules. All changes represent incremental patch/minor versions.
plugins/governance/go.mod (1)
17-112: LGTM! Dependency updates synchronized with repository-wide refresh.The version bumps across AWS SDK v2, Bytedance Sonic, OpenAPI components, OpenTelemetry, client libraries, and golang.org packages mirror the consistent dependency refresh pattern throughout the PR. All updates are incremental patch/minor versions.
plugins/logging/go.mod (1)
6-111: LGTM! Dependency updates consistent with PR goals.The version bumps promote Bytedance Sonic to a direct dependency (v1.14.2) and update all related indirect dependencies consistently with the broader repository refresh. All changes are incremental patch/minor versions, supporting the PR's goal of switching to Sonic for improved JSON serialization performance.
transports/go.mod (1)
7-79: Dependency bumps look consistent; rely on tests to catch any behavioral shifts.Patch/minor upgrades for sonic, AWS SDK v2, fasthttp, OpenAPI, vector DB clients, Mongo driver, OTEL, and x/* libs are coherent with the rest of the stack. No obvious API breaks here; just make sure full transport + e2e tests have been run with these versions to surface any regressions.
Also applies to: 90-90, 105-106, 112-113, 116-133
plugins/telemetry/go.mod (1)
9-10: Telemetry dependency versions aligned with core; validate OTEL/exporter flows.Telemetry’s go.mod mirrors the transport/core bumps (fasthttp, AWS SDK v2, sonic, OpenAPI, qdrant, redis, weaviate, Mongo driver, OTEL, x/*, gRPC/protobuf). This alignment is good; just confirm your OTEL pipelines and exporters behave as expected under these versions.
Also applies to: 17-18, 38-39, 45-65, 75-75, 78-79, 88-89, 94-95, 98-103, 104-113
core/schemas/bifrost.go (1)
142-145: Tried-key context keys are well-defined and consistent with usage.
BifrostContextKeyTriedKeyIDs/BifrostContextKeyTriedKeyNamescorrectly document the intended[]stringpayload and “set by bifrost” semantics, matching how the worker and plugins read/write these values.core/bifrost.go (3)
73-74: Key sorter initialization and algorithm selection are correct and fail fast on misconfig.The new
keySorterfield and theInitswitch onconfig.KeySorterAlgorithmcorrectly default toKeySorterWeightedRandom, validate the presence ofCustomKeySorterwhenKeySorterCustomis requested, and error on unknown algorithms. This gives deterministic behavior and clear failures for bad configs.Also applies to: 124-139
331-333: Retry helper’s attempt index threading looks sound and preserves prior semantics.
executeRequestWithRetriesnow passes the currentattemptsvalue into therequestHandler, and all shown call sites (includingListModelsRequeston Lines 331–333) either use it or explicitly ignore it. The loop still:
- Updates
BifrostContextKeyNumberOfRetrieseach iteration.- Applies backoff only on retries (
attempts > 0).- Breaks on success, bifrost-internal errors, or non-retryable conditions.
This matches the previous behavior while enabling attempt-aware handlers.
Also applies to: 2767-2842
2734-2762: Key selection + tried-keys tracking across retries is correct and index-safe.
selectKeyAndUpdateTriedKeyscorrectly:
- Clamps the attempt index into the valid
[0,len(keys)-1]range.- Sets
BifrostContextKeySelectedKeyID/Namefor the current attempt.- For
attempts > 0, appends the previous key’s ID/Name to thetriedKey*slices and writes those slices back to the context, so “tried keys” represent keys attempted before the current one.The worker’s use of this helper in both streaming and non‑streaming retry closures initializes the slices once per request and reuses them across retries, giving ordered tried‑key history without risk of out‑of‑bounds panics.
Also applies to: 2933-2945
ui/lib/types/logs.ts (1)
267-269: LogEntry tried-key fields align with backend schema.Adding
tried_key_ids?: string[]andtried_key_names?: string[]matches the JSON fields exposed fromlogstore.Log(tried_key_ids,tried_key_names) and gives the UI enough type info to render them safely.plugins/telemetry/main.go (4)
9-9: LGTM!The
stringsimport is correctly added to support joining tried key slices for label formatting.
312-313: LGTM!Correctly retrieves tried key slices using the appropriate helper function and constants.
333-334: Implementation note: comma-separated values in labels.While the implementation correctly joins the tried key slices, note that comma-separated values in Prometheus labels make querying more difficult (e.g., filtering for a specific key within the list). This compounds the cardinality concern raised for lines 101-102.
307-321: Inconsistency between governance and bifrost context key usage patterns - clarification needed.Lines 310-316 use defined constants (e.g.,
schemas.BifrostContextKeySelectedKeyID), while lines 307-308 and 318-321 use string literals wrapped inschemas.BifrostContextKey()for governance keys. However, investigation shows that governance context keys (bf-governance-virtual-key-*,bf-governance-team-*,bf-governance-customer-*) are intentionally used as string literals consistently across the codebase (governance/resolver.go, plugins/logging/main.go, plugins/otel/main.go), rather than being defined as constants.This appears to be an architectural pattern: core bifrost framework keys use constants, while governance plugin keys use string literals. Clarify whether this inconsistency in the same code block is intentional or if governance keys should be constantized (and if so, whether they belong in core/schemas/bifrost.go or in the governance plugin).
plugins/otel/go.mod (3)
110-110: The bytedance/sonic v1.14.2 update is stable and suitable for this use case.v1.14.2 has no documented breaking changes. The OTEL plugin's usage—basic
Unmarshal()andMarshal()operations for configuration objects—is unaffected by the known edge-case issues in the sonic repository (JIT bugs, overflow edge cases in complex unmarshalling scenarios, and ast.Node API limitations). The version aligns with the PR's performance improvement objectives.
8-9: No action needed. The gRPC v1.77.0 and Protobuf v1.36.11 updates are compatible with the codebase. The breaking changes in gRPC v1.77.0 (mem.Reader interface and old pick_first load-balancer removal) do not affect the otel plugin, which only uses standard client APIs (grpc.NewClient, grpc.ClientConn, grpc.WithTransportCredentials). Protobuf v1.36.11 introduces no breaking changes, only minor maintenance updates and feature improvements.
92-94: No action needed. The OpenTelemetry v1.39.0 update is compatible with this plugin. The Go version requirement (Go >= 1.24) is already satisfied by Go 1.25.5 in plugins/otel/go.mod. The plugin uses OTEL proto packages for data definitions, which are stable across minor versions and have no breaking changes. The minor version bump from v1.38.x to v1.39.0 maintains backward compatibility.
Pratham-Mishra04
force-pushed
the
12-23-feat_added_responses_mcp_tool_execute_endpoint
branch
from
4a1a08e to
d6c403b
Compare
Pratham-Mishra04
force-pushed
the
12-23-feat_added_retries_across_multiple_keys
branch
from
6e9445e to
d42c851
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
plugins/logging/main.go (1)
59-76: Tried-key propagation through logging looks consistent; verify pool reset covers new fieldsThe logging plugin now:
- Captures
selectedKeyID/Name,triedKeyIDs/Names,virtualKeyID/Name, andnumberOfRetriesfrom the context via the bifrost helpers, and- Threads
TriedKeyIDs/TriedKeyNamesthrough all update paths (updateLogEntryandupdateStreamingLogEntry) alongside the existing key metadata.That gives logstore (and the UI) a complete picture of which keys were attempted vs. ultimately selected in both streaming and non-streaming flows.
Since
LogMessageinstances are pooled, just make sure the reset logic in your pool helpers (e.g.,getLogMessage/putLogMessageinplugins/logging/pool.go) also clearsTriedKeyIDsandTriedKeyNames(e.g., by zeroing the struct or explicitly nil-ing the slices) so future reuses can’t accidentally retain old slice contents.Also applies to: 223-300, 351-388, 403-421, 447-463, 580-595
core/bifrost.go (1)
73-74: Retry-aware key ordering and tried-key context semantics look sound; ensure keySorter is always set via InitThe core wiring for multi-key retries is well put together:
Initenforces a valid key sorter choice (defaulting to WeightedRandom) and assignsbifrost.keySorter, soselectKeyFromProviderForModelcan rely on it.selectKeyFromProviderForModelreturns a sorted slice of supported keys (respecting model, deployment, and batch/file constraints), and honors direct-key and key-name overrides before falling back to the sorter.executeRequestWithRetriesthreads the current attempt index through to the request handler and annotates the context withbifrost-number-of-retries, which downstream plugins consume.selectKeyAndUpdateTriedKeysuses that attempt index to:
- Choose the current key (with clamping once you run out of distinct keys), and
- Accumulate previously used keys into
tried_key_ids/tried_key_namesin the context, so the final context state cleanly separates “selected key” from “prior attempts”.Given this, telemetry/logging/logstore/UI should all see consistent selected vs. tried-key data across retries and fallbacks.
The only thing to keep in mind going forward is that
selectKeyFromProviderForModelassumesbifrost.keySorteris non-nil; today that’s guaranteed because all production construction flows go throughInit, but if you ever add alternate constructors or extendReloadConfigto touch key-sorting behavior, make sure they also set or updatekeySorteraccordingly to avoid panics.Also applies to: 124-139, 2734-2762, 2767-2804, 2868-2917, 2932-2945, 3334-3441, 3443-3558
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (11)
core/go.sumis excluded by!**/*.sumframework/go.sumis excluded by!**/*.sumplugins/governance/go.sumis excluded by!**/*.sumplugins/jsonparser/go.sumis excluded by!**/*.sumplugins/logging/go.sumis excluded by!**/*.sumplugins/maxim/go.sumis excluded by!**/*.sumplugins/mocker/go.sumis excluded by!**/*.sumplugins/otel/go.sumis excluded by!**/*.sumplugins/semanticcache/go.sumis excluded by!**/*.sumplugins/telemetry/go.sumis excluded by!**/*.sumtransports/go.sumis excluded by!**/*.sum
📒 Files selected for processing (29)
core/bifrost.gocore/bifrost_test.gocore/go.modcore/key-sorting-algos/weighted-random.gocore/schemas/bifrost.gocore/schemas/context.gocore/utils.goframework/go.modframework/logstore/migrations.goframework/logstore/tables.goplugins/governance/go.modplugins/jsonparser/go.modplugins/logging/go.modplugins/logging/main.goplugins/logging/operations.goplugins/logging/pool.goplugins/logging/utils.goplugins/maxim/go.modplugins/mocker/go.modplugins/otel/converter.goplugins/otel/go.modplugins/otel/main.goplugins/semanticcache/go.modplugins/telemetry/go.modplugins/telemetry/main.goplugins/telemetry/utils.gotransports/go.modui/app/workspace/logs/views/logDetailsSheet.tsxui/lib/types/logs.ts
💤 Files with no reviewable changes (2)
- plugins/telemetry/utils.go
- plugins/logging/utils.go
🚧 Files skipped from review as they are similar to previous changes (11)
- framework/logstore/migrations.go
- plugins/logging/pool.go
- plugins/mocker/go.mod
- ui/lib/types/logs.ts
- core/schemas/context.go
- core/schemas/bifrost.go
- plugins/logging/operations.go
- plugins/jsonparser/go.mod
- core/bifrost_test.go
- core/go.mod
- plugins/maxim/go.mod
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/otel/main.goframework/logstore/tables.gocore/key-sorting-algos/weighted-random.goplugins/governance/go.modplugins/telemetry/main.goplugins/otel/converter.gocore/bifrost.goplugins/telemetry/go.modcore/utils.goui/app/workspace/logs/views/logDetailsSheet.tsxplugins/logging/main.goplugins/otel/go.modframework/go.modtransports/go.modplugins/logging/go.modplugins/semanticcache/go.mod
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/otel/main.goframework/logstore/tables.gocore/key-sorting-algos/weighted-random.goplugins/telemetry/main.goplugins/otel/converter.gocore/bifrost.gocore/utils.goplugins/logging/main.go
🧬 Code graph analysis (4)
plugins/otel/main.go (2)
core/utils.go (1)
GetStringSliceFromContext(303-310)core/schemas/bifrost.go (2)
BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)
core/key-sorting-algos/weighted-random.go (2)
core/schemas/account.go (1)
Key(8-20)core/schemas/bifrost.go (1)
ModelProvider(46-46)
plugins/telemetry/main.go (2)
core/utils.go (3)
GetStringFromContext(293-300)GetStringSliceFromContext(303-310)GetIntFromContext(313-320)core/schemas/bifrost.go (7)
BifrostContextKey(131-131)BifrostContextKeySelectedKeyID(140-140)BifrostContextKeySelectedKeyName(141-141)BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)BifrostContextKeyNumberOfRetries(144-144)BifrostContextKeyFallbackIndex(145-145)
ui/app/workspace/logs/views/logDetailsSheet.tsx (1)
ui/app/workspace/logs/views/logEntryDetailsView.tsx (1)
LogEntryDetailsView(15-49)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (13)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (11)
plugins/semanticcache/go.mod (3)
38-39: Bytedance/sonic v1.14.2 is stable and recommended; verify stack-wide compatibility.The sonic library upgrade to v1.14.2 (released 2025-10-28) is a maintenance patch with no documented breaking changes. The sonic/loader v0.4.0 bump is the matching version from that release. However, per coding guidelines, verify that these dependency versions maintain compatibility across other modules in the PR stack and align with the serialization refactor implementation.
96-109: Verify stdlib and google.golang.org module versions align with the PR stack.Multiple stdlib modules (golang.org/x/crypto, net, oauth2, sync, sys, text) and google.golang.org modules (genproto, grpc, protobuf) have been updated. General compatibility checks indicate that google.golang.org/grpc v1.77.0 and golang.org/x/sync v0.19.0 are compatible, and go.opentelemetry.io/otel v1.39.0 is compatible with gRPC v1.77.0. However, verify these specific versions are compatible with the broader
bifrost/coreandbifrost/frameworkmodules in your stack, particularly for concurrency, networking, and gRPC instrumentation handling.
18-19: AWS SDK v2 module updates are safe — no breaking changes or security issues identified.Config v1.32.6, credentials v1.19.6, s3 v1.94.0, and sso v1.30.8 are routine patch and feature releases with no breaking changes. These updates pose no risk to the retry/key handling refactor.
plugins/otel/main.go (1)
252-253: LGTM! Clean integration of tried-key tracking.The implementation correctly extracts tried key IDs and names from context and propagates them to OTEL spans for both streaming and non-streaming code paths. The pattern mirrors the existing virtual/selected key handling, ensuring consistency.
Also applies to: 292-293, 317-318
core/utils.go (2)
302-310: LGTM! Safe context extraction utility.The
GetStringSliceFromContextfunction follows the established pattern of other context extraction utilities (GetStringFromContext,GetIntFromContext) and safely handles missing or mistyped values by returning an empty slice.
410-418: LGTM! Standard utility function.The
clampfunction correctly constrains an integer value to the inclusive range [min, max]. Implementation is straightforward and correct.plugins/telemetry/go.mod (1)
9-113: No further action needed. The listed dependency versions (sonic v1.14.2, fasthttp v1.68.0, OpenTelemetry v1.39.0, etc.) have no known public security vulnerabilities. Standard practice of using recent, stable versions already mitigates security risk.ui/app/workspace/logs/views/logDetailsSheet.tsx (1)
271-277: Tried-keys UI wiring looks correctConditional rendering and simple comma-joined formatting for
tried_key_ids/tried_key_namesis consistent with the rest of the details view and should surface the new observability fields as intended.framework/logstore/tables.go (1)
86-88: Tried-key persistence and sonic-based (de)serialization look consistentThe new
TriedKeyIDs/TriedKeyNamesstring columns plus their[]stringparsed counterparts integrate cleanly:
- SerializeFields marshals the parsed slices only when non-nil and writes JSON strings into the DB fields.
- DeserializeFields unmarshals when the DB fields are non-empty and safely falls back to empty slices on error.
- Patterns match the existing JSON-backed fields (InputHistory, Tools, CacheDebug), keeping the model’s behavior predictable.
No issues from a storage/round-trip perspective.
Also applies to: 136-137, 169-307, 317-435
plugins/otel/converter.go (1)
640-652: OTEL span/resource enrichment for tried keys is wired correctlyExtending
completeResourceSpanwithtriedKeyIDs/triedKeyNamesand emitting them both asgen_ai.*span attributes and bare resource attributes (when non-empty) aligns with the existing virtual/selected-key metadata pattern and should make tried-key sequences visible in traces without disrupting existing consumers.Also applies to: 1068-1073, 1088-1101
plugins/telemetry/main.go (1)
9-9: Prometheus tried-key labels are correctly sourced from contextAdding
tried_key_ids/tried_key_namestodefaultBifrostLabelsand populating them viaGetStringSliceFromContext+strings.Joinin PostHook cleanly plugs tried-key sequences into the upstream metrics without changing the existing label ordering or helper logic. For non-retry paths the joined value will be empty, which is acceptable for label values.Also applies to: 92-109, 307-321, 325-341
| // Compute scores; handle w<=0 cheaply | ||
| // w<0 is usually a config bug; if you truly want speed, you can drop this check. | ||
| for i := 0; i < n; i++ { | ||
| w := keys[i].Weight | ||
| if w < 0 { | ||
| return nil, fmt.Errorf("key %q has negative weight %f", keys[i].ID, w) | ||
| } | ||
| if w == 0 { | ||
| scores[i] = math.Inf(-1) // push to the end | ||
| continue | ||
| } | ||
|
|
||
| // u in (0,1]; ensure non-zero so log(u) is finite. | ||
| u := r.Float64() | ||
| for u == 0 { | ||
| u = r.Float64() | ||
| } | ||
| scores[i] = math.Log(u) / w | ||
| } | ||
|
|
||
| // Sort keys in-place by score descending (higher score first) | ||
| sort.Slice(keys, func(i, j int) bool { | ||
| return scores[i] > scores[j] | ||
| }) | ||
|
|
||
| // Optional compromise: uniformly shuffle the tail of zero-weight keys | ||
| // so they don't always appear in the same relative order. | ||
| // Find first -Inf (zero-weight) position. | ||
| z := n | ||
| for i := 0; i < n; i++ { | ||
| if math.IsInf(scores[i], -1) { | ||
| z = i | ||
| break | ||
| } | ||
| } | ||
| // Fisher–Yates on keys[z:] | ||
| for i := n - 1; i > z; i-- { | ||
| j := z + r.Intn(i-z+1) | ||
| keys[i], keys[j] = keys[j], keys[i] | ||
| } | ||
|
|
There was a problem hiding this comment.
Bug in weighted sort: keys and scores get out of sync during sort.Slice
sort.Slice(keys, func(i, j int) bool { return scores[i] > scores[j] }) reorders keys but never swaps the corresponding entries in scores. After the first swap, indices no longer refer to the same key/score pair, so:
- The ordering is no longer by the per-key scores you computed (the comparator depends only on fixed index-based
scores[i]/scores[j]), and - The zero-weight tail detection (
math.IsInf(scores[i], -1)) no longer lines up with the keys slice.
This breaks the intended weighted-random permutation and can also misidentify which keys are in the zero-weight tail.
You should sort keys and scores together so they stay aligned. For example:
Proposed fix: custom sorter that swaps both keys and scores
@@
- // Sort keys in-place by score descending (higher score first)
- sort.Slice(keys, func(i, j int) bool {
- return scores[i] > scores[j]
- })
+ // Sort keys and scores in-place by score descending (higher score first)
+ type keyedScores struct {
+ keys []schemas.Key
+ scores []float64
+ }
+
+ sort.Sort(keyedScores{
+ keys: keys,
+ scores: scores,
+ })
@@
- // Find first -Inf (zero-weight) position.
+ // Find first -Inf (zero-weight) position. This now correctly lines up
+ // with the sorted keys because scores are swapped alongside keys.and define the sort interface just above WeightedRandomKeySorter:
type keyedScores struct {
keys []schemas.Key
scores []float64
}
func (s keyedScores) Len() int { return len(s.keys) }
func (s keyedScores) Less(i, j int) bool {
return s.scores[i] > s.scores[j] // higher score first
}
func (s keyedScores) Swap(i, j int) {
s.keys[i], s.keys[j] = s.keys[j], s.keys[i]
s.scores[i], s.scores[j] = s.scores[j], s.scores[i]
}This keeps keys and scores aligned for both the sort comparator and the subsequent zero-weight-tail scan.
🤖 Prompt for AI Agents
In core/key-sorting-algos/weighted-random.go around lines 63 to 103, the
sort.Slice call only reorders keys but not the parallel scores slice, which
desynchronizes key-score pairs and breaks the zero-weight tail detection; fix by
sorting keys and scores together (e.g., create a small keyedScores struct
holding keys[] and scores[] and implement Len, Less (scores[i]>scores[j]) and
Swap that swaps both keys and scores), then call sort.Sort on that keyedScores
so the score ordering stays aligned with keys and the subsequent
scan/Fisher–Yates on the zero-weight tail operates on the correct indices.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/maxim/go.mod (1)
1-114: Verify that upstream dependency versions (core v1.2.40, framework v1.1.50) are intentional for this PR.The ~50 dependency updates in this file are transitive cascades from
core v1.2.40andframework v1.1.50, not plugin-specific updates. While the Bytedance Sonic dependencies appear here, the maxim plugin itself still uses standard libraryencoding/jsonrather than Sonic for JSON serialization.Confirm that the upstream version bumps are intentional for this "retries across multiple keys" feature, particularly since this PR is part of a larger stack.
🧹 Nitpick comments (9)
core/key-sorting-algos/weighted-random.go (1)
75-81: Consider adding a comment about the u≠0 loop.While the probability of
Float64()returning exactly 0.0 repeatedly is astronomically low, thefor u == 0loop is theoretically unbounded. Consider adding a comment explaining why this is safe (finite precision, vanishingly small probability).🔎 Suggested comment
// u in (0,1]; ensure non-zero so log(u) is finite. + // Float64() returns [0.0, 1.0), probability of exact 0.0 is ~2^-53 u := r.Float64() for u == 0 { u = r.Float64() }plugins/logging/main.go (1)
368-387: End‑to‑end tried‑key propagation into log updates is wired correctlyContext extraction of selected/tried/virtual key info and passing
TriedKeyIDs/TriedKeyNamesthrough all update paths (error, streaming final, regular) cleanly feeds the new DB columns and UI without changing existing behavior.If you ever see hot-path pressure here, you could reuse a single temporary slice copy for
TriedKeyIDs/TriedKeyNamesto decouple from the context slice and reduce chances of holding onto larger backing arrays via the pool, but it's not required for correctness.Also applies to: 413-415, 455-457, 586-587
framework/logstore/tables.go (1)
86-87: Tried‑key persistence via JSON string columns is well‑integrated
- Adding
TriedKeyIDs/TriedKeyNamesstring columns plus their[]stringparsed counterparts follows the same virtual/serialized pattern as other complex fields.- SerializeFields/DeserializeFields for the new fields mirror existing logic (marshal when Parsed is non‑nil; unmarshal when column is non‑empty), so DB <→ in‑memory representation stays symmetric.
If deserialization of tried‑key JSON becomes noisy in production, consider adding lightweight structured logging around those specific unmarshal failures, similar to how higher layers log cache/debug issues, but current “fail open with empty slice” behavior is acceptable.
Also applies to: 136-137, 169-176, 177-184, 185-192, 193-200, 201-208, 209-216, 217-224, 225-232, 233-240, 241-248, 249-256, 257-264, 265-276, 277-284, 285-292, 293-308, 317-323, 324-330, 331-337, 338-344, 345-351, 352-358, 359-365, 366-372, 373-379, 380-386, 387-393, 395-401, 402-408, 409-415, 416-422, 423-435
core/schemas/bifrost.go (1)
16-29: KeySorter types are coherent with intended usageThe
KeySorterFuncandKeySorterabstractions match how key sorting is wired incore/bifrost.go(ctx pointer + key slice + provider/model), and the constants give a clear enum‑style surface for config.If you ever broaden this API, consider documenting explicitly that
ctxmay be mutated by sorters (e.g., to attach diagnostics), since the pointer receiver implies that but the comment currently only discusses ordering, not side effects.plugins/logging/operations.go (2)
48-83: Non‑stream update path correctly persists tried keysPassing
triedKeyIDs/triedKeyNamesintoupdateLogEntryand serializing them via a temporarylogstore.Logkeeps DB representation consistent with the main model; failures are logged and scoped to those fields only.You could reuse the later
tempEntryinstance instead of allocating separate ones for tried‑key IDs and names, but the current approach is straightforward and not performance‑critical.
205-240: Streaming update path mirrors tried‑key handling from regular updatesThe streaming update function handles
triedKeyIDs/triedKeyNameswith the same JSON‑through‑Log pattern as the non‑stream path, ensuring consistent persistence for both response shapes.Same as above: consider reusing a single
tempEntryfor tried‑key serialization to trim a couple of allocations per call if this turns hot.core/bifrost.go (3)
17-18: Key sorter wiring and defaulting in Init look solid
Bifrostnow holds akeySorterfunction, chosen at init time based onKeySorterAlgorithmwith a sensible default ofweighted_random.- Custom algorithms are gated on
CustomKeySorter != nil, and unsupported algorithm names fail fast during Init, which is preferable to latent panics.If you anticipate hot‑reloading of key‑sorting behavior in future, consider teaching
ReloadConfigto updatekeySorteras well, or documenting that sorter choice is fixed at Init.Also applies to: 55-74, 124-140
2734-2762: selectKeyAndUpdateTriedKeys correctly tracks selected and previously tried keys
- For each attempt, the helper sets
BifrostContextKeySelectedKeyID/Namefrom the sortedkeysslice and, for attempts > 0, appends the previous key to the tried‑key slices and writes them to context.- This yields
triedKeyIDs/namesas “keys attempted before the current one,” matching the intent to surface fallback history in logging/OTEL.- The usage in
requestWorkerfor both streaming and non‑streaming flows is race‑free (single goroutine per request) and bounded byMaxRetries.If you ever want
tried_key_*to represent all attempted keys including the final one, you could additionally appendselectedKeyafter the loop terminates; current semantics (only previous keys) are internally consistent and may already match your observability needs.Also applies to: 2932-2945
3445-3559: selectKeyFromProviderForModel: slice return + pluggable sorter behave as intended
- The method now:
- Handles direct keys and skip‑selection flags up front.
- Filters and validates keys (enabled, batch/file flags, model/deployment support) as before.
- Returns a
[]Keyordered bybifrost.keySorter, with direct and key‑name selection short‑circuiting to single‑element slices.- Error messages remain specific for “no keys” vs “no batch‑enabled keys” and Azure/Bedrock/Vertex deployment mismatches.
For batch/file multi‑key operations, you might want to skip invoking the key sorter entirely and rely purely on deterministic ID sorting (or document that the sorter’s ordering is also applied there), to avoid confusion about the meaning of “selected” vs “used” keys in those flows.
📜 Review details
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (11)
core/go.sumis excluded by!**/*.sumframework/go.sumis excluded by!**/*.sumplugins/governance/go.sumis excluded by!**/*.sumplugins/jsonparser/go.sumis excluded by!**/*.sumplugins/logging/go.sumis excluded by!**/*.sumplugins/maxim/go.sumis excluded by!**/*.sumplugins/mocker/go.sumis excluded by!**/*.sumplugins/otel/go.sumis excluded by!**/*.sumplugins/semanticcache/go.sumis excluded by!**/*.sumplugins/telemetry/go.sumis excluded by!**/*.sumtransports/go.sumis excluded by!**/*.sum
📒 Files selected for processing (29)
core/bifrost.gocore/bifrost_test.gocore/go.modcore/key-sorting-algos/weighted-random.gocore/schemas/bifrost.gocore/schemas/context.gocore/utils.goframework/go.modframework/logstore/migrations.goframework/logstore/tables.goplugins/governance/go.modplugins/jsonparser/go.modplugins/logging/go.modplugins/logging/main.goplugins/logging/operations.goplugins/logging/pool.goplugins/logging/utils.goplugins/maxim/go.modplugins/mocker/go.modplugins/otel/converter.goplugins/otel/go.modplugins/otel/main.goplugins/semanticcache/go.modplugins/telemetry/go.modplugins/telemetry/main.goplugins/telemetry/utils.gotransports/go.modui/app/workspace/logs/views/logDetailsSheet.tsxui/lib/types/logs.ts
💤 Files with no reviewable changes (2)
- plugins/logging/utils.go
- plugins/telemetry/utils.go
🚧 Files skipped from review as they are similar to previous changes (9)
- plugins/logging/pool.go
- core/bifrost_test.go
- core/schemas/context.go
- ui/app/workspace/logs/views/logDetailsSheet.tsx
- framework/logstore/migrations.go
- plugins/jsonparser/go.mod
- plugins/telemetry/main.go
- plugins/governance/go.mod
- plugins/mocker/go.mod
🧰 Additional context used
📓 Path-based instructions (1)
**
⚙️ CodeRabbit configuration file
always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)
Files:
plugins/logging/main.goplugins/otel/converter.goplugins/otel/main.goplugins/logging/operations.gocore/utils.goframework/go.modplugins/otel/go.modplugins/maxim/go.modframework/logstore/tables.goplugins/telemetry/go.modplugins/logging/go.modui/lib/types/logs.tscore/bifrost.gocore/go.modplugins/semanticcache/go.modtransports/go.modcore/schemas/bifrost.gocore/key-sorting-algos/weighted-random.go
🧠 Learnings (1)
📚 Learning: 2025-12-09T17:07:42.007Z
Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.
Applied to files:
plugins/logging/main.goplugins/otel/converter.goplugins/otel/main.goplugins/logging/operations.gocore/utils.goframework/logstore/tables.gocore/bifrost.gocore/schemas/bifrost.gocore/key-sorting-algos/weighted-random.go
🧬 Code graph analysis (6)
plugins/logging/main.go (2)
core/schemas/bifrost.go (7)
BifrostContextKeyFallbackIndex(145-145)BifrostContextKeySelectedKeyID(140-140)BifrostContextKeySelectedKeyName(141-141)BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)BifrostContextKey(131-131)BifrostContextKeyNumberOfRetries(144-144)core/utils.go (3)
GetIntFromContext(313-320)GetStringFromContext(293-300)GetStringSliceFromContext(303-310)
plugins/otel/main.go (2)
core/utils.go (1)
GetStringSliceFromContext(303-310)core/schemas/bifrost.go (2)
BifrostContextKeyTriedKeyIDs(142-142)BifrostContextKeyTriedKeyNames(143-143)
plugins/logging/operations.go (1)
framework/logstore/tables.go (2)
Log(75-142)Log(145-147)
framework/logstore/tables.go (2)
ui/lib/types/logs.ts (5)
SpeechInput(11-13)TranscriptionInput(15-17)CacheDebug(213-222)ChatMessage(108-118)ResponsesMessage(441-457)core/schemas/chatcompletions.go (1)
ChatMessage(478-487)
core/schemas/bifrost.go (2)
core/schemas/account.go (1)
Key(8-20)core/schemas/mcp.go (1)
MCPConfig(14-23)
core/key-sorting-algos/weighted-random.go (2)
core/schemas/account.go (1)
Key(8-20)core/schemas/bifrost.go (1)
ModelProvider(46-46)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (20)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (22)
plugins/semanticcache/go.mod (1)
18-114: All indirect dependencies are transitive from upstream modules; no direct changes required.All ~100 indirect dependencies in this go.mod, including bytedance/sonic, are transitively pulled from the 5 direct dependencies (bifrost/core, bifrost/framework, bifrost/plugins/mocker). The plugin itself does not directly import sonic, so these updates are expected to flow through when the upstream modules are bumped in the stack PRs (#1160, #1161). This is standard Go module behavior in a multi-module project structure.
transports/go.mod (5)
12-19: Verify bifrost internal module versions align with PR stack.The internal bifrost module versions (core, framework, plugins) should be consistent across all PRs in the stack (#1160, #1161, #1163). Confirm that the versions specified here match the expected state after all PRs in the stack are merged.
The current file specifies:
- github.com/maximhq/bifrost/core v1.2.40
- github.com/maximhq/bifrost/framework v1.1.50
- github.com/maximhq/bifrost/plugins/governance v1.3.51
- github.com/maximhq/bifrost/plugins/logging v1.3.51
- github.com/maximhq/bifrost/plugins/maxim v1.4.51
- github.com/maximhq/bifrost/plugins/otel v1.0.50
- github.com/maximhq/bifrost/plugins/semanticcache v1.3.50
- github.com/maximhq/bifrost/plugins/telemetry v1.3.50
Verify these versions are the expected outputs from the downstack PRs in this stack.
59-79: Verify go-openapi component compatibility.The transports/go.mod file includes go-openapi dependencies across multiple versions: strfmt (v0.25.0), swag with submodules (v0.25.4), and validate (v0.25.1) are in the v0.25.x series, while analysis (v0.24.2), errors (v0.22.5), jsonpointer (v0.22.4), jsonreference (v0.21.4), loads (v0.23.2), runtime (v0.29.2), and spec (v0.22.2) are at various earlier versions. This wide version spread across go-openapi modules should be verified for compatibility.
Likely an incorrect or invalid review comment.
106-106: No documented breaking changes exist between v9.14.0 and v9.17.2. The version jump introduces new commands, performance/pool fixes, and Redis 8.4 support without public API breaking changes. Note that v9.17.2 adds typed errors which may affect error handling if retry logic relies on exact error message strings.
7-7: bytedance/sonic v1.14.2 is a safe upgrade.The patch bump from v1.14.1 contains only bug fixes (encoder recursion, decoder panic handling, validate skip alignment) and internal optimizations, with no breaking changes. The upgrade aligns with the PR's JSON serialization switch.
11-11: No breaking changes identified in this upgrade.The jump from v0.41.1 to v0.43.2 introduces new optional features (Session-specific resources, Roots support, custom HTTP headers) and robustness improvements (bufio.Scanner → bufio.Reader for large message handling), but no breaking changes. The bifrost-http transport code uses only stable APIs (server.NewMCPServer, AddTool, HandleMessage, and result helpers) that remain unchanged between these versions and are fully compatible.
Likely an incorrect or invalid review comment.
core/utils.go (2)
302-310: LGTM! Consistent with existing context helpers.The implementation follows the same pattern as
GetStringFromContextandGetIntFromContext. Returning an empty slice instead ofnilis a reasonable default that prevents nil-pointer issues in range loops.
410-418: LGTM! Simple and correct bounds enforcement.The
clampfunction correctly enforces inclusive bounds. Using an unexported function is appropriate for an internal utility.core/key-sorting-algos/weighted-random.go (2)
63-81: LGTM! Score computation and zero-weight handling are correct.The weighted random sampling algorithm using
log(u)/wis mathematically sound and correctly prioritizes higher-weighted keys. Zero-weight keys are appropriately pushed to the end with-Infscores.
18-31: Excellent use of sync.Pool for performance optimization.The pooling of RNG instances and score buffers minimizes allocations and GC pressure, which is critical for a hot path like key selection.
framework/go.mod (1)
1-117: Dependency updates align with ecosystem upgrades.These version bumps are consistent with the updates in
core/go.modand support the broader changes in the PR (key sorting, tried-key context propagation, OTEL span enrichment).plugins/otel/go.mod (1)
1-118: LGTM! OTEL and protobuf updates support span enrichment.The gRPC and protobuf version bumps align with the OTEL span enrichment changes mentioned in the AI summary (adding tried_key_ids and tried_key_names to resource spans).
plugins/telemetry/go.mod (1)
1-118: LGTM! Dependency updates support telemetry enhancements.These updates align with the telemetry changes in the PR that propagate tried-key context through metrics and labels.
plugins/logging/go.mod (1)
1-112: LGTM! Updates support logging plugin enhancements.The sonic upgrade is particularly relevant given the PR's switch from
encoding/jsontobytedance/sonicfor improved serialization performance.ui/lib/types/logs.ts (1)
267-268: LGTM! Schema extension for tried-key observability.The new optional fields align with the backend changes mentioned in the AI summary (framework/logstore/tables.go, plugins/logging/main.go) and enable displaying tried keys in the UI log details sheet.
core/go.mod (1)
8-23: No compatibility issues found with updated dependencies.Verification confirms that bytedance/sonic v1.14.2, mark3labs/mcp-go v0.43.2, and valyala/fasthttp v1.68.0 contain no breaking changes. These versions include only backwards-compatible bug fixes (encoder/decoder improvements in sonic, notification handling and JSON unmarshalling fixes in mcp-go, and named return bug fixes in fasthttp) and are safe to integrate.
plugins/otel/main.go (1)
252-254: Tried-key context threading into OTEL spans looks correctReading triedKeyIDs/names via
GetStringSliceFromContextand passing them intocompleteResourceSpanin both streaming-final and non-stream paths is type-safe and matches the new context keys; no additional changes needed here.Also applies to: 282-295, 307-325
plugins/logging/main.go (1)
66-67: Logging struct/context extraction changes are consistent
- Adding
TriedKeyIDs/TriedKeyNamestoLogMessageand populatingTranscriptionInputplusFallbackIndexvia helpers keeps logging data model aligned with new retry semantics.- No behavioral risks spotted; pool reuse remains safe as fields are overwritten per use.
Also applies to: 279-279, 295-300
plugins/otel/converter.go (1)
640-652: Tried‑key OTEL attributes are added consistentlyThe new
triedKeyIDs/triedKeyNamesparameters and their emission as span and resource attributes are conditionally applied (len>0) and follow the existing pattern of comma-joining multi-valued fields; looks good.Also applies to: 1068-1073, 1096-1101
core/schemas/bifrost.go (1)
38-43: New tried‑key context keys and config fields fit the existing schema
- Adding
KeySorterAlgorithm/CustomKeySortertoBifrostConfigis backward compatible (zero value → default algorithm) and aligns with the new sorter plumbing.BifrostContextKeyTriedKeyIDs/BifrostContextKeyTriedKeyNamesare clearly documented as[]stringfor “keys tried in order” and match how plugins fetch them.Also applies to: 135-145
core/bifrost.go (2)
331-334: Retry handler signature change preserves behavior while exposing attempt indexUpdating
executeRequestWithRetriesto passattemptsintorequestHandler(and adjusting the ListModels path accordingly) keeps the existing retry/backoff behavior while allowing callers to make attempt‑aware decisions; no issues spotted.Also applies to: 2768-2804
2868-2917: Key retrieval and error handling around selectKeyFromProviderForModel remain safeSwitching to a slice‑returning
selectKeyFromProviderForModelfor single‑key ops preserves the earlier validation behavior (no keys → error, batch/file special‑casing) while supplying a deterministic, sorted key list for retry selection; error propagation toreq.Erris unchanged.
| func WeightedRandomKeySorter( | ||
| _ *context.Context, | ||
| keys []schemas.Key, | ||
| _ schemas.ModelProvider, | ||
| _ string, | ||
| ) ([]schemas.Key, error) { |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Pass context by value, not pointer.
Go convention is to pass context.Context by value, not by pointer. The context package is designed to be passed by value and is safe for concurrent use.
🔎 Proposed fix
func WeightedRandomKeySorter(
- _ *context.Context,
+ _ context.Context,
keys []schemas.Key,
_ schemas.ModelProvider,
_ string,
) ([]schemas.Key, error) {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func WeightedRandomKeySorter( | |
| _ *context.Context, | |
| keys []schemas.Key, | |
| _ schemas.ModelProvider, | |
| _ string, | |
| ) ([]schemas.Key, error) { | |
| func WeightedRandomKeySorter( | |
| _ context.Context, | |
| keys []schemas.Key, | |
| _ schemas.ModelProvider, | |
| _ string, | |
| ) ([]schemas.Key, error) { |
🤖 Prompt for AI Agents
In core/key-sorting-algos/weighted-random.go around lines 35 to 40, the function
currently accepts a *context.Context pointer; change the signature to accept
context.Context by value (e.g., ctx context.Context) and update any internal
uses from dereferencing the pointer to using ctx directly; also update all call
sites to pass the context value (not &ctx) and remove any nil pointer checks
related to the context; ensure the package imports "context" if not already.
Summary
Implement a more efficient key selection algorithm with improved retry handling. This PR replaces the previous random key selection with a sorted key approach, allowing for more deterministic fallback behavior when retrying requests.
Changes
KeySelectorwithKeySorterFuncthat returns keys in sorted order rather than a single keytried_key_idsandtried_key_namescolumnsencoding/jsontobytedance/sonicfor better serialization performanceType of change
Affected areas
How to test
Test key selection and retry behavior with multiple keys:
Breaking changes
The
KeySelectorinterface has been replaced withKeySorterFuncand the configuration has changed. Custom key selectors will need to be updated to the new interface.Security considerations
No security implications. This change only affects how API keys are selected and retried.
Checklist