fix: preserve non-pricing model pool entries across pricing reloads

[Pratham-Mishra04]

Summary

When populateModelPoolFromPricingData runs on a reload (e.g., gossip ReloadFromDB or ForceReloadPricing), it previously wiped modelPool and unfilteredModelPool entirely before rebuilding from pricing data. This caused models contributed by per-provider list-models output and allowed_models key entries to be silently dropped on every reload, creating drift between the initial state and the reloaded state.

Changes

• Before wiping the model pools, a snapshot of the current modelPool and unfilteredModelPool is taken.
• After the pricing-data rebuild completes, the snapshot is unioned back in, restoring any models that were added via UpsertModelDataForProvider / UpsertUnfilteredModelDataForProvider but are absent from the pricing sheet.
• Pricing entries from the rebuild take precedence on duplicates (they are already written before the union step).
• Models explicitly removed via DeleteModelDataForProvider are correctly excluded because that method strips the provider from the live map before this function runs.
• baseModelIndex is intentionally not preserved across reloads, as aliases outside the pricing sheet have no canonical base-model entry and fall through to algorithmic stripping.

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

Trigger a pricing reload (via gossip ReloadFromDB or ForceReloadPricing) on a running instance that has providers with list-models output or allowed_models entries. Verify that models contributed by those sources are still present in the model pool after the reload completes.

go test ./framework/modelcatalog/...

Breaking changes

• Yes
• No

Related issues

Security considerations

None.

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

• Bug Fixes
  • Fixed an issue where model data was being lost during pricing data reloads. Non-pricing model contributions are now preserved while pricing entries remain authoritative.

[coderabbitai]

📝 Walkthrough

Walkthrough

The populateModelPoolFromPricingData function in framework/modelcatalog/main.go now preserves non-pricing model contributions across pricing data rebuilds. A snapshot of existing model pools is captured before the rebuild clears them, then unioned back afterward, with pricing entries winning on conflicts.

Changes

Model Pool Preservation in Pricing Rebuild

Layer / File(s)	Summary
Snapshot and union model pools framework/modelcatalog/main.go	populateModelPoolFromPricingData snapshots modelPool and unfilteredModelPool before clearing them, then unions the snapshot back after rebuilding from pricing data, re-adding only models not already present per provider to ensure pricing entries take precedence while preserving non-pricing contributions.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Poem

🐰 A snapshot, a wipe, then a union so neat,
Non-pricing pools dance through the pricing rebuild feat,
Pricing wins conflicts, but old friends stay—
Model pools bloom, both old and new way! 🌱

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately captures the main change: preserving non-pricing model pool entries across pricing reloads, which directly addresses the core bug being fixed.
Description check	✅ Passed	The pull request description follows the template structure with all major sections completed: Summary, Changes, Type of change (Bug fix selected), Affected areas (Core Go selected), How to test provided, Breaking changes addressed, and Checklist items documented.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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 docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 06-02-fix_repopulate_in_memory_model_pool_entry_to_avoid_losing_custom_models

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.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

_{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]

• fix: preserve non-pricing model pool entries across pricing reloads #3999 👈 (View in Graphite)
• main

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

[Pratham-Mishra04]

@coderabbitai full-review

[greptile-apps]

Confidence Score: 3/5

The change fixes a real drift bug but introduces a mirrored staleness problem: models removed from the upstream pricing sheet will re-appear in the pool after every reload, remaining visible to callers until the next per-provider list-models refresh.

The snapshot-and-union approach correctly restores non-pricing models on reload, which is the stated goal. However, because the snapshot captures the entire model pool — including models whose only source was the previous pricing rebuild — any model deleted from the remote pricing sheet will be re-added by the union and will linger in the pool. This is a present defect on the changed path: after a reload that removes a model from pricing data, that model remains advertised as available. There are also no tests for the new reload behavior, leaving the fix invisible to the test suite.

framework/modelcatalog/main.go — the union loop and the absence of a dedicated reload test.

Important Files Changed

Filename	Overview
framework/modelcatalog/main.go	Adds a snapshot-and-union mechanism in populateModelPoolFromPricingData to preserve non-pricing model pool entries across reloads; the union inadvertently re-adds models that were legitimately removed from the remote pricing sheet, and no regression test is included.

_{Reviews (1): Last reviewed commit: "fix: repopulate in memory model pool ent..." | Re-trigger Greptile}

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Fetch PR details and diff
gh pr view 3999 --repo maximhq/bifrost

echo "=== PR DIFF ==="
gh pr diff 3999 --repo maximhq/bifrost

Length of output: 6093

---

🏁 Script executed:

#!/bin/bash
# Read the full modelcatalog main.go file
cat framework/modelcatalog/main.go

echo "=== TESTS ==="
fd -e go . framework/modelcatalog/ --exec cat {}

