refactor: replace aggregated model discovery shims with per-key live cache fanout via OnKeyAdded/Updated/Deleted

[Pratham-Mishra04]

Summary

Replaces the single aggregated per-provider live model cache with a per-(provider, keyID) cache. Previously, all keys for a provider were merged into a single live entry keyed by "". Now each key gets its own entry, and key lifecycle events (OnKeyAdded, OnKeyUpdated, OnKeyDeleted) trigger targeted fetches or invalidations for only the affected key rather than re-fetching all keys for the provider.

Changes

• Removed shims.go and its deprecated UpsertModelDataForProvider, UpsertUnfilteredModelDataForProvider, and DeleteModelDataForProvider methods. The extractModelIDs helper was moved inline.
• Added UpsertLiveFromResponse to ModelCatalog as the canonical way to push a BifrostListModelsResponse into the live cache for a specific key.
• Added OnKeyAdded, OnKeyUpdated, and OnKeyDeleted to the ServerCallbacks interface and ModelsManager interface, with implementations in BifrostHTTPServer. Key create/update/delete handlers in provider_keys.go now call these instead of the old attemptModelDiscovery.
• Replaced populateModelPoolWithListModels with RefreshLiveModelsForProvider (fans out per-key in parallel) and FetchAndStoreLiveForKey (issues filtered + unfiltered list-models for a single key concurrently).
• Bootstrap now calls ReplaceKeyConfig to seed the full keyconfig snapshot before fanning out per-provider live fetches.
• ReloadProvider now calls SetKeyConfigForProvider + InvalidateLiveProvider + RefreshLiveModelsForProvider instead of re-running the full aggregated list-models flow.
• RemoveProvider now calls InvalidateLiveProvider + RemoveKeyConfigForProvider.
• ForceReloadPricing and ReloadPricingFromDBAndPopulateModelPool no longer trigger a list-models refresh — pricing reload is now pricing-only.
• Keyless providers continue to use the "" sentinel for their live cache entry. An isKeylessProvider helper centralizes that check.
• Test fixtures updated to use modelcatalog.NewTestCatalog(nil) instead of &modelcatalog.ModelCatalog{}, and mockModelsManager updated to implement the new OnKey* methods.

Type of change

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

Affected areas

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

How to test

go version
go test ./...

Verify that adding, updating, and deleting a provider key triggers only the expected number of list-models calls (2 per key rather than 2×N). Confirm that ReloadProvider correctly invalidates stale entries and re-fetches only for the current key set.

Screenshots/Recordings

N/A

Breaking changes

• Yes
• No

ServerCallbacks and ModelsManager interfaces have three new required methods: OnKeyAdded, OnKeyUpdated, and OnKeyDeleted. Any external implementations of these interfaces must add these methods. The deprecated shim methods (UpsertModelDataForProvider, UpsertUnfilteredModelDataForProvider, DeleteModelDataForProvider) have been removed.

Related issues

N/A

Security considerations

No new auth, secrets, or PII handling introduced. Key validation (BifrostContextKeyValidateKeys) is preserved in FetchAndStoreLiveForKey.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

Summary by CodeRabbit

• New Features

  • Provider key lifecycle events (add/update/delete) now trigger automatic model-catalog synchronization.
  • Live model upsert from provider responses to populate per-key live caches.

• Refactor

  • Model discovery/refresh redesigned to run per-provider per-key for faster, incremental updates.
  • Deprecated compatibility methods that aggregated provider data removed.

• Behavior Changes

  • Catalog-refresh hook failures are logged and do not change API responses.
  • Pricing reloads no longer trigger model-pool refresh.
  • Startup seeds catalogs from provider key snapshots and refreshes live models per key.

• Tests

  • Adjusted tests to use per-key catalog seeding; removed a couple obsolete wildcard tests.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Refactors model-catalog refresh to be key-driven: adds UpsertLiveFromResponse, key lifecycle callbacks, routes handler key events to callbacks, introduces per-key parallel refresh primitives, updates bootstrap and provider reload flows, and decouples pricing reload from model listing.

Changes

Key-Driven Model Catalog Refresh

