fix: filter list models by virtual key providers#3187
fix: filter list models by virtual key providers#3187princepal9120 wants to merge 3 commits intomaximhq:mainfrom
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
✅ Files skipped from review due to trivial changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds virtual-key-aware provider filtering: a new request helper resolves virtual keys to allowed providers and stores them in request context; the HTTP list-models handler short-circuits using that filter; core ListAllModels consults context via a new unexported helper to restrict which providers are queried; tests added for the helper. ChangesVirtual-key → provider filtering
Sequence DiagramsequenceDiagram
participant Client
participant HTTP as HTTP Handler
participant VK as VK Resolver (list_models_vk)
participant Config as Config Store
participant Context as Request Context
participant Core as Bifrost Core (ListAllModels)
Client->>HTTP: GET /models (no provider param, bearer token)
HTTP->>VK: Check for virtual key
VK->>Config: GetVirtualKeyByValue(token)
Config-->>VK: Virtual key config (providers)
VK->>Context: Set BifrostContextKeyAvailableProviders
VK-->>HTTP: Return (filter applied)
HTTP->>Core: ListAllModels(ctx)
Core->>Context: Read BifrostContextKeyAvailableProviders
Core->>Core: filterProvidersByContext()
Core->>Core: Query only filtered providers
Core-->>HTTP: Aggregated models
HTTP-->>Client: Return filtered models
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly Related Issues
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.11.4)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Review rate limit: 5/8 reviews remaining, refill in 19 minutes and 11 seconds.Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/bifrost.go (1)
611-614: ⚡ Quick winHarden malformed
available-providershandling to fail closed.At Line 611–614, if the context value is present but wrongly typed, the code falls back to
providerKeys(all configured providers). That re-opens broad fan-out in malformed-context cases. Prefer fail-closed when the key is present but invalid.Suggested patch
func filterProvidersByContext(ctx *schemas.BifrostContext, providerKeys []schemas.ModelProvider) []schemas.ModelProvider { if ctx == nil { return providerKeys } - availableProviders, ok := ctx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - if !ok { + rawAvailableProviders := ctx.Value(schemas.BifrostContextKeyAvailableProviders) + if rawAvailableProviders == nil { return providerKeys } + availableProviders, ok := rawAvailableProviders.([]schemas.ModelProvider) + if !ok { + return []schemas.ModelProvider{} + } if len(availableProviders) == 0 || len(providerKeys) == 0 { return []schemas.ModelProvider{} }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/bifrost.go` around lines 611 - 614, When reading ctx.Value(schemas.BifrostContextKeyAvailableProviders), change the logic to distinguish "key absent" vs "key present but wrong type": first assign raw := ctx.Value(...); if raw == nil keep the existing behavior and return providerKeys; if raw is non-nil but the type assertion to []schemas.ModelProvider fails (availableProviders, ok := raw.([]schemas.ModelProvider)), return a fail-closed value (e.g., nil or empty slice) instead of falling back to providerKeys. Update the handling around availableProviders and providerKeys accordingly so malformed context values do not reopen broad fan-out.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@core/bifrost.go`:
- Around line 611-614: When reading
ctx.Value(schemas.BifrostContextKeyAvailableProviders), change the logic to
distinguish "key absent" vs "key present but wrong type": first assign raw :=
ctx.Value(...); if raw == nil keep the existing behavior and return
providerKeys; if raw is non-nil but the type assertion to
[]schemas.ModelProvider fails (availableProviders, ok :=
raw.([]schemas.ModelProvider)), return a fail-closed value (e.g., nil or empty
slice) instead of falling back to providerKeys. Update the handling around
availableProviders and providerKeys accordingly so malformed context values do
not reopen broad fan-out.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e785d1a4-381e-499b-b7b4-6a2d0c1bb632
📒 Files selected for processing (4)
core/bifrost.gocore/bifrost_test.gotransports/bifrost-http/handlers/inference.gotransports/bifrost-http/handlers/list_models_vk.go
Confidence Score: 4/5Safe to merge with awareness of the unresolved Core filtering logic in transports/bifrost-http/handlers/list_models_vk.go — Important Files Changed
Reviews (3): Last reviewed commit: "docs: clarify provider filter context ow..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| availableProviders = append(availableProviders, schemas.ModelProvider(provider)) | ||
| } | ||
|
|
||
| bifrostCtx.SetValue(schemas.BifrostContextKeyAvailableProviders, availableProviders) |
There was a problem hiding this comment.
schemas.BifrostContextKeyAvailableProviders is annotated in the schema as // []ModelProvider (set by bifrost - DO NOT SET THIS MANUALLY). This PR sets it directly in a transport-layer handler, mirroring an identical pattern already present in transports/bifrost-http/integrations/router.go:775. The comment appears to be aimed at external callers, but it's misleading for internal components and could cause confusion for future contributors. Consider either updating the schema comment to clarify that internal bifrost components may set this key, or establishing a setter helper that makes the intent explicit.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| func (h *CompletionHandler) applyListModelsVirtualKeyProviderFilter(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext) bool { | ||
| vkValue := governanceplugin.ParseVirtualKeyFromFastHTTPRequest(ctx) | ||
| if vkValue == nil { | ||
| return true | ||
| } | ||
|
|
||
| trimmedVKValue := strings.TrimSpace(*vkValue) | ||
| if trimmedVKValue == "" { | ||
| return true | ||
| } | ||
|
|
||
| if h.config == nil || h.config.ConfigStore == nil { | ||
| SendError(ctx, fasthttp.StatusServiceUnavailable, "database store unavailable") | ||
| return false | ||
| } | ||
|
|
||
| vk, err := h.config.ConfigStore.GetVirtualKeyByValue(ctx, trimmedVKValue) | ||
| if err != nil { | ||
| if errors.Is(err, configstore.ErrNotFound) { | ||
| return true | ||
| } | ||
| SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to resolve virtual key: %v", err)) | ||
| return false | ||
| } | ||
| if vk == nil { | ||
| return true | ||
| } | ||
|
|
||
| availableProviders := make([]schemas.ModelProvider, 0, len(vk.ProviderConfigs)) | ||
| for _, providerConfig := range vk.ProviderConfigs { | ||
| provider := strings.TrimSpace(providerConfig.Provider) | ||
| if provider == "" { | ||
| continue | ||
| } | ||
| availableProviders = append(availableProviders, schemas.ModelProvider(provider)) | ||
| } | ||
|
|
||
| bifrostCtx.SetValue(schemas.BifrostContextKeyAvailableProviders, availableProviders) | ||
| return true | ||
| } |
There was a problem hiding this comment.
applyListModelsVirtualKeyProviderFilter
The existing TestListModels_VKFilter* tests in providers_test.go exercise ProviderHandler, not CompletionHandler. The new applyListModelsVirtualKeyProviderFilter on CompletionHandler — including the not-found fallthrough, the config-store-unavailable 503 path, and the inactive-VK edge case — has no direct test coverage. Adding a list_models_vk_test.go (or extending governance_test.go) would keep coverage consistent with the rest of the handlers package.
| } | ||
|
|
||
| availableProviders := make([]schemas.ModelProvider, 0, len(vk.ProviderConfigs)) | ||
| for _, providerConfig := range vk.ProviderConfigs { | ||
| provider := strings.TrimSpace(providerConfig.Provider) | ||
| if provider == "" { | ||
| continue | ||
| } | ||
| availableProviders = append(availableProviders, schemas.ModelProvider(provider)) | ||
| } | ||
|
|
||
| bifrostCtx.SetValue(schemas.BifrostContextKeyAvailableProviders, availableProviders) | ||
| return true |
There was a problem hiding this comment.
After retrieving the VK the function never checks vk.IsActive. Every other place that consumes a VK checks this flag before honoring it — resolver.go:261 returns DecisionVirtualKeyBlocked, and main.go:406 drops VK governance entirely when !IsActive. Without the same guard here, an inactive VK still scopes the fan-out to its provider configs, producing incorrect results (empty or restricted model list) instead of the expected behavior (treat the key as absent or return an explicit error).
| } | |
| availableProviders := make([]schemas.ModelProvider, 0, len(vk.ProviderConfigs)) | |
| for _, providerConfig := range vk.ProviderConfigs { | |
| provider := strings.TrimSpace(providerConfig.Provider) | |
| if provider == "" { | |
| continue | |
| } | |
| availableProviders = append(availableProviders, schemas.ModelProvider(provider)) | |
| } | |
| bifrostCtx.SetValue(schemas.BifrostContextKeyAvailableProviders, availableProviders) | |
| return true | |
| if vk == nil || !vk.IsActive { | |
| return true | |
| } | |
| availableProviders := make([]schemas.ModelProvider, 0, len(vk.ProviderConfigs)) | |
| for _, providerConfig := range vk.ProviderConfigs { | |
| provider := strings.TrimSpace(providerConfig.Provider) | |
| if provider == "" { | |
| continue | |
| } | |
| availableProviders = append(availableProviders, schemas.ModelProvider(provider)) | |
| } | |
| bifrostCtx.SetValue(schemas.BifrostContextKeyAvailableProviders, availableProviders) | |
| return true |
1 participant
Summary
/v1/modelsfan-out to the providers allowed by the request virtual key.ListAllModelshonorBifrostContextKeyAvailableProvidersso disallowed providers are skipped before provider-level list calls run.Why
When
/v1/modelsis called with a virtual key scoped to one provider, Bifrost previously attempted to list models from every configured provider. Governance then rejected the disallowed providers, creating noisy expected errors in logs and inflating error rates.Validation
cd core && go test . -run TestFilterProvidersByContext -count=1cd transports && go test ./bifrost-http/handlers -run TestListModels_VKFilter -count=1Duplicate check
#2887and the issue title before opening. No open PR was found.