echo "=== CHECK IMPORTS ==="
head -50 framework/modelcatalog/main.go

Length of output: 274759

---

---

@Pratham-Mishra04 Here is a thorough review of the changes.

---

Review Summary

The fix is logically correct — the snapshot-then-union approach properly handles the three cases stated in the PR description (list-models additions, allowed_models aliases, and DeleteModelDataForProvider removals). A few items worth addressing before merge:

---

1. 🔴 No regression test for the fixed behaviour

The PR checklist acknowledges this, but a test is important here because the bug is purely about state surviving a function call. Without it, a future refactor of populateModelPoolFromPricingData could silently reintroduce the drift.

A minimal test skeleton, given the existing newTestCatalog / noOpLogger helpers in the package:

func TestPopulateModelPool_PreservesNonPricingEntriesOnReload(t *testing.T) {
    mc := newTestCatalog(nil, nil)
    mc.logger = noOpLogger{}

    // Simulate a non-pricing model added via UpsertModelDataForProvider
    mc.mu.Lock()
    mc.modelPool[schemas.OpenAI] = []string{"gpt-4o", "custom-model"}
    mc.unfilteredModelPool[schemas.OpenAI] = []string{"gpt-4o", "custom-model"}
    mc.mu.Unlock()

    // Add a pricing entry only for gpt-4o
    mc.pricingData[makeKey("gpt-4o", "openai", "chat")] = configstoreTables.TableModelPricing{
        Model: "gpt-4o", Provider: "openai", Mode: "chat",
    }

    // Simulate a reload
    mc.populateModelPoolFromPricingData()

    models := mc.GetModelsForProvider(schemas.OpenAI)
    assert.Contains(t, models, "gpt-4o")
    assert.Contains(t, models, "custom-model", "non-pricing model must survive reload")

    unfiltered := mc.GetUnfilteredModelsForProvider(schemas.OpenAI)
    assert.Contains(t, unfiltered, "custom-model", "non-pricing model must survive reload in unfiltered pool")
}

func TestPopulateModelPool_DeletedProviderNotRestoredOnReload(t *testing.T) {
    mc := newTestCatalog(nil, nil)
    mc.logger = noOpLogger{}

    mc.mu.Lock()
    mc.modelPool[schemas.OpenAI] = []string{"gpt-4o"}
    mc.unfilteredModelPool[schemas.OpenAI] = []string{"gpt-4o"}
    mc.mu.Unlock()

    // Simulate DeleteModelDataForProvider being called before the reload
    mc.DeleteModelDataForProvider(schemas.OpenAI)

    mc.populateModelPoolFromPricingData()

    // Provider was explicitly deleted — must not be restored from snapshot
    models := mc.GetModelsForProvider(schemas.OpenAI)
    assert.Empty(t, models, "explicitly deleted provider must not be restored on reload")
}

---

2. 🟡 O(n²) union due to slices.Contains

The union loop is:

for provider, models := range previousModelPool {
    for _, m := range models {
        if !slices.Contains(mc.modelPool[provider], m) {  // O(n) scan per model
            mc.modelPool[provider] = append(mc.modelPool[provider], m)
        }
    }
}

For providers like OpenRouter or Vertex that can carry thousands of models, this becomes O(m × n) per provider. Consider building a set from the already-rebuilt slice once per provider:

for provider, models := range previousModelPool {
    existing := make(map[string]struct{}, len(mc.modelPool[provider]))
    for _, m := range mc.modelPool[provider] {
        existing[m] = struct{}{}
    }
    for _, m := range models {
        if _, ok := existing[m]; !ok {
            mc.modelPool[provider] = append(mc.modelPool[provider], m)
        }
    }
}
// same pattern for previousUnfilteredModelPool

This reduces the inner check to O(1) and the whole pass to O(m + n).

---

3. 🟢 Shared slice between modelPool and unfilteredModelPool (minor)

After the pricing rebuild:

mc.modelPool[provider] = models
mc.unfilteredModelPool[provider] = models   // same backing array

Go's append will create separate backing arrays when capacity is exceeded, so the union step's independent appends are safe in practice. However, a reader who isn't familiar with Go slice semantics might find this surprising. A defensive copy (append([]string(nil), models...) or slices.Clone(models)) on one assignment would make the intent explicit and eliminate the dependency on append-on-overflow behaviour.

---

4. 🟢 baseModelIndex gap acknowledged — worth a comment in UpsertModelDataForProvider

Models restored from the snapshot won't have baseModelIndex entries, so getBaseModelNameUnsafe falls back to algorithmic stripping for them. This is correctly documented in the populateModelPoolFromPricingData comment. It may be worth a brief inline comment near the UpsertModelDataForProvider call sites noting that callers should not rely on catalog-based base-model resolution for models that are not in the pricing sheet, to help future contributors understand the constraint.