Layer / File(s)	Summary
Model Catalog API Evolution framework/modelcatalog/pool.go, plugins/governance/httptransportprehook_test.go, framework/modelcatalog/models.go	New UpsertLiveFromResponse writes provider/key-specific live entries from BifrostListModelsResponse; minor comment edit in GetProvidersForModel; tests updated to use the new helper.
Key Lifecycle Callback Contracts transports/bifrost-http/handlers/providers.go, transports/bifrost-http/server/server.go	Add OnKeyAdded, OnKeyUpdated, OnKeyDeleted callbacks to ModelsManager and ServerCallbacks.
Handler Key Event Routing transports/bifrost-http/handlers/provider_keys.go	Provider key create/update/delete handlers now call h.modelsManager lifecycle hooks instead of attempting immediate discovery; add/update are skipped for keyless providers and hook errors are logged as warnings.
Handler Test Updates transports/bifrost-http/handlers/providers_test.go	mockModelsManager implements new no-op callbacks; several listModelDetails tests switch to modelcatalog.NewTestCatalog(nil) for catalog setup.
Server Callback Implementation & Provider Cleanup transports/bifrost-http/server/server.go	BifrostHTTPServer implements key lifecycle callbacks, invalidates provider live cache and updates catalog key-config; RemoveProvider removes key-config and invalidates cache; helper for keyless sentinel added.
Live Models Refresh Primitives transports/bifrost-http/server/server.go	Add RefreshLiveModelsForProvider to fan out per-key refresh and FetchAndStoreLiveForKey to run filtered/unfiltered discovery in parallel and upsert results into the catalog with selective error handling.
Pricing Decoupling & Bootstrap transports/bifrost-http/server/server.go	ForceReloadPricing and ReloadPricingFromDBAndPopulateModelPool no longer trigger list-models/model-pool refresh; Bootstrap snapshots provider keys into the catalog then concurrently calls RefreshLiveModelsForProvider.
Governance test cleanup plugins/governance/resolver_test.go	Remove framework/modelcatalog import and two wildcard/catalog-opaque provider test cases.

Sequence Diagram(s)

sequenceDiagram
  participant ProviderHTTPHandler
  participant ModelsManager
  participant BifrostHTTPServer
  participant ModelCatalog
  ProviderHTTPHandler->>ModelsManager: OnKeyAdded(ctx, provider, key)
  ModelsManager->>BifrostHTTPServer: Trigger RefreshLiveModelsForProvider(provider, keys)
  BifrostHTTPServer->>ModelCatalog: FetchAndStoreLiveForKey(provider, keyID) (filtered + unfiltered)
  ModelCatalog->>ModelCatalog: UpsertLiveFromResponse(provider, keyID, unfiltered, resp)

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested reviewers

• akshaydeo
• danpiths
• roroghost17

"I’m a rabbit with a tiny hat,
Keys hop in, the catalog chats,
Per-key fetches split and play,
Live cache blossoms in the day,
Hooray — the models find their way!"

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main refactoring: replacing aggregated model discovery shims with per-key live cache fanout via key lifecycle callbacks.
Description check	✅ Passed	The description is comprehensive, covering summary, detailed changes, type of change, affected areas, testing instructions, breaking changes with migration guidance, and security considerations.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 06-04-feat_wire_up_modelcatalog_composer

---

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

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Pratham-Mishra04]

• feat: add load balancer settings page with enterprise fallback view #4167
• feat: add core routing engine audit trail for retry and fallback transitions #4067
• refactor: replace aggregated model discovery shims with per-key live cache fanout via OnKeyAdded/Updated/Deleted #4038 👈 (View in Graphite)
• feat: add KeyID filter to ListModelsRequest to scope calls to a single key #4036 : 1 other dependent PR (#4037 )
• feat: adds datasheet store to modelcatalog #4035
• feat: add live model cache store and port keyconfig regression tests for alias/model isolation #4034
• feat: add keyconfig.Store for per-key config aggregation with allow/block/alias views #4033
• docs: document rich alias object form, routing_info response block, and pricing lookup precedence #4022
• refactor: replace resolvePricing positional model args with RoutingInfo and add backward-compat fallback for legacy ExtraFields #4021
• feat: add RoutingInfo to response/error extra fields with fallback and key alias signals #4020
• feat: replace flat aliases table with rich deployments table supporting per-deployment model family, canonical name, and provider-specific overrides #4019
• feat: add per-alias Replicate use deployments override and provider-scoped alias sub-config validation #4018
• feat: add per-alias Vertex project/region overrides and context-aware model family resolution #4017
• feat: propagate context into Bedrock region/ARN/model-family resolution for per-alias overrides #4016
• feat: add per-alias Azure endpoint, API version, and Anthropic version overrides with context-aware family resolution #4015
• refactor: promote KeyAliases values from string to AliasConfig with backward-compatible JSON marshaling #4014
• feat: enforce routing allowlist in core and propagate VK provider constraints via context key #3936 : 1 other dependent PR (#3711 )
• docs: document PreRequestHook routing phase, per-request vs per-attempt semantics, and updated plugin sequencing #3935
• refactor: move model-catalog provider resolution into a modelcatalogresolver PreRequestHook plugin #3934
• refactor: migrate governance routing from HTTPTransportPreHook to PreRequestHook #3933
• feat: add PreRequestHook to LLMPlugin interface as a once-per-request routing phase before PreLLMHook #3932
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

