feat: adds team budget and rl collection exporters

[roroghost17]

Summary

Exposes two new methods on LocalGovernanceStore — CollectTeamBudgets and CollectTeamRateLimits — so the enterprise layer can include team-level governance objects when building a user→team→business-unit hierarchy, in the same way collectBudgetsFromHierarchy folds them into the VK hierarchy.

Changes

• Added CollectTeamBudgets which resolves a team's configured budgets from the hot budgets map, returning live objects that reflect current usage counters and recent edits.
• Added CollectTeamRateLimits which resolves a team's configured rate limit (at most one) from the hot rate-limits map.
• Both methods mirror the existing read patterns in CheckTeamBudget and CheckTeamRateLimit respectively, and return nil gracefully when the team is unknown, has no associated governance objects, or the IDs cannot be resolved.
• Both methods are exported to allow the enterprise hierarchy collector to consume them without duplicating the store lookup logic.

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 test ./plugins/governance/...

Verify that a team with configured budgets and a rate limit returns the expected live objects from both methods, and that teams with no budgets, no rate limit, or unknown IDs return nil without panicking.

Breaking changes

• Yes
• No

Related issues

Security considerations

No new auth surfaces introduced. Both methods read from the existing in-memory store maps and follow the same access patterns as the existing Check* methods.

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
  • Team budget collection: retrieve live budget entries associated with a team for enterprise governance workflows.
  • Team rate-limit collection: access active rate-limit configurations tied to a team to support governance and enforcement.

[CLAassistant]

All committers have signed the CLA.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: ade0dd38-f018-45d3-97bc-3c5b7f8e6dd8

📥 Commits

Reviewing files that changed from the base of the PR and between 9529196 and b7d6813.

📒 Files selected for processing (1)

• plugins/governance/store.go

---

📝 Walkthrough

Walkthrough

Adds two exported helpers to LocalGovernanceStore: CollectTeamBudgets and CollectTeamRateLimits. Each loads a team from the in-memory teams map and resolves referenced budget or rate-limit IDs from gs.budgets / gs.rateLimits, returning live TableBudget/TableRateLimit pointers or nil when unresolved.

Changes

Team Governance Helper Methods

Layer / File(s)	Summary
Team budget and rate-limit collectors plugins/governance/store.go	CollectTeamBudgets resolves a team's budget objects from the budgets cache by budget IDs, and CollectTeamRateLimits resolves the team's rate-limit objects from the rate-limits cache; both return nil for empty team IDs, unknown teams, or missing referenced governance entries.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Suggested reviewers

• akshaydeo

Poem

A rabbit hops through governance halls,
Peeks in the caches, answers calls,
Budgets and limits it gently finds,
Nil-safe paths and tidy minds,
Hooray for helpers in memory-lined walls! 🐰✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Description check	❓ Inconclusive	The description covers most required sections but lacks evidence that tests were added or updated, and the checklist shows unfulfilled items despite the PR being open.	Complete the checklist items by adding tests for the new methods (covering edge cases like empty teamID, unknown teams, and missing cache entries) and verify all CI requirements are met before merge.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: adding two exported helper methods for collecting team budgets and rate limits.
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-03-feat_adds_team_budget_and_rl_collection_exporters

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.}

[roroghost17]

• fix: broaden Mantle model matching to all gpt variants #4091
• test: remove TestCollectApplicableGovernanceIDs_UserScopedModelConfigs #4090
• feat: add business unit CRUD, team assignment, and governance endpoints to OpenAPI spec #4082
• feat: add inline loading shell to #root before React mounts #4063
• fix: reduce FTS index left() cap from 800k to 250k chars to stay within tsvector limit #4057
• feat: add alias column to materialized view and filter support #4078
• feat: add "View logs" menu item to customer, team, and virtual key tables #4073
• feat: extract BudgetDisplay component with progress bar and tooltip #4055
• feat: add clickable links for routing rule, selected key, virtual key, and prompt in log detail views #4054
• fix: standardize page container layout with no-padding-parent and 100dvh sizing #4046
• fix: add colored tooltip values and subtle cursor highlight to dashboard charts #4044
• feat: add customer detail sheet and remove inline usage section from customer sheet #4026
• fix: replicates vk hierarchy flow for ghost nodes reconcilation #4088
• feat: supports multiple customers logs and usage tracking #4066
• fix: fixes virtual key model configs UI #4052
• fix: fixes cascade deletes for model configs #4051
• feat: adds support for multiple team and bu units in logstore #4041
• feat: adds team budget and rl collection exporters #4040 👈 (View in Graphite)
• dev

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

[greptile-apps]

Confidence Score: 4/5

Safe to merge; the new methods read from existing in-memory maps with the same guards as the Check* siblings and are correctly wired into the interface.

