From 78b230d5b9508b6e6701a5a611120bfb185ea1cd Mon Sep 17 00:00:00 2001 From: Anuj Parihar Date: Fri, 5 Jun 2026 11:29:03 +0530 Subject: [PATCH] feat: make use of http pre-hook instead of handler for model filter --- transports/bifrost-http/handlers/inference.go | 3 - .../bifrost-http/handlers/list_models_vk.go | 57 ------ .../handlers/list_models_vk_test.go | 167 ------------------ 3 files changed, 227 deletions(-) delete mode 100644 transports/bifrost-http/handlers/list_models_vk.go delete mode 100644 transports/bifrost-http/handlers/list_models_vk_test.go diff --git a/transports/bifrost-http/handlers/inference.go b/transports/bifrost-http/handlers/inference.go index 6e4afb6339..b3207738ef 100644 --- a/transports/bifrost-http/handlers/inference.go +++ b/transports/bifrost-http/handlers/inference.go @@ -799,9 +799,6 @@ func (h *CompletionHandler) listModels(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusBadRequest, "Failed to convert context") return } - if provider == "" && !h.applyListModelsVirtualKeyProviderFilter(ctx, bifrostCtx) { - return - } var resp *schemas.BifrostListModelsResponse var bifrostErr *schemas.BifrostError diff --git a/transports/bifrost-http/handlers/list_models_vk.go b/transports/bifrost-http/handlers/list_models_vk.go deleted file mode 100644 index 14b0a589d6..0000000000 --- a/transports/bifrost-http/handlers/list_models_vk.go +++ /dev/null @@ -1,57 +0,0 @@ -package handlers - -import ( - "errors" - "fmt" - "strings" - - "github.com/maximhq/bifrost/core/schemas" - "github.com/maximhq/bifrost/framework/configstore" - governanceplugin "github.com/maximhq/bifrost/plugins/governance" - "github.com/valyala/fasthttp" -) - -// applyListModelsVirtualKeyProviderFilter narrows provider fan-out for GET /v1/models -// when the request is made with a virtual key. Without this, ListAllModels asks every -// configured provider to list models and governance rejects providers outside the VK, -// creating noisy, expected errors in request logs. -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 || !vk.IsActiveValue() { - 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 -} diff --git a/transports/bifrost-http/handlers/list_models_vk_test.go b/transports/bifrost-http/handlers/list_models_vk_test.go deleted file mode 100644 index 1d493c6d20..0000000000 --- a/transports/bifrost-http/handlers/list_models_vk_test.go +++ /dev/null @@ -1,167 +0,0 @@ -package handlers - -import ( - "context" - "errors" - "strings" - "testing" - "time" - - "github.com/maximhq/bifrost/core/schemas" - "github.com/maximhq/bifrost/framework/configstore" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" - "github.com/maximhq/bifrost/transports/bifrost-http/lib" - "github.com/valyala/fasthttp" -) - -type mockListModelsVKConfigStore struct { - configstore.ConfigStore - vk *configstoreTables.TableVirtualKey - err error -} - -func (m *mockListModelsVKConfigStore) GetVirtualKeyByValue(_ context.Context, _ string) (*configstoreTables.TableVirtualKey, error) { - return m.vk, m.err -} - -func TestApplyListModelsVirtualKeyProviderFilterSetsActiveVKProviders(t *testing.T) { - h := &CompletionHandler{ - config: &lib.Config{ - ConfigStore: &mockListModelsVKConfigStore{vk: &configstoreTables.TableVirtualKey{ - Value: "sk-bf-active", - IsActive: schemas.Ptr(true), - ProviderConfigs: []configstoreTables.TableVirtualKeyProviderConfig{ - {Provider: "openai"}, - {Provider: " anthropic "}, - {Provider: ""}, - }, - }}, - }, - } - - ctx := &fasthttp.RequestCtx{} - ctx.Request.Header.Set("Authorization", "Bearer sk-bf-active") - bifrostCtx := schemas.NewBifrostContext(context.Background(), time.Time{}) - - if ok := h.applyListModelsVirtualKeyProviderFilter(ctx, bifrostCtx); !ok { - t.Fatalf("expected active VK to apply provider filter") - } - got, ok := bifrostCtx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - if !ok { - t.Fatalf("expected available providers to be set") - } - want := []schemas.ModelProvider{schemas.OpenAI, schemas.Anthropic} - if len(got) != len(want) { - t.Fatalf("expected providers %#v, got %#v", want, got) - } - for i := range want { - if got[i] != want[i] { - t.Fatalf("expected providers %#v, got %#v", want, got) - } - } -} - -func TestApplyListModelsVirtualKeyProviderFilterReturnsErrorOnLookupFailure(t *testing.T) { - h := &CompletionHandler{ - config: &lib.Config{ - ConfigStore: &mockListModelsVKConfigStore{err: errors.New("database unavailable")}, - }, - } - - ctx := &fasthttp.RequestCtx{} - ctx.Request.Header.Set("Authorization", "Bearer sk-bf-active") - bifrostCtx := schemas.NewBifrostContext(context.Background(), time.Time{}) - - if ok := h.applyListModelsVirtualKeyProviderFilter(ctx, bifrostCtx); ok { - t.Fatalf("expected lookup error to fail request") - } - if got := ctx.Response.StatusCode(); got != fasthttp.StatusInternalServerError { - t.Fatalf("expected status %d, got %d", fasthttp.StatusInternalServerError, got) - } - if body := string(ctx.Response.Body()); !strings.Contains(body, "Failed to resolve virtual key") { - t.Fatalf("expected virtual key lookup error response, got %q", body) - } -} - -func TestApplyListModelsVirtualKeyProviderFilterReturnsUnavailableWithoutConfigStore(t *testing.T) { - h := &CompletionHandler{config: &lib.Config{}} - - ctx := &fasthttp.RequestCtx{} - ctx.Request.Header.Set("Authorization", "Bearer sk-bf-active") - bifrostCtx := schemas.NewBifrostContext(context.Background(), time.Time{}) - - if ok := h.applyListModelsVirtualKeyProviderFilter(ctx, bifrostCtx); ok { - t.Fatalf("expected missing config store to fail request") - } - if got := ctx.Response.StatusCode(); got != fasthttp.StatusServiceUnavailable { - t.Fatalf("expected status %d, got %d", fasthttp.StatusServiceUnavailable, got) - } - if body := string(ctx.Response.Body()); !strings.Contains(body, "database store unavailable") { - t.Fatalf("expected unavailable response, got %q", body) - } -} - -func TestApplyListModelsVirtualKeyProviderFilterSkipsWhenVKNotFound(t *testing.T) { - t.Run("nil return from store", func(t *testing.T) { - h := &CompletionHandler{ - config: &lib.Config{ - ConfigStore: &mockListModelsVKConfigStore{}, - }, - } - - ctx := &fasthttp.RequestCtx{} - ctx.Request.Header.Set("Authorization", "Bearer sk-bf-missing") - bifrostCtx := schemas.NewBifrostContext(context.Background(), time.Time{}) - - if ok := h.applyListModelsVirtualKeyProviderFilter(ctx, bifrostCtx); !ok { - t.Fatalf("expected missing VK to be ignored without failing request") - } - if got := bifrostCtx.Value(schemas.BifrostContextKeyAvailableProviders); got != nil { - t.Fatalf("expected missing VK not to set available providers, got %#v", got) - } - }) - - t.Run("ErrNotFound from store", func(t *testing.T) { - h := &CompletionHandler{ - config: &lib.Config{ - ConfigStore: &mockListModelsVKConfigStore{err: configstore.ErrNotFound}, - }, - } - - ctx := &fasthttp.RequestCtx{} - ctx.Request.Header.Set("Authorization", "Bearer sk-bf-missing") - bifrostCtx := schemas.NewBifrostContext(context.Background(), time.Time{}) - - if ok := h.applyListModelsVirtualKeyProviderFilter(ctx, bifrostCtx); !ok { - t.Fatalf("expected ErrNotFound to be ignored without failing request") - } - if got := bifrostCtx.Value(schemas.BifrostContextKeyAvailableProviders); got != nil { - t.Fatalf("expected ErrNotFound not to set available providers, got %#v", got) - } - }) -} - -func TestApplyListModelsVirtualKeyProviderFilterSkipsInactiveVK(t *testing.T) { - h := &CompletionHandler{ - config: &lib.Config{ - ConfigStore: &mockListModelsVKConfigStore{vk: &configstoreTables.TableVirtualKey{ - Value: "sk-bf-inactive", - IsActive: schemas.Ptr(false), - ProviderConfigs: []configstoreTables.TableVirtualKeyProviderConfig{ - {Provider: "openai"}, - }, - }}, - }, - } - - ctx := &fasthttp.RequestCtx{} - ctx.Request.Header.Set("Authorization", "Bearer sk-bf-inactive") - bifrostCtx := schemas.NewBifrostContext(context.Background(), time.Time{}) - - if ok := h.applyListModelsVirtualKeyProviderFilter(ctx, bifrostCtx); !ok { - t.Fatalf("expected inactive VK to be ignored without failing request") - } - if got := bifrostCtx.Value(schemas.BifrostContextKeyAvailableProviders); got != nil { - t.Fatalf("expected inactive VK not to set available providers, got %#v", got) - } -}