[greptile-apps]

Confidence Score: 5/5

Safe to merge; the refactor is well-structured with no behavioral regressions in the changed code paths.

All changed production paths are logically correct: keyconfig seeding, per-key live cache fanout, lifecycle callbacks, and the keyless/non-keyless guards are consistent. The only concern is two governance resolver tests that were deleted instead of updated, leaving wildcard + opaque-provider blacklist ordering without direct coverage — but the underlying logic is unchanged.

plugins/governance/resolver_test.go — two behavioral tests were removed rather than updated to the new catalog API.

Important Files Changed

Filename	Overview
framework/modelcatalog/pool.go	Adds UpsertLiveFromResponse — nil-safe wrapper that extracts model IDs from a BifrostListModelsResponse and forwards to the live store; clean and well-tested.
framework/modelcatalog/pool_test.go	New test file with thorough coverage for UpsertLiveFromResponse and extractModelIDs — nil guard, prefix stripping, gateway nesting, foreign-prefix filtering, dedup, and per-key invalidation.
framework/modelcatalog/shims.go	Deleted file — removes the deprecated UpsertModelDataForProvider, UpsertUnfilteredModelDataForProvider, and DeleteModelDataForProvider shims as planned.
transports/bifrost-http/server/server.go	Core refactor: Bootstrap seeds keyconfig via ReplaceKeyConfig then fans out per-provider; ReloadProvider delegates to RefreshLiveModelsForProvider; OnKeyAdded/Updated/Deleted target only the affected key; FetchAndStoreLiveForKey properly validates keys and propagates key statuses; pricing reload is now pricing-only. No concurrency issues; guard for non-keyless empty-key providers is correct.
transports/bifrost-http/handlers/provider_keys.go	Replaces attemptModelDiscovery calls with OnKeyAdded/Updated/Deleted; keyless guard on create/update is consistent; delete handler correctly lacks the guard since OnKeyDeleted makes no network calls and the keyless path is already blocked at deleteProviderKey entry.
plugins/governance/resolver_test.go	Two governance resolver tests deleted instead of being updated to use the new API; the blacklist-before-wildcard ordering on catalog-opaque providers is no longer directly tested.
plugins/governance/httptransportprehook_test.go	Migrated UpsertModelDataForProvider call to UpsertLiveFromResponse with explicit keyID sentinel "" — correct for keyless/aggregate usage in tests.
transports/bifrost-http/handlers/providers.go	Extends ModelsManager interface with OnKeyAdded, OnKeyUpdated, OnKeyDeleted — straightforward addition.
transports/bifrost-http/handlers/providers_test.go	Mock updated to implement the new OnKey* interface methods; ModelCatalog field initialised via NewTestCatalog(nil) instead of zero-value struct literal.