The logic is sound and the two previously identified issues (missing interface signatures, non-nil empty slice return) are both addressed. The only gap is the absence of unit tests for the new exported methods, which leaves the nil-vs-empty and resolution-failure edge cases unverified by automated tests.

plugins/governance/store.go — the new Collect* methods lack test coverage.

Important Files Changed

Filename	Overview
plugins/governance/store.go	Adds CollectTeamBudgets and CollectTeamRateLimits to both the GovernanceStore interface and LocalGovernanceStore; both methods correctly mirror the CheckTeam* read patterns, guard against empty-slice returns, and are nil-safe. No unit tests were added for the new methods.

_{Reviews (2): Last reviewed commit: "feat: adds team budget and rl collection..." | Re-trigger Greptile}

[roroghost17]

@coderabbitai full-review

[coderabbitai]

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Fetch the PR diff
gh pr view 4040 --repo maximhq/bifrost --json title,body,files,additions,deletions
gh pr diff 4040 --repo maximhq/bifrost

Length of output: 5049

---

🏁 Script executed:

#!/bin/bash
# Get CheckTeamBudget, CheckTeamRateLimit, LoadBudget, LoadRateLimit implementations
rg -n "func.*LocalGovernanceStore.*CheckTeamBudget|func.*LocalGovernanceStore.*CheckTeamRateLimit|func.*LocalGovernanceStore.*LoadBudget|func.*LocalGovernanceStore.*LoadRateLimit|GovernanceStore interface" plugins/governance/store.go -A 30

# Check if there's an interface definition that these methods should be added to
rg -n "GovernanceStore\|GovernancePlugin\|interface" plugins/governance/ -l

Length of output: 8951

---

🏁 Script executed:

#!/bin/bash
# Check if GovernanceStore interface includes the new methods, and if there are test files
rg -n "CollectTeamBudgets\|CollectTeamRateLimits" plugins/governance/ -C2

# Look at the interface to see lines 99-240 for full scope
sed -n '99,240p' plugins/governance/store.go

# Check for test files and existing team-related tests
fd "_test.go" plugins/governance/
rg -n "CollectTeam\|CheckTeamBudget\|CheckTeamRateLimit" plugins/governance/ --type go

Length of output: 12189

---

---

@roroghost17 — full review below.

---