---

Overall: The bug fix is sound. Addressing the missing test (item 1) and the O(n²) union (item 2) would make this production-ready. Items 3 and 4 are minor polish.

[akshaydeo]

@coderabbitai full review

[coderabbitai]

✅ Actions performed

Full review triggered.

[coderabbitai]

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)

framework/modelcatalog/main.go (1)

409-494: ⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Add regression tests for snapshot-union model pool rebuild behavior
No framework/modelcatalog tests exercise populateModelPoolFromPricingData’s snapshot-then-union logic (existing tests focus on base model/name/capabilities and pricing computations; the only UpsertModelDataForProvider usage found is in a governance prehook test and doesn’t cover rebuild/persistence/deletion). Add tests covering:

• non-pricing models preserved across the rebuild (via UpsertModelDataForProvider / UpsertUnfilteredModelDataForProvider)
• pricing models present after rebuild
• models removed via DeleteModelDataForProvider remain excluded after rebuild

🤖 Prompt for 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.

In `@framework/modelcatalog/main.go` around lines 409 - 494, Add regression tests
that exercise ModelCatalog.populateModelPoolFromPricingData to verify
snapshot-then-union behavior: create a ModelCatalog test instance, seed
mc.pricingData with a few pricing entries, then call UpsertModelDataForProvider
and UpsertUnfilteredModelDataForProvider to add non-pricing models and also call
DeleteModelDataForProvider to remove one provider/model; invoke
populateModelPoolFromPricingData (or the public flow that triggers it) and
assert three things: (1) pricing models from mc.pricingData are present in
mc.modelPool/mc.unfilteredModelPool, (2) non-pricing models added via
UpsertModelDataForProvider / UpsertUnfilteredModelDataForProvider are preserved
after the rebuild, and (3) any models removed via DeleteModelDataForProvider
remain absent after the rebuild; use the concrete function names
(populateModelPoolFromPricingData, UpsertModelDataForProvider,
UpsertUnfilteredModelDataForProvider, DeleteModelDataForProvider) to locate code
and add table-driven tests covering multiple providers and mix of
duplicate/unique models.

🤖 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/main.go`:
- Line 463: The assignment mc.unfilteredModelPool[provider] = models shares the
slice backing array with mc.modelPool[provider]; change the assignment to create
a defensive copy (e.g., copy into a new slice using append([]Model(nil),
models...) or make+copy) when setting both mc.modelPool[provider] and
mc.unfilteredModelPool[provider] so each map entry has its own backing array and
future appends to one won't affect the other; update the code around where
models is assigned into mc.modelPool and mc.unfilteredModelPool (referencing
mc.modelPool, mc.unfilteredModelPool, provider, models) to use the copy
approach.
- Around line 472-485: The current union loops iterate
previousModelPool/previousUnfilteredModelPool and call slices.Contains against
mc.modelPool/mc.unfilteredModelPool which makes the merge O(m×n); instead, for
each provider build a lookup set (map[T]struct{}) of existing entries from
mc.modelPool[provider] (and separately for mc.unfilteredModelPool[provider])
then iterate the snapshot models and append only if the lookup misses; update
the loops that reference slices.Contains to use these per-provider maps to
reduce complexity to O(m+n) per provider.

---

Outside diff comments:
In `@framework/modelcatalog/main.go`:
- Around line 409-494: Add regression tests that exercise
ModelCatalog.populateModelPoolFromPricingData to verify snapshot-then-union
behavior: create a ModelCatalog test instance, seed mc.pricingData with a few
pricing entries, then call UpsertModelDataForProvider and
UpsertUnfilteredModelDataForProvider to add non-pricing models and also call
DeleteModelDataForProvider to remove one provider/model; invoke
populateModelPoolFromPricingData (or the public flow that triggers it) and
assert three things: (1) pricing models from mc.pricingData are present in
mc.modelPool/mc.unfilteredModelPool, (2) non-pricing models added via
UpsertModelDataForProvider / UpsertUnfilteredModelDataForProvider are preserved
after the rebuild, and (3) any models removed via DeleteModelDataForProvider
remain absent after the rebuild; use the concrete function names
(populateModelPoolFromPricingData, UpsertModelDataForProvider,
UpsertUnfilteredModelDataForProvider, DeleteModelDataForProvider) to locate code
and add table-driven tests covering multiple providers and mix of
duplicate/unique models.

🪄 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: d980c285-64d5-4faf-819b-8f2ee0750507

📥 Commits

Reviewing files that changed from the base of the PR and between a70594c and 658136d.

📒 Files selected for processing (1)

• framework/modelcatalog/main.go

[akshaydeo]

Merge activity

• Jun 2, 5:50 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• Jun 2, 5:50 PM UTC: @akshaydeo merged this pull request with Graphite.