refactor: replace aggregated model discovery shims with per-key live cache fanout via OnKeyAdded/Updated/Deleted - #4194
Conversation
|
Warning Review limit reached
More reviews will be available in 33 minutes and 45 seconds. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (9)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
|
Confidence Score: 3/5The per-key fanout logic is sound, but the clear-before-fetch ordering in ReloadProvider can silently leave the live cache empty after a transient list-models failure, and FetchAndStoreLiveForKey goroutines access s.Config without a nil guard. Two concrete defects in the hot path: ReloadProvider calls InvalidateLiveProvider unconditionally before spawning fetches so if any per-key fetch fails the live entries are gone and the function still returns nil hiding the degradation; FetchAndStoreLiveForKey goroutines dereference s.Config.ModelCatalog without a nil guard creating a latent panic path. transports/bifrost-http/server/server.go — specifically the ReloadProvider invalidation ordering and FetchAndStoreLiveForKey nil safety. Important Files Changed
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
74900ce to
c46a122
Compare
267d94f to
5952bf3
Compare
c46a122 to
2a1f701
Compare
Merge activity
|
2a1f701 to
6d9fa38
Compare
| // Refresh keyconfig from the current key list, then drop any stale live | ||
| // entries (for keys removed in this update) before refetching per-key. | ||
| s.Config.ModelCatalog.SetKeyConfigForProvider(provider, inMemoryKeys) | ||
| s.Config.ModelCatalog.InvalidateLiveProvider(provider) | ||
| if hasNoKeys { | ||
| logger.Warn("model discovery skipped for provider %s: no keys configured", provider) | ||
| } else { | ||
| s.Config.ModelCatalog.UpsertUnfilteredModelDataForProvider(provider, unfilteredModels) | ||
| s.RefreshLiveModelsForProvider(ctx, provider, inMemoryKeys) | ||
| } | ||
| return updatedProvider, nil |
There was a problem hiding this comment.
InvalidateLiveProvider clears all existing entries for the provider before RefreshLiveModelsForProvider runs. If any per-key FetchAndStoreLiveForKey call fails (errors are logged and swallowed, not propagated), the affected key's live entry is never written back. ReloadProvider then returns (updatedProvider, nil), so the caller has no signal that the cache is now empty. Until the next successful reload, GetModelsForProvider returns nothing for those keys, and routing silently degrades to the static catalog.
The old code wrote atomically (invalidate + upsert happened inside UpsertModelDataForProvider in a single shim call), so a failed list-models call still left a populated entry. A safer ordering would be to collect the per-key results first and only replace the live entries that were successfully refreshed — leave entries for keys whose fetch failed in place rather than pre-emptively clearing them.
…cache fanout via `OnKeyAdded/Updated/Deleted` (#4194) ## Summary Replaces the coarse provider-level model discovery (one aggregated live entry per provider) with per-key live cache entries, so adding, updating, or deleting a single key triggers at most 2 list-models calls for that key rather than 2×N calls across all keys. Removes the deprecated compatibility shims (`UpsertModelDataForProvider`, `UpsertUnfilteredModelDataForProvider`, `DeleteModelDataForProvider`) and replaces them with `UpsertLiveFromResponse`, `OnKeyAdded`, `OnKeyUpdated`, and `OnKeyDeleted`. ## Changes - **`modelcatalog/pool.go`**: Added `UpsertLiveFromResponse`, which extracts and deduplicates model IDs from a `BifrostListModelsResponse` before writing to the live cache. A nil response is a no-op, preventing accidental cache eviction. - **`modelcatalog/shims.go`**: Deleted. The three deprecated shim methods are gone; call sites now use the per-key API directly. - **`modelcatalog/pool_test.go`**: New test file covering `UpsertLiveFromResponse` (nil no-op, happy path), `extractModelIDs` (prefix stripping, gateway nested prefixes, foreign prefix filtering, nil input, deduplication), `InvalidateLive`, and `InvalidateLiveProvider`. - **`server/server.go`**: Replaced `populateModelPoolWithListModels` (one aggregated entry per provider) with `RefreshLiveModelsForProvider` (fans out per key in parallel) and `FetchAndStoreLiveForKey` (issues filtered + unfiltered list-models for a single key). `ReloadProvider` now reads keys from the in-memory store, calls `SetKeyConfigForProvider` + `InvalidateLiveProvider`, then delegates to `RefreshLiveModelsForProvider`. `ForceReloadPricing` and `ReloadPricingFromDBAndPopulateModelPool` no longer trigger a full model pool refresh — pricing reload is now pricing-only. `RemoveProvider` calls `InvalidateLiveProvider` + `RemoveKeyConfigForProvider` instead of the deleted shim. Added `OnKeyAdded`, `OnKeyUpdated`, `OnKeyDeleted` to `ServerCallbacks` and implemented them on `BifrostHTTPServer`. - **`handlers/provider_keys.go`**: Key create/update/delete handlers now call `modelsManager.OnKeyAdded/OnKeyUpdated/OnKeyDeleted` instead of `attemptModelDiscovery`. Keyless providers skip the add/update path. - **`handlers/providers.go`**: Extended `ModelsManager` interface with `OnKeyAdded`, `OnKeyUpdated`, `OnKeyDeleted`. - **`handlers/providers_test.go`** and **`governance/httptransportprehook_test.go`**: Updated to use `UpsertLiveFromResponse` and `NewTestCatalog(nil)` instead of the removed shims and bare struct literals. - **`governance/resolver_test.go`**: Removed two tests that depended on the shim API; equivalent coverage exists in the catalog-level tests. - **`modelcatalog/models.go`**: Removed stale comment referencing the pre-refactor file. ## Type of change - [ ] Bug fix - [ ] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./framework/modelcatalog/... go test ./transports/bifrost-http/... go test ./plugins/governance/... go test ./... ``` After adding a provider key via `POST /api/providers/{provider}/keys`, verify that only the new key's models are fetched (2 list-models calls) rather than a full provider refresh. After deleting a key, confirm the deleted key's live entries are evicted while other keys' entries remain intact. ## Screenshots/Recordings N/A ## Breaking changes - [x] Yes - [ ] No `UpsertModelDataForProvider`, `UpsertUnfilteredModelDataForProvider`, and `DeleteModelDataForProvider` are removed from `ModelCatalog`. Any external code calling these methods must migrate to `UpsertLiveFromResponse` / `UpsertLive` / `InvalidateLiveProvider`. `ServerCallbacks` now requires `OnKeyAdded`, `OnKeyUpdated`, and `OnKeyDeleted` — implementors must add these three methods. ## Related issues N/A ## Security considerations No new auth surfaces. Key validation (`BifrostContextKeyValidateKeys`) is preserved in `FetchAndStoreLiveForKey`, maintaining the same key-validation behavior at boot, after key add, and after provider reload. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
2 participants
Summary
Replaces the coarse provider-level model discovery (one aggregated live entry per provider) with per-key live cache entries, so adding, updating, or deleting a single key triggers at most 2 list-models calls for that key rather than 2×N calls across all keys. Removes the deprecated compatibility shims (
UpsertModelDataForProvider,UpsertUnfilteredModelDataForProvider,DeleteModelDataForProvider) and replaces them withUpsertLiveFromResponse,OnKeyAdded,OnKeyUpdated, andOnKeyDeleted.Changes
modelcatalog/pool.go: AddedUpsertLiveFromResponse, which extracts and deduplicates model IDs from aBifrostListModelsResponsebefore writing to the live cache. A nil response is a no-op, preventing accidental cache eviction.modelcatalog/shims.go: Deleted. The three deprecated shim methods are gone; call sites now use the per-key API directly.modelcatalog/pool_test.go: New test file coveringUpsertLiveFromResponse(nil no-op, happy path),extractModelIDs(prefix stripping, gateway nested prefixes, foreign prefix filtering, nil input, deduplication),InvalidateLive, andInvalidateLiveProvider.server/server.go: ReplacedpopulateModelPoolWithListModels(one aggregated entry per provider) withRefreshLiveModelsForProvider(fans out per key in parallel) andFetchAndStoreLiveForKey(issues filtered + unfiltered list-models for a single key).ReloadProvidernow reads keys from the in-memory store, callsSetKeyConfigForProvider+InvalidateLiveProvider, then delegates toRefreshLiveModelsForProvider.ForceReloadPricingandReloadPricingFromDBAndPopulateModelPoolno longer trigger a full model pool refresh — pricing reload is now pricing-only.RemoveProvidercallsInvalidateLiveProvider+RemoveKeyConfigForProviderinstead of the deleted shim. AddedOnKeyAdded,OnKeyUpdated,OnKeyDeletedtoServerCallbacksand implemented them onBifrostHTTPServer.handlers/provider_keys.go: Key create/update/delete handlers now callmodelsManager.OnKeyAdded/OnKeyUpdated/OnKeyDeletedinstead ofattemptModelDiscovery. Keyless providers skip the add/update path.handlers/providers.go: ExtendedModelsManagerinterface withOnKeyAdded,OnKeyUpdated,OnKeyDeleted.handlers/providers_test.goandgovernance/httptransportprehook_test.go: Updated to useUpsertLiveFromResponseandNewTestCatalog(nil)instead of the removed shims and bare struct literals.governance/resolver_test.go: Removed two tests that depended on the shim API; equivalent coverage exists in the catalog-level tests.modelcatalog/models.go: Removed stale comment referencing the pre-refactor file.Type of change
Affected areas
How to test
After adding a provider key via
POST /api/providers/{provider}/keys, verify that only the new key's models are fetched (2 list-models calls) rather than a full provider refresh. After deleting a key, confirm the deleted key's live entries are evicted while other keys' entries remain intact.Screenshots/Recordings
N/A
Breaking changes
UpsertModelDataForProvider,UpsertUnfilteredModelDataForProvider, andDeleteModelDataForProviderare removed fromModelCatalog. Any external code calling these methods must migrate toUpsertLiveFromResponse/UpsertLive/InvalidateLiveProvider.ServerCallbacksnow requiresOnKeyAdded,OnKeyUpdated, andOnKeyDeleted— implementors must add these three methods.Related issues
N/A
Security considerations
No new auth surfaces. Key validation (
BifrostContextKeyValidateKeys) is preserved inFetchAndStoreLiveForKey, maintaining the same key-validation behavior at boot, after key add, and after provider reload.Checklist
docs/contributing/README.mdand followed the guidelines