Review: feat: adds team budget and rl collection exporters (#4040)

The implementation is clean and the logic correctly mirrors the existing CheckTeamBudget/CheckTeamRateLimit read patterns. A few items to address:

---

🔴 GovernanceStore interface not updated

Both new methods are added only as concrete *LocalGovernanceStore methods but are absent from the GovernanceStore interface. Every other CheckTeam*/Check* method is in the interface; this breaks the pattern. More importantly, the PR says these are consumed by the enterprise hierarchy layer — if that layer holds a GovernanceStore interface value (rather than a concrete *LocalGovernanceStore), it cannot call CollectTeamBudgets or CollectTeamRateLimits without a type assertion, which would silently break for any alternative store implementation (mocks, test doubles, future implementations).

// GovernanceStore interface — add the two new methods in the "Team level" section:
// Team level collection exporters (for enterprise hierarchy assembly)
CollectTeamBudgets(ctx context.Context, teamID string) []*configstoreTables.TableBudget
CollectTeamRateLimits(ctx context.Context, teamID string) []*configstoreTables.TableRateLimit

---

🟡 Doc comment inaccuracy: empty slice vs nil

CollectTeamBudgets states "Returns nil when the team is unknown or has no budgets", but when team.Budgets is non-empty yet every budget ID is a hot-cache miss, the function returns a non-nil empty slice []. Callers that check == nil would see a false positive, and callers that check len(...) == 0 would be fine. Consider either:

• Adding a final if len(list) == 0 { return nil } to make the behaviour match the documented contract, or
• Updating the doc to "Returns nil or an empty slice when …"

---

🟡 No tests added

plugins/governance/store_test.go and plugins/governance/storeconcurrency_test.go exist and cover the existing Check* methods. The checklist item "I added/updated tests where appropriate" is unchecked. At minimum, table-driven cases for the following should be covered:

Scenario	CollectTeamBudgets	CollectTeamRateLimits
Empty teamID	nil	nil
Unknown teamID	nil	nil
Team with no budgets / no RL	nil/empty	nil
Team with budgets all in hot cache	non-empty slice	—
Team with RL in hot cache	—	single-element slice
Team with budgets, all cache-miss	empty/nil	—

---

✅ What looks good

• Guard logic and nil-safety are correct and slightly more defensive than the mirrored Check* counterparts (explicit team == nil check, explicit empty teamID guard in CollectTeamRateLimits).
• Returning a slice for CollectTeamRateLimits (even though at most one entry) is consistent with EntityWiseRateLimits's map-of-slice shape, making it straightforward for the hierarchy collector to consume.
• ctx threading through to LoadBudget/LoadRateLimit is correct for API consistency even though the underlying sync.Map calls don't use it today.
• No new auth surfaces, purely read-path additions.

---

🧠 Learnings used

Learnt from: roroghost17
Repo: maximhq/bifrost PR: 3998
File: framework/configstore/migrations.go:9574-9591
Timestamp: 2026-06-02T16:47:23.453Z
Learning: In maximhq/bifrost reviews, full-table scans in one-time boot migrations should generally be treated as non-blocking optimizations unless there is a clear scale, locking, memory, or correctness risk. Do not escalate such scans as major issues solely because they are broader than necessary; note them as optional improvements when appropriate.

---

Learnt from: roroghost17
Repo: maximhq/bifrost PR: 4070
File: plugins/governance/store.go:2897-2912
Timestamp: 2026-06-04T23:22:40.769Z
Learning: In maximhq/bifrost `plugins/governance/store.go`, `CreateTeamInMemory` intentionally shallow-copies the incoming `*TableTeam` (storing `&clone` for the team and `clone.RateLimit` pointer directly into `gs.rateLimits`). This is safe because the function is only reached via `ReloadTeam`, which passes a fresh `ConfigStore.GetTeam` result that callers only serialize (`SendJSON`) after the store call — it is never mutated post-store. Teams have no `hydrateVKGovernance`-style reverse-map step, so the `RateLimit` pointer aliasing has no observable effect. Team budgets are already stored as independent copies via `b := clone.Budgets[i]; gs.budgets.Store(b.ID, &b)`. Do not flag the `RateLimit` shallow-copy in `CreateTeamInMemory` as a defect unless a new caller that mutates the team after storing is introduced.

---

Learnt from: CR
Repo: maximhq/bifrost PR: 0
File: .greptile/rules.md:0-0
Timestamp: 2026-06-03T21:47:30.819Z
Learning: Applies to plugins/governance/**/*.go : Review budget, rate-limit, virtual key, and RBAC paths for fail-closed behavior where security is involved.

---

Learnt from: roroghost17
Repo: maximhq/bifrost PR: 3998
File: plugins/governance/store.go:642-653
Timestamp: 2026-06-02T16:47:33.210Z
Learning: In `plugins/governance/store.go` (`GetGovernanceData` customer snapshot section), the fallback to the original stale budget entry when `gs.budgets.Load(b.ID)` misses is intentional. Customer budgets are owned via a `CustomerID` FK with `OnDelete:CASCADE`, so a missing live cache entry during snapshot rebuild is more likely a transient timing window than a deleted row. Silently dropping it could make the customer appear budget-less. This differs from the VK/team budget refresh pattern (which drops missing entries), because VK/team associations have a different ownership model. Do not flag this fallback as an inconsistency or suggest aligning it with the VK/team drop-on-miss pattern.

✅ Action performed

Full review finished.

[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 `@plugins/governance/store.go`:
- Around line 1270-1295: The LocalGovernanceStore method CollectTeamBudgets(ctx
context.Context, teamID string) []*configstoreTables.TableBudget is implemented
but missing from the GovernanceStore interface; update the GovernanceStore
interface (the "Team level" section) to include the signature
CollectTeamBudgets(ctx context.Context, teamID string)
[]*configstoreTables.TableBudget so callers using the GovernanceStore interface
can call it without type assertions and the implementation satisfies the
interface.
- Around line 1297-1318: The GovernanceStore interface is missing the
CollectTeamRateLimits method that's implemented on LocalGovernanceStore,
preventing callers with the interface type from invoking it; add the signature
CollectTeamRateLimits(ctx context.Context, teamID string)
[]*configstoreTables.TableRateLimit to the GovernanceStore interface in the
"Team level" section so the interface matches the concrete type and callers can
use CollectTeamRateLimits without a type assertion (ensure the method name and
signature exactly match the LocalGovernanceStore implementation).
- Around line 1272-1273: The doc for the function that "Returns nil when the
team is unknown or has no budgets" conflicts with behavior: when
gs.LoadBudget(ctx, b.ID) yields no budgets the function currently returns an
empty slice; update the function (the scope that builds and returns list, e.g.,
the method that iterates team.Budgets and appends loaded budgets) to return nil
instead of an empty slice by adding a check like "if len(list) == 0 { return
nil, nil }" just before the current return, so callers receive nil when no
budgets were loaded; adjust only the return behavior, not the doc.

🪄 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: 39aec221-5c16-439c-8515-23b5b139daa0

📥 Commits

Reviewing files that changed from the base of the PR and between d754185 and 9529196.

📒 Files selected for processing (1)

• plugins/governance/store.go

[akshaydeo]

Merge activity

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