_{Reviews (13): Last reviewed commit: "feat: wire up modelcatalog composer" | Re-trigger Greptile}

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@transports/bifrost-http/server/server.go`:
- Around line 908-922: RefreshLiveModelsForProvider launches an unbounded
goroutine per key causing spikes; limit concurrency by using a bounded worker
pool/semaphore (e.g., a buffered channel or worker goroutines) when iterating
keys, acquire a token before calling FetchAndStoreLiveForKey and release it
after the call, still using the existing sync.WaitGroup to wait for completion;
update RefreshLiveModelsForProvider to use this semaphore/worker pattern
(referencing RefreshLiveModelsForProvider and FetchAndStoreLiveForKey) so at
most N concurrent key refreshes run (choose a sensible default like 10 or make
it configurable).
- Around line 1573-1578: The loop unconditionally calls
RefreshLiveModelsForProvider for providers with zero keys, causing an unintended
keyless fallback; update the provider loop to skip calling
s.RefreshLiveModelsForProvider when providerConfig.Keys is empty and the
provider is not declared keyless (use the provider config flag corresponding to
custom_provider_config.is_key_less, e.g., providerConfig.IsKeyLess or similar).
In practice, before wg.Add/starting the goroutine, check if
len(providerConfig.Keys) == 0 && !providerConfig.IsKeyLess then continue;
otherwise proceed to spawn the goroutine calling
s.RefreshLiveModelsForProvider(ctx, provider, providerConfig.Keys).
- Around line 930-970: FetchAndStoreLiveForKey always calls ListModelsRequest
even when a provider has custom_provider_config.allowed_requests.list_models ==
false; before starting the goroutines in FetchAndStoreLiveForKey, check the
provider's custom config (custom_provider_config.allowed_requests.list_models)
from the server config and short-circuit/return (or skip spawning both
ListModelsRequest goroutines) when list_models is disabled for that provider so
no ListModelsRequest() calls are issued; update code paths around
FetchAndStoreLiveForKey, the two goroutines that call
s.Client.ListModelsRequest, and any helper used to read provider configs (e.g.,
the server config accessors) to perform this check.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: b2534ba7-40a5-41ae-a47c-5e1ffc97402d

📥 Commits

Reviewing files that changed from the base of the PR and between 8056692 and 489baca.

📒 Files selected for processing (8)

• framework/modelcatalog/models.go
• framework/modelcatalog/pool.go
• framework/modelcatalog/shims.go
• plugins/governance/httptransportprehook_test.go
• transports/bifrost-http/handlers/provider_keys.go
• transports/bifrost-http/handlers/providers.go
• transports/bifrost-http/handlers/providers_test.go
• transports/bifrost-http/server/server.go

💤 Files with no reviewable changes (1)

• framework/modelcatalog/shims.go

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@framework/modelcatalog/pool.go`:
- Around line 19-21: The UpsertLiveFromResponse method can clear the cache when
resp is nil because extractModelIDs returns nil; add an early nil guard in
ModelCatalog.UpsertLiveFromResponse to return immediately if resp is nil (and
optionally if resp.Models == nil) before calling mc.live.Upsert so you don't
pass a nil/empty slice into mc.live.Upsert; reference the function
ModelCatalog.UpsertLiveFromResponse, the helper extractModelIDs, and the call
mc.live.Upsert to locate where to add the check.

In `@transports/bifrost-http/server/server.go`:
- Around line 950-999: The single shared bfCtx is mutated by SetValue and then
reused concurrently by both goroutines—create a fresh BifrostContext for each
goroutine instead: inside each goroutine call schemas.NewBifrostContext(ctx,
time.Now().Add(15*time.Second)), set the same flags with
SetValue(schemas.BifrostContextKeySkipPluginPipeline, true) and
SetValue(schemas.BifrostContextKeyValidateKeys, true), defer Cancel() on that
per-goroutine context, and pass that new context to s.Client.ListModelsRequest;
keep the rest of the logic (keyIDPtr, UpsertLiveFromResponse, updateKeyStatus)
unchanged and ensure both goroutines use independent contexts to avoid
shared-state races involving NewBifrostContext and SetValue.
- Around line 630-644: The code is silently ignoring errors from
GetProviderKeysRaw which causes SetKeyConfigForProvider(provider, nil) and
subsequent InvalidateLiveProvider/skip discovery; update ReloadProvider and the
three OnKey* callbacks to capture the error returned by
s.Config.GetProviderKeysRaw(provider) (e.g. inMemoryKeys, err := ...), and if
err != nil return or propagate that error (or handle it explicitly) instead of
treating keys as empty—do not call SetKeyConfigForProvider(provider, nil) or
call InvalidateLiveProvider when the lookup failed; only call
SetKeyConfigForProvider/InvalidateLiveProvider/RefreshLiveModelsForProvider when
GetProviderKeysRaw succeeds so the catalog stays in sync.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 6d79d08d-1908-4082-9c68-b2d8878ca592

📥 Commits

Reviewing files that changed from the base of the PR and between 3e436ba and 40bebde.

📒 Files selected for processing (8)

• framework/modelcatalog/models.go
• framework/modelcatalog/pool.go
• plugins/governance/httptransportprehook_test.go
• plugins/governance/resolver_test.go
• transports/bifrost-http/handlers/provider_keys.go
• transports/bifrost-http/handlers/providers.go
• transports/bifrost-http/handlers/providers_test.go
• transports/bifrost-http/server/server.go

💤 Files with no reviewable changes (1)

• plugins/governance/resolver_test.go

[akshaydeo]

Merge activity

• Jun 9, 5:17 AM UTC: A user started a stack merge that includes this pull request via Graphite.
• Jun 9, 5:35 AM UTC: @akshaydeo merged this pull request with Graphite.