feat: governance plugin refactor #1154

Pratham-Mishra04 · 2025-12-22T10:22:12Z

Summary

Briefly explain the purpose of this PR and the problem it solves.

Changes

What was changed and why
Any notable design decisions or trade-offs

Type of change

Affected areas

How to test

Describe the steps to validate this change. Include commands and expected outcomes.

# Core/Transports
go version
go test ./...

# UI
cd ui
pnpm i || npm i
pnpm test || npm test
pnpm build || npm run build

If adding new configs or environment variables, document them here.

Screenshots/Recordings

If UI changes, add before/after screenshots or short clips.

Breaking changes

Yes
No

If yes, describe impact and migration instructions.

Related issues

Link related issues and discussions. Example: Closes #123

Security considerations

Note any security implications (auth, secrets, PII, sandboxing, etc.).

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

Pratham-Mishra04 · 2025-12-22T10:22:35Z

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

coderabbitai · 2025-12-22T10:22:39Z

📝 Walkthrough

Summary by CodeRabbit

New Features
- GET endpoints for budgets and rate-limits with a "from_memory" option.
- Enterprise override hooks to customize governance and pricing loading.
Improvements
- Reworked in-memory governance for hierarchical rate-limit and budget tracking, faster resets/dumps, and immediate visibility after create/update/delete.
- Transactional deletions and a retry helper for more reliable config operations.
- Broad dependency updates across platform components.
Bug Fixes
- Clarified budget-exceeded error message.
Documentation
- Changelog formatting cleaned up.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

Walkthrough

Refactors governance into a GovernanceStore interface with a LocalGovernanceStore and in-memory hierarchical rate-limit/budget handling; threads EnterpriseOverrides through config/server/plugin/pricing loading; adds transactional cascade deletions and retry helpers in ConfigStore; introduces GORM AfterFind runtime fields; adds HTTP endpoints for budgets and rate-limits; many dependency updates.

Changes

Cohort / File(s)	Summary
Editorconfig & docs `\.editorconfig`, `docs/changelogs/v1.3.47.mdx`, `docs/features/governance/virtual-keys.mdx`, `ui/lib/constants/logs.ts`	Add Go-specific .editorconfig block; formatting tweaks in changelog and TS constants; adjust Budget error message text.
Core `core/go.mod`, `core/mcp.go`, `core/providers/nebius/nebius.go`, `core/schemas/context.go`	Dependency bumps; remove strict type assertion in in-process MCP client creation; initialize `NebiusProvider.sendBackRawRequest` from config; lazily initialize `BifrostContext.userValues` in `SetValue`.
Framework configstore (persistence) `framework/configstore/rdb.go`, `framework/configstore/store.go`	Add transactional cascade deletions for providers/teams/customers/virtual-keys and provider configs; add `GetRateLimits` and `RetryOnNotFound` APIs; update ConfigStore interface signatures.
Framework DB models `framework/configstore/tables/budget.go`, `framework/configstore/tables/mcp.go`, `framework/configstore/tables/ratelimit.go`	Add runtime-only fields (`LastDBUsage`, `LastDBTokenUsage`, `LastDBRequestUsage`) with `AfterFind` hooks; remove isCodeModeClient injection into MCP config hash.
Framework module updates `framework/go.mod`	Upgrade framework dependencies (redis, qdrant, weaviate, OpenTelemetry, golang.org/x/*, grpc/protobuf, mongo driver, etc.).
Model catalog / pricing `framework/modelcatalog/main.go`, `framework/modelcatalog/sync.go`	Add `ShouldSyncPricingFunc` predicate to `ModelCatalog` and thread into `Init`; `syncPricing` may early-exit when predicate returns false.
Governance plugin core `plugins/governance/main.go`, `plugins/governance/resolver.go`, `plugins/governance/store.go`, `plugins/governance/tracker.go`, `plugins/governance/utils.go`	Major rewrite: introduce `GovernanceStore` interface and `LocalGovernanceStore` (formerly GovernanceStore type), add `GovernanceData`, in-memory rate-limit/budget caches and hierarchical checks, new in-memory CRUD/update/reset/dump APIs; update resolver/tracker to use in-memory flows; add `equalPtr` helper; adjust many public signatures.
Plugin dependency updates `plugins/*/go.mod` (many)	Widespread go.mod version bumps across plugins (AWS SDK v2, bytedance/sonic, fasthttp, mark3labs/mcp-go, go-openapi, opentelemetry, golang.org/x/*, grpc/protobuf, mongo driver, qdrant, redis, klauspost/compress); some indirect removals.
HTTP handlers & routes `transports/bifrost-http/handlers/governance.go`, `transports/bifrost-http/handlers/middlewares.go`	Add `GetGovernanceData()` to GovernanceManager; new GET endpoints for budgets and rate-limits; support `from_memory` query to return in-memory data; `TransportInterceptorMiddleware` now accepts `EnterpriseOverrides` and can bypass interception.
HTTP server & lib `transports/bifrost-http/lib/lib.go`, `transports/bifrost-http/lib/config.go`, `transports/bifrost-http/lib/config_test.go`, `transports/bifrost-http/server/server.go`	Add `EnterpriseOverrides` interface and thread it through `LoadConfig`, `LoadPlugins`, `RegisterAPIRoutes`, plugin/pricing loaders; add `DBLookupMaxRetries`/`DBLookupDelay` constants; update tests/mocks (`MockConfigStore.RetryOnNotFound`, `GetRateLimits`); add server methods to load governance/pricing and expose `GetGovernanceData`.
Transport module updates `transports/go.mod`	Update transport-level dependencies (sonic, fasthttp, mark3labs/mcp-go, OpenTelemetry, go-openapi, golang.org/x/*, grpc/protobuf, etc.).

Sequence Diagram(s)

sequenceDiagram
    participant Client as HTTP Client
    participant Server as BifrostHTTPServer
    participant Config as Config Loader
    participant Enterprise as EnterpriseOverrides
    participant GovPlugin as Governance Plugin
    participant LocalStore as LocalGovernanceStore
    participant DB as ConfigStore (DB)

    Client->>Server: Start / bootstrap
    Server->>Config: LoadConfig(ctx, dir, EnterpriseOverrides)
    Config->>Enterprise: GetGovernancePluginName / LoadPricingManager
    alt EnterpriseOverrides present
        Enterprise-->>Config: plugin name / pricing manager
    else
        Config-->>Config: default modelcatalog.Init
    end
    Config-->>Server: Config ready

    Server->>Server: LoadPlugins(ctx, config, EnterpriseOverrides)
    Server->>GovPlugin: LoadGovernancePlugin(...)
    GovPlugin->>LocalStore: NewLocalGovernanceStore(...)
    LocalStore->>DB: Load VKs/Teams/Budgets/RateLimits
    DB-->>LocalStore: Data
    LocalStore-->>GovPlugin: GovernanceData (in-memory)

    Client->>Server: GET /api/governance/budgets?from_memory=true
    Server->>GovPlugin: GetGovernanceData()
    GovPlugin->>LocalStore: GetGovernanceData()
    LocalStore-->>GovPlugin: In-memory budgets
    GovPlugin-->>Server: budgets
    Server-->>Client: budgets (from cache)

    Client->>Server: Request requiring governance
    Server->>GovPlugin: EvaluateRequest(...)
    GovPlugin->>LocalStore: CheckRateLimit(...) / CheckBudget(...)
    LocalStore-->>GovPlugin: Decisions
    GovPlugin-->>Server: allow/deny
    Server-->>Client: response

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Poem

🐇 I hopped through code with tabs and cheer,

Caches woke up, and budgets now steer,
Hooks peeked at DB footprints in the night,
Plugins learned their names by enterprise light,
A rabbit’s patch — small, nimble, and bright.

Pre-merge checks and finishing touches

❌ Failed checks (3 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is a template with no substantive content filled in; all sections remain unchecked and unfilled except for the template headers themselves.	Complete the description by filling in the Summary, Changes, Type of change selection, Affected areas, testing steps, breaking changes declaration, related issues, and security considerations.
Linked Issues check	⚠️ Warning	The PR contains extensive governance plugin refactoring (store interface redesign, hierarchical budget/rate-limit checking, in-memory caching, new resolver patterns) but is linked to issue #123 (Files API Support), which is unrelated to governance plugin changes.	Verify the correct linked issue; if governance refactoring is the actual objective, link to the appropriate issue instead of the Files API Support issue.
Out of Scope Changes check	⚠️ Warning	The PR contains extensive out-of-scope changes relative to the linked Files API Support issue (#123): governance plugin refactoring, in-memory store redesign, hierarchical budget/rate-limit checks, and multiple infrastructure-level updates unrelated to file uploads or RAG workflows.	Link the PR to the correct governance refactoring issue or clarify whether Files API Support is an actual objective; remove or justify unrelated changes.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'feat: governance plugin refactor' accurately describes the main change observed in the PR—a comprehensive refactoring of the governance plugin system with structural improvements to interfaces, stores, and manager patterns.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment
Commit unit tests in branch 12-08-feat_governance_plugin_refactors

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

coderabbitai

Actionable comments posted: 12

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (6)

plugins/semanticcache/go.mod (1)

18-109: Consider testing the compatibility of indirect dependency updates before merging.

The redis/go-redis v9.17 release updates the minimum required Go version to 1.21, and it should work with any Redis 7.0+. The mark3labs/mcp-go library introduced breaking changes in v0.29.0 regarding tool request handling, where request.Params.Arguments is no longer directly accessible as a map and must use the new GetArguments() method instead. Since the go.mod specifies v0.43.2 (after v0.29.0), any code using mcp-go must have already been adapted to these changes. For the go-openapi modules transitioning from v0.24.x to v0.25.x, ensure that any code relying on these indirect dependencies has been tested with the new versions.
transports/bifrost-http/handlers/governance.go (3)
896-896: Inconsistent error handling for ReloadVirtualKey.

Line 896 ignores the error from ReloadVirtualKey, while line 403 properly checks and handles it with a fallback. For consistency, consider handling the error here as well:
🔎 Suggested fix for consistency
-	h.governanceManager.ReloadVirtualKey(ctx, vk.ID)
+	_, err = h.governanceManager.ReloadVirtualKey(ctx, vk.ID)
+	if err != nil {
+		logger.Error("failed to reload virtual key: %v", err)
+	}
917-921: Complete the error comment explanation.

The comments on lines 919, 1173, and 1405 are incomplete: "But we ignore this error because its not" — not what? Please clarify why these in-memory removal errors are non-critical. For example, "because it's not critical if the in-memory cache is inconsistent" or "because the DB deletion already succeeded."
🔎 Suggested comment completion
-	// But we ignore this error because its not
+	// But we ignore this error because it's not critical; DB deletion succeeded and cache will resync
Also applies to: 1171-1175, 1403-1407

486-490: Verify the nil handling logic for TeamID and CustomerID.

The current implementation clears both TeamID and CustomerID when both fields are nil in the request. However, for pointer fields in Go, JSON unmarshaling cannot distinguish between:

Field omitted from request (should leave existing value unchanged)

Field explicitly set to null (should clear the value)

Both cases result in nil pointers, so the current logic at lines 486-490 will clear both associations even when the fields were simply omitted from the update request.

💡 Potential solutions

Option 1: Remove lines 486-490 entirely if the intent is to only update fields that are explicitly provided.

Option 2: Use a three-state pattern (pointer to pointer, or separate "present" bool flags) to distinguish between omitted, null, and valued.

Option 3: Document this behavior clearly in the API spec so users know that omitting both fields clears the associations.

The logic at lines 478-485 correctly updates individual fields when provided, but the lines 486-490 may cause unexpected clearing.
transports/bifrost-http/server/server.go (1)
212-322: Nil-safe governance plugin name in switch case

LoadPlugin takes an EnterpriseOverrides parameter that can be nil. The switch case case EnterpriseOverrides.GetGovernancePluginName(): evaluates the method call immediately upon entering the switch, causing a nil pointer dereference panic if EnterpriseOverrides is nil and name doesn't match an earlier case.

Precompute the governance plugin name before the switch to avoid this:
governanceName := governance.PluginName
if EnterpriseOverrides != nil {
    governanceName = EnterpriseOverrides.GetGovernancePluginName()
}

switch name {
case telemetry.PluginName:
    ...
case logging.PluginName:
    ...
case governanceName:
    ...
}
This allows nil EnterpriseOverrides without risking a panic, while preserving override behavior for enterprise builds.
plugins/governance/store.go (1)
764-818: Variable shadowing: inner loops overwrite outer loop index i.

The inner loops (lines 767, 773, 779, 785) all declare for i := range ..., shadowing the outer loop's i at line 764. After the inner loops complete, the outer i retains the inner loop's final value, causing virtualKeys[i] = *vk at line 818 to write to the wrong index.
🔎 Proposed fix - use different variable names
 	for i := range virtualKeys {
 		vk := &virtualKeys[i]

-		for i := range teams {
-			if vk.TeamID != nil && teams[i].ID == *vk.TeamID {
-				vk.Team = &teams[i]
+		for j := range teams {
+			if vk.TeamID != nil && teams[j].ID == *vk.TeamID {
+				vk.Team = &teams[j]
 			}
 		}

-		for i := range customers {
-			if vk.CustomerID != nil && customers[i].ID == *vk.CustomerID {
-				vk.Customer = &customers[i]
+		for j := range customers {
+			if vk.CustomerID != nil && customers[j].ID == *vk.CustomerID {
+				vk.Customer = &customers[j]
 			}
 		}

-		for i := range budgets {
-			if vk.BudgetID != nil && budgets[i].ID == *vk.BudgetID {
-				vk.Budget = &budgets[i]
+		for j := range budgets {
+			if vk.BudgetID != nil && budgets[j].ID == *vk.BudgetID {
+				vk.Budget = &budgets[j]
 			}
 		}

-		for i := range rateLimits {
-			if vk.RateLimitID != nil && rateLimits[i].ID == *vk.RateLimitID {
-				vk.RateLimit = &rateLimits[i]
+		for j := range rateLimits {
+			if vk.RateLimitID != nil && rateLimits[j].ID == *vk.RateLimitID {
+				vk.RateLimit = &rateLimits[j]
 			}
 		}
Also fix the nested loops at lines 798 and 808 (use k instead of reusing variable names).

🧹 Nitpick comments (11)

plugins/governance/utils.go (1)
18-25: LGTM! Clean and correct pointer comparison utility.

The implementation correctly handles all nil cases before dereferencing, and the generic constraint is appropriate. The logic is sound:

Returns true when both pointers are nil

Returns false when only one is nil

Safely dereferences and compares values when both are non-nil

Consider verifying that this utility is actively used in the governance refactor and has test coverage:
#!/bin/bash
# Description: Verify equalPtr usage and test coverage in the governance module

# Search for equalPtr usage in the codebase
echo "=== equalPtr usage ==="
rg -n 'equalPtr' --type go -g '!*.md'

# Search for test coverage of equalPtr
echo -e "\n=== equalPtr tests ==="
rg -n 'equalPtr|TestEqual' --type go -g '*_test.go' plugins/governance/
transports/bifrost-http/handlers/governance.go (1)
188-197: Consider extracting the from_memory pattern into a helper.

The from_memory query parameter handling is duplicated across 8 handlers with identical logic. While this works, extracting it into a helper function would reduce duplication and make future changes easier.
💡 Example refactor

You could extract a helper like:
func shouldUseMemory(ctx *fasthttp.RequestCtx) bool {
    return string(ctx.QueryArgs().Peek("from_memory")) == "true"
}
Or for more type safety, consider parsing boolean query parameters properly to handle values like "1", "True", "false", etc.
Also applies to: 418-432, 942-964, 1047-1060, 1193-1202, 1281-1294, 1425-1434, 1449-1458
transports/bifrost-http/lib/config_test.go (2)

130-137: MockConfigStore additions are functionally correct but underutilize internal state

RetryOnNotFound correctly satisfies the new ConfigStore contract and preserves the underlying function’s behavior, but it ignores maxRetries and retryDelay. That’s fine for most unit tests, yet if you later add tests around retry/backoff behavior this stub won’t reflect production semantics.

GetRateLimits currently returns an empty slice instead of reading from m.governanceConfig.RateLimits, unlike CreateRateLimit which populates that field. Wiring it through to the stored state would make the mock more faithful without much effort.

Consider minimally delegating GetRateLimits to m.governanceConfig now, and only enriching RetryOnNotFound if/when tests need retry semantics.

Also applies to: 194-196, 317-319

6675-6685: Confirm LoadConfig’s third parameter is safely nilable in all test call sites

All updated tests now call LoadConfig(ctx, tempDir, nil) to satisfy the extended signature that threads EnterpriseOverrides. This is fine as long as the EnterpriseOverrides type is an interface or pointer (or an alias that accepts nil); if it is (or becomes) a concrete struct, these calls will stop compiling or require a zero value instead.

Please double‑check the EnterpriseOverrides type in config.go and confirm that nil is the intended “no overrides” sentinel for both production and tests. If you expect to exercise overrides in tests later, consider adding a small helper like loadConfigNoOverrides(t, dir) to centralize the default argument.

Also applies to: 6713-6717, 6724-6733, 6756-6760, 6774-6778, 6808-6812, 6873-6883, 6935-6940, 6970-6978, 7020-7047, 7087-7097, 7147-7155, 7217-7223, 7257-7277, 7298-7338, 7360-7368, 7417-7440, 7497-7532, 7598-7615, 7644-7648, 7665-7673, 7718-7724, 7760-7793, 7861-7865, 7927-7939, 7972-8073, 8081-8321, 8351-8447, 86629-86638, 8717-8722, 8795-8803, 8897-8904, 8961-8968, 9019-9026, 9089-9096, 9190-9198, 9217-9225, 9256-9264, 9290-9298, 9339-9347, 9472-9477, 95828-95838, 9639-9647, 9688-9696, 9751-9759, 9798-9806, 9854-9862, 9909-9917, 9970-9978, 10027-10035

plugins/governance/tracker.go (1)

127-148: Startup reset DB transaction looks correct but always updates all usage fields

The new startup reset flow—with ResetExpired*InMemory followed by a transactional UPDATE that only touches usage and *_last_reset columns and skips hooks—is a solid improvement to avoid clobbering max_limit / reset_duration. Note that you always send all four usage/last_reset columns in updates even if only token or request limits changed; this is harmless but slightly noisier. If you care about minimizing writes, you could conditionally add keys to updates based on which side actually expired.

Also applies to: 221-255

transports/bifrost-http/lib/config.go (1)

47-50: EnterpriseOverrides integration for pricing looks sound

The new EnterpriseOverrides plumbing through LoadConfig, initFrameworkConfigFromFile, and initDefaultFrameworkConfig cleanly lets enterprise builds override pricing‑manager initialization while preserving the previous modelcatalog.Init behavior when overrides are nil. No functional issues stand out here; the only minor nit is that DBLookupMaxRetries/DBLookupDelay are config‑layer constants but are only used from the server package—if you later add more DB retry callers, keeping them here is fine; otherwise consider co-locating them with the retry call site.

Also applies to: 275-355, 1505-1546, 1574-1629, 1866-1927

transports/bifrost-http/server/server.go (1)

324-439: LoadPlugins/ReloadPlugin: EnterpriseOverrides threading looks good once LoadPlugin is nil‑safe

Passing EnterpriseOverrides through LoadPlugins and ReloadPlugin keeps governance and any enterprise overrides centralized without changing behaviour for other plugins. Once LoadPlugin is made safe for EnterpriseOverrides == nil, these call sites will also behave sensibly in tests or non‑server contexts that might pass nil.

Also applies to: 840-887

plugins/governance/store.go (4)

41-63: Large interface may benefit from segregation in the future.

The GovernanceStore interface has 22 methods, which is quite large. While acceptable for now given the coherent domain, consider splitting into smaller interfaces (e.g., BudgetChecker, RateLimitChecker, GovernanceUpdater) if this interface grows further or if consumers only need subsets of functionality.

87-140: Consider returning copies to prevent external mutation.

GetGovernanceData() returns pointers to the actual stored objects. External callers can mutate these, potentially causing data corruption. If immutability is required, consider returning deep copies.

830-834: Reassigning sync.Map fields is not atomic.

Replacing the sync.Map instances (lines 830-834) is not thread-safe if other goroutines are concurrently accessing them. This appears safe during initialization, but if a runtime refresh feature is added later, this would need synchronization (e.g., a mutex or atomic pointer swap).

1037-1149: Consider extracting helper methods for readability.

UpdateVirtualKeyInMemory is ~112 lines with deeply nested logic. Consider extracting sections like "update VK budget," "update VK rate limit," and "update provider configs" into separate helper methods to improve readability and testability.

📜 Review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between bd2355b and bc1bdce.

⛔ Files ignored due to path filters (11)

core/go.sum is excluded by !**/*.sum
framework/go.sum is excluded by !**/*.sum
plugins/governance/go.sum is excluded by !**/*.sum
plugins/jsonparser/go.sum is excluded by !**/*.sum
plugins/logging/go.sum is excluded by !**/*.sum
plugins/maxim/go.sum is excluded by !**/*.sum
plugins/mocker/go.sum is excluded by !**/*.sum
plugins/otel/go.sum is excluded by !**/*.sum
plugins/semanticcache/go.sum is excluded by !**/*.sum
plugins/telemetry/go.sum is excluded by !**/*.sum
transports/go.sum is excluded by !**/*.sum

📒 Files selected for processing (36)

.editorconfig
core/go.mod
core/mcp.go
core/providers/nebius/nebius.go
core/schemas/context.go
docs/changelogs/v1.3.47.mdx
docs/features/governance/virtual-keys.mdx
framework/configstore/rdb.go
framework/configstore/store.go
framework/configstore/tables/budget.go
framework/configstore/tables/mcp.go
framework/configstore/tables/ratelimit.go
framework/go.mod
framework/modelcatalog/main.go
framework/modelcatalog/sync.go
plugins/governance/go.mod
plugins/governance/main.go
plugins/governance/resolver.go
plugins/governance/store.go
plugins/governance/tracker.go
plugins/governance/utils.go
plugins/jsonparser/go.mod
plugins/logging/go.mod
plugins/maxim/go.mod
plugins/mocker/go.mod
plugins/otel/go.mod
plugins/semanticcache/go.mod
plugins/telemetry/go.mod
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/lib/lib.go
transports/bifrost-http/server/server.go
transports/go.mod
ui/lib/constants/logs.ts

💤 Files with no reviewable changes (1)

framework/configstore/tables/mcp.go

🧰 Additional context used

📓 Path-based instructions (1)

**

⚙️ CodeRabbit configuration file

always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)

Files:

core/schemas/context.go
framework/modelcatalog/sync.go
framework/configstore/tables/ratelimit.go
core/mcp.go
transports/bifrost-http/lib/lib.go
framework/configstore/store.go
core/providers/nebius/nebius.go
transports/bifrost-http/handlers/middlewares.go
plugins/mocker/go.mod
framework/configstore/rdb.go
framework/configstore/tables/budget.go
docs/features/governance/virtual-keys.mdx
plugins/logging/go.mod
plugins/telemetry/go.mod
transports/bifrost-http/handlers/governance.go
plugins/maxim/go.mod
plugins/semanticcache/go.mod
framework/go.mod
plugins/governance/resolver.go
framework/modelcatalog/main.go
plugins/jsonparser/go.mod
ui/lib/constants/logs.ts
plugins/governance/utils.go
plugins/governance/go.mod
plugins/otel/go.mod
docs/changelogs/v1.3.47.mdx
core/go.mod
transports/go.mod
plugins/governance/main.go
transports/bifrost-http/lib/config.go
plugins/governance/tracker.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/server/server.go
plugins/governance/store.go

🧠 Learnings (3)

📚 Learning: 2025-12-09T17:07:42.007Z

Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.

Applied to files:

core/schemas/context.go
framework/modelcatalog/sync.go
framework/configstore/tables/ratelimit.go
core/mcp.go
transports/bifrost-http/lib/lib.go
framework/configstore/store.go
core/providers/nebius/nebius.go
transports/bifrost-http/handlers/middlewares.go
framework/configstore/rdb.go
framework/configstore/tables/budget.go
transports/bifrost-http/handlers/governance.go
plugins/governance/resolver.go
framework/modelcatalog/main.go
plugins/governance/utils.go
plugins/governance/main.go
transports/bifrost-http/lib/config.go
plugins/governance/tracker.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/server/server.go
plugins/governance/store.go

📚 Learning: 2025-12-12T08:25:02.629Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.

Applied to files:

transports/bifrost-http/lib/lib.go
transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/server/server.go

📚 Learning: 2025-12-19T09:26:54.961Z

Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/utils/utils.go:1050-1051
Timestamp: 2025-12-19T09:26:54.961Z
Learning: Update streaming end-marker handling so HuggingFace is treated as a non-[DONE] provider for backends that do not emit a DONE marker (e.g., meta llama on novita). In core/providers/utils/utils.go, adjust ProviderSendsDoneMarker() (or related logic) to detect providers that may not emit DONE and avoid relying on DONE as the sole end signal. Add tests to cover both DONE-emitting and non-DONE backends, with clear documentation in code comments explaining the rationale and any fallback behavior.

Applied to files:

core/providers/nebius/nebius.go

🧬 Code graph analysis (12)

transports/bifrost-http/lib/lib.go (5)

transports/bifrost-http/lib/config.go (1)

Config (189-220)

plugins/semanticcache/main.go (2)

Config (25-43)

Plugin (137-144)

core/schemas/plugin.go (1)

Plugin (61-87)

framework/configstore/store.go (1)

ConfigStore (18-157)

framework/modelcatalog/main.go (1)

ModelCatalog (23-46)

framework/configstore/store.go (1)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

transports/bifrost-http/handlers/middlewares.go (2)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

transports/bifrost-http/lib/middleware.go (1)

BifrostHTTPMiddleware (7-7)

framework/configstore/rdb.go (5)

framework/configstore/tables/virtualkey.go (6)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

TableVirtualKeyMCPConfig (81-87)

TableVirtualKeyMCPConfig (90-92)

framework/configstore/tables/budget.go (2)

TableBudget (11-27)

TableBudget (30-30)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

framework/configstore/tables/team.go (2)

TableTeam (12-38)

TableTeam (41-41)

framework/configstore/tables/customer.go (2)

TableCustomer (6-22)

TableCustomer (25-25)

framework/configstore/tables/budget.go (1)

framework/configstore/tables/utils.go (1)

ParseDuration (9-43)

transports/bifrost-http/handlers/governance.go (2)

plugins/governance/store.go (1)

GovernanceData (33-39)

framework/configstore/tables/virtualkey.go (2)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

plugins/governance/resolver.go (4)

plugins/governance/store.go (1)

GovernanceStore (41-63)

framework/configstore/tables/virtualkey.go (2)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

core/schemas/bifrost.go (1)

ModelProvider (32-32)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

framework/modelcatalog/main.go (2)

transports/bifrost-http/lib/config.go (1)

Config (189-220)

framework/configstore/tables/modelpricing.go (2)

TableModelPricing (4-41)

TableModelPricing (44-44)

transports/bifrost-http/lib/config.go (1)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

plugins/governance/tracker.go (3)

plugins/governance/store.go (1)

GovernanceStore (41-63)

plugins/governance/resolver.go (1)

BudgetResolver (64-67)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

transports/bifrost-http/lib/config_test.go (1)

transports/bifrost-http/lib/config.go (1)

LoadConfig (275-355)

transports/bifrost-http/server/server.go (4)

plugins/governance/store.go (1)

GovernanceData (33-39)

plugins/governance/main.go (5)

Config (32-34)

PluginName (21-21)

BaseGovernancePlugin (40-47)

GovernancePlugin (50-69)

Init (104-158)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

transports/bifrost-http/handlers/middlewares.go (1)

TransportInterceptorMiddleware (49-144)

🔇 Additional comments (46)

plugins/logging/go.mod (1)

1-112: Request visibility into actual source code changes beyond dependency updates.

The provided file contains only dependency version bumps (AWS SDK v2, Sonic, go-openapi, OpenTelemetry, gRPC, protobuf, and stdlib packages), but the PR objectives and AI summary describe significant refactoring work:

Governance plugin refactor (GovernanceStore → LocalGovernanceStore)

Deadlock fix (per commit message)

Database model enhancements with virtual runtime fields and AfterFind hooks

EnterpriseOverrides pattern introduction

New governance API endpoints for budgets and rate-limits

To provide a comprehensive review aligned with the coding guidelines ("see all changes in light of the whole stack"), I need visibility into the actual source code changes, not just the dependency manifest. The go.mod changes alone do not reveal the correctness, security posture, or potential breaking changes associated with the refactoring work.

Please provide:

Source code files implementing the governance plugin refactor (e.g., new/modified GovernanceStore, LocalGovernanceStore implementations)

Files addressing the deadlock fix mentioned in commit messages

Database model changes (files with virtual runtime fields and AfterFind hooks)

New/modified API endpoint handlers (governance budgets, rate-limits)

Any breaking changes or migration notes

Without these files, I cannot assess the full scope of the refactoring against best practices, security, and correctness.

plugins/semanticcache/go.mod (1)

1-114: Clarify scope: confirm this file represents only transitive dependency updates for the governance plugin refactor.

The provided file contains only indirect dependency version bumps with no direct code changes. However, the PR objectives describe concrete code changes: GovernanceStore interface, LocalGovernanceStore implementation, database model enhancements, and new governance API endpoints. This file should be reviewed alongside the actual refactored source files to validate that dependency upgrades align with and support the refactored code.

Can you confirm:

Is this the only file changed in the semanticcache plugin for this PR, or are there additional source files (e.g., implementations of GovernanceStore, API handlers) that should be reviewed?

Are the dependency bumps prerequisite updates for the governance refactor, or are they independent maintenance updates?

plugins/mocker/go.mod (2)

1-59: Clarify PR scope: governance refactor vs. transitive dependency updates.

This file contains only indirect transitive dependency version updates with no changes to the module's direct dependencies or its own code. However, the PR title and objectives focus on "governance plugin refactor" and "Files API Support." The relationship between these dependency updates and the stated PR objectives is unclear.

Per the coding guidelines, ensure that you are reviewing this change in the context of the full PR stack and that all modules (not just the mocker plugin) have aligned dependency versions, particularly for critical packages like AWS SDK (used for S3/Files API) and the core framework.

Confirm:

Are similar dependency updates present across other modules (governance, core, framework, etc.) in this PR?

Does the AWS SDK update relate to the Files API Support objective mentioned in the PR?

Is this a coordinated, repository-wide dependency refresh, or are there specific code changes in this PR that necessitate these updates?

15-16: Verify mcp-go breaking changes impact; sonic/loader update appears safe.

MCP-Go has documented breaking changes starting from v0.29.0 that changed how tool request arguments are handled, requiring migration from direct access to the GetArguments() method. The bump from v0.41.1 → v0.43.2 spans multiple releases and warrants verification that any code using mcp-go has been updated to handle these API changes.

In contrast, the sonic/loader v0.4.0 update appears to be a maintenance release without documented breaking changes. This update is lower risk.

Verify that any direct usage of mcp-go's tool request handling in the mocker plugin accommodates the breaking changes introduced in the v0.29.0+ versions now included in this dependency bump.

core/providers/nebius/nebius.go (1)

52-52: LGTM!

The initialization of sendBackRawRequest is consistent with the pattern used for sendBackRawResponse on line 53.

.editorconfig (1)

8-10: LGTM!

The Go-specific editor configuration follows standard Go conventions (tabs with size 4).

docs/features/governance/virtual-keys.mdx (1)

618-618: LGTM!

The simplified error message is clearer and more concise by removing the redundant "Budget check failed: " prefix.

core/schemas/context.go (1)

174-176: LGTM!

The nil check prevents potential panics if userValues map is uninitialized. While the constructor (line 55) initializes the map, this defensive check protects against edge cases where the struct might be created without using the constructor.

ui/lib/constants/logs.ts (1)

153-153: LGTM!

Formatting cleanup with no functional changes.

docs/changelogs/v1.3.47.mdx (1)

20-41: LGTM!

Formatting normalization adds consistent spacing after list item markers throughout the changelog.

core/mcp.go (1)

1140-1140: Verify type of config.InProcessServer and confirm NewInProcessClient parameter compatibility.

The NewInProcessClient function requires a *server.MCPServer parameter. To ensure the type change is safe, confirm that config.InProcessServer is already typed as *server.MCPServer and that removing any previous type assertion doesn't introduce a type mismatch or defer validation to the client constructor.

plugins/jsonparser/go.mod (1)

12-13: Verify that dependency versions exist and are free from known vulnerabilities.

Multiple dependency versions have been updated. Ensure all versions are published on pkg.go.dev and do not contain known security advisories. Pay special attention to fasthttp versions (should be >= 1.34.0 to avoid directory traversal CVE-2022-21221, and >= 1.53.0 to avoid open redirect issues).

framework/modelcatalog/main.go (2)

81-85: LGTM! Well-designed optional callback pattern.

The ShouldSyncPricingFunc type provides a clean, optional hook for external control over pricing synchronization. The documentation clearly indicates it's optional and can be nil.

88-108: LGTM! Callback properly wired into initialization.

The Init function signature is correctly updated to accept the shouldSyncPricingFunc parameter, and it's properly assigned to the ModelCatalog instance at line 107. This enables external callers to influence pricing sync behavior.

framework/modelcatalog/sync.go (1)

61-66: LGTM! Callback properly gated with nil check.

The early-exit logic correctly checks if shouldSyncPricingFunc is non-nil before invoking it. If the callback returns false, pricing sync is cancelled with appropriate debug logging. This maintains backward compatibility when the callback is nil.

core/go.mod (1)

8-22: Dependency versions verified—all exist and are secure.

All 15 dependencies have been verified on the Go module proxy and confirmed available. Security advisory review found no known vulnerabilities in the specified versions:

AWS SDK v2 components, Smithy, and other libraries are at current safe versions

golang.org/x/oauth2 v0.34.0 and golang.org/x/text v0.32.0 contain all historical security fixes

No further action needed.

plugins/otel/go.mod (2)

5-117: Dependency version bumps look consistent.

The version bumps across grpc, protobuf, AWS SDK, OpenAPI tooling, and otel packages are minor/patch increments and align with the broader dependency refresh described in the PR summary.

3-3: No action needed—Go 1.25.5 is valid and current.

Go 1.25.5 was released on December 2, 2025 and is a legitimate, stable version. No change to the Go version directive is necessary.

Likely an incorrect or invalid review comment.

framework/configstore/store.go (3)

7-7: LGTM: Time import added for RetryOnNotFound.

The time import is correctly added to support the new RetryOnNotFound method signature.

103-108: LGTM: GetRateLimits interface method.

The new GetRateLimits method follows the established pattern of GetBudgets and other collection retrieval methods in the interface.

146-148: LGTM: RetryOnNotFound interface method.

The retry wrapper with configurable maxRetries and retryDelay parameters provides flexibility for callers to tune retry behavior based on their use case.

framework/configstore/tables/ratelimit.go (2)

32-36: LGTM: Virtual runtime fields for tracking DB state.

The LastDBTokenUsage and LastDBRequestUsage fields with gorm:"-" json:"-" correctly implement non-persistent, non-serialized fields. This mirrors the pattern established in TableBudget.LastDBUsage and enables tracking delta changes between in-memory and persisted state.

82-88: LGTM: AfterFind hook implementation.

The AfterFind hook correctly captures the database state at fetch time into the virtual fields, enabling the caller to compute deltas between current in-memory usage and the last-known DB values.

framework/configstore/tables/budget.go (2)

24-27: LGTM: Virtual runtime field for budget tracking.

The LastDBUsage field follows the same pattern as the rate limit virtual fields, enabling delta tracking between in-memory and persisted budget usage.

47-52: LGTM: AfterFind hook for budget.

The hook correctly captures the persisted CurrentUsage value into LastDBUsage at fetch time.

transports/bifrost-http/handlers/middlewares.go (1)

59-69: LGTM: Governance plugin name resolution via EnterpriseOverrides.

The change to use enterpriseOverrides.GetGovernancePluginName() instead of a hardcoded constant decouples the middleware from the concrete governance plugin, enabling enterprise-specific customization.

transports/bifrost-http/lib/lib.go (1)

17-22: LGTM: EnterpriseOverrides interface design.

The interface is well-focused with three methods that provide clear extension points for enterprise customization:

Plugin name resolution

Governance plugin loading

Pricing manager loading

This follows Go's interface design principles of keeping interfaces small and purposeful.

framework/configstore/rdb.go (5)

1399-1464: LGTM: Transactional cascade deletion for virtual keys.

The refactored DeleteVirtualKey properly handles the deletion cascade:

Preloads ProviderConfigs to access related entities

Cleans up join table entries

Deletes associated budgets and rate limits from provider configs

Deletes provider and MCP configs

Deletes the virtual key's own budget and rate limit

Finally deletes the virtual key itself

This ensures referential integrity is maintained throughout the deletion process.

1665-1699: LGTM: Transactional team deletion with FK nullification.

The cascade logic correctly:

Nullifies team_id on virtual keys (preserving them rather than deleting)

Deletes the team's associated budget

Deletes the team

This prevents orphaned references while preserving dependent virtual keys.

1754-1792: LGTM: Transactional customer deletion with FK nullification.

Similar to team deletion, this properly:

Nullifies customer_id on virtual keys and teams

Deletes the customer's budget

Deletes the customer

This maintains consistency with the team deletion pattern.

1794-1804: LGTM: GetRateLimits implementation.

The implementation follows the established pattern of GetBudgets with proper ErrNotFound handling.

2154-2179: The syntax for attempt := range maxRetries (line 2157) is compatible with the project's Go versions. All go.mod files declare Go 1.24.3 or higher (with 1.25.5 being the standard across most modules), and these versions support range over integers, which was introduced in Go 1.22.

transports/bifrost-http/handlers/governance.go (5)

25-25: LGTM! Clean interface extension.

The addition of GetGovernanceData() to the GovernanceManager interface enables in-memory data access, which aligns with the new from_memory query parameter support throughout the handlers.

42-44: LGTM! Good defensive validation.

The nil check for the governance manager prevents potential nil pointer dereferences and provides clear error messaging.

178-181: LGTM! Read-only routes for aggregate views.

The new budget and rate-limit endpoints provide read-only access, which is appropriate since these resources are managed through their parent entities (Virtual Keys, Teams, Customers, Provider Configs).

698-711: LGTM! Well-documented nil handling for KeyIDs.

The comment clearly explains the three-state logic for KeyIDs:

nil (field omitted): leaves existing keys unchanged

empty slice (field present but empty): clears all keys

non-empty slice: updates to the specified keys

This design correctly supports partial updates.

1423-1469: LGTM! Clean implementation of new handlers.

The getBudgets and getRateLimits handlers follow the established patterns in this file, including proper error handling and from_memory parameter support.

transports/bifrost-http/lib/config_test.go (3)

627-633: testLogger stub correctly tracks the expanded logger interface

The no-op implementations for Debug/Info/Warn/Error/Fatal and SetLevel cleanly satisfy the updated logger interface and keep SQLite integration tests quiet; there’s nothing to change here.

1832-1844: Key comparison broadened to include name equality looks intentional

Changing the duplicate-key detection to dbKeyHash == fileKeyHash || fileKey.Name == dbKey.Name ensures a DB key isn’t treated as “file-missing” (and thus preserved as dashboard-only) when only value/models/weight changed but the logical key identity (name) matches. This aligns with the test’s intent of preserving genuinely dashboard‑added keys without duplicating or dropping renamed ones.

2200-2207: Comment and message formatting tweaks have no behavioral impact

The various additions of // SAME/// CHANGED annotations and expanded t.Errorf messages in these ranges only improve readability and diagnostics; they don’t affect test logic or runtime behavior. No action needed.

Also applies to: 3012-3015, 3361-3366, 3617-3623, 3731-3740, 3872-3890, 41483-41516

transports/bifrost-http/server/server.go (3)

201-210: GovernanceInMemoryStore now safely reads providers under lock

Exporting Config and adding RLock/RUnlock around Config.Providers access in GetConfiguredProviders is a good fix; it avoids races with provider updates while keeping the hot path cheap (map read only). No further issues here.

662-675: GetGovernanceData implementation is safe and matches new GovernanceStore API

Using PluginsMutex.RLock() while resolving the governance plugin and then delegating to GetGovernanceStore().GetGovernanceData() is the right shape for exposing in‑memory governance state via ServerCallbacks. Returning nil on lookup failure or missing BaseGovernancePlugin support is also a reasonable, non‑crashing fallback.

957-1041: EnterpriseOverrides wiring in Bootstrap/API routes and middleware is coherent

Bootstrapping now passes the server instance as EnterpriseOverrides into LoadConfig, LoadPlugins, RegisterAPIRoutes, and TransportInterceptorMiddleware. Together with GetGovernancePluginName, LoadGovernancePlugin, and LoadPricingManager on BifrostHTTPServer, this cleanly centralizes any future enterprise overrides without changing OSS behaviour.

Once the governance plugin name case in LoadPlugin is made nil‑safe, this wiring looks solid.

Also applies to: 1075-1101, 1117-1125, 1239-1249

plugins/governance/store.go (3)

65-85: LGTM - Constructor follows good patterns.

The constructor properly handles both database-backed and config-memory initialization paths with appropriate error handling.

610-623: Good deadlock handling for multi-node scenarios.

The deadlock detection and graceful retry-next-cycle approach is a solid pattern for distributed systems where multiple nodes may compete for the same rows.

1440-1445: Verify equalPtr function is defined.

The function equalPtr is called at lines 1440-1445 to compare pointer values but its definition location is unclear. Confirm this helper function is defined elsewhere in the package, imported from another package, or needs to be added.

plugins/governance/main.go

plugins/governance/resolver.go

plugins/governance/store.go

plugins/governance/tracker.go

transports/bifrost-http/handlers/middlewares.go

transports/bifrost-http/server/server.go

coderabbitai

Actionable comments posted: 2

♻️ Duplicate comments (1)

transports/bifrost-http/handlers/middlewares.go (1)

49-73: LGTM! Nil-check for enterpriseOverrides properly addresses previous concern.

The added nil-guard at lines 58-61 prevents the nil pointer dereference that was flagged in a previous review. The control flow is correct: check plugins first, then enterpriseOverrides, then governance plugin presence.

🧹 Nitpick comments (5)

plugins/governance/resolver.go (1)

194-217: checkRateLimitHierarchy discards decision on success path but reconstructs rateLimitInfo on error.

The decision returned by CheckRateLimit at line 196 is used in the error result but is not validated on success. Additionally, the rate limit info is reconstructed by iterating provider configs (lines 199-207) rather than being returned from CheckRateLimit. Consider having CheckRateLimit return the relevant rate limit info to avoid this reconstruction.

plugins/governance/tracker.go (1)

47-49: Clarify/guard Dump behavior when running in in‑memory mode*

Init explicitly supports configStore == nil (governance runs in in‑memory‑only mode), but resetExpiredCounters unconditionally calls DumpRateLimits/DumpBudgets on every workerInterval tick. If those methods assume a non‑nil backing config store, this path could either panic or do unnecessary work when persistence is disabled.

Consider either:

guarding the dump calls with if t.configStore != nil { … }, or

ensuring the GovernanceStore’s DumpRateLimits/DumpBudgets implementations are explicit no‑ops when constructed with a nil configStore.

This keeps the “memory‑only” mode cheap and avoids surprises if the dump methods evolve.

Also applies to: 104-110, 127-148
transports/bifrost-http/server/server.go (3)
201-210: Guard MCPServerHandler usage in governance reload/remove helpers

The new getGovernancePlugin helper plus in‑memory updates for VK/Team/Customer look good and fix the earlier unsynchronized plugin access. One remaining edge case: ReloadVirtualKey and RemoveVirtualKey call s.MCPServerHandler.SyncVKMCPServer / DeleteVKMCPServer without a nil‑check.

In setups where:

MCP is disabled, or

ReloadVirtualKey/RemoveVirtualKey are invoked before RegisterAPIRoutes initializes MCPServerHandler,

these calls could panic.

Consider a defensive guard:
Suggested nil‑check around MCPServerHandler calls
 func (s *BifrostHTTPServer) ReloadVirtualKey(ctx context.Context, id string) (*tables.TableVirtualKey, error) {
@@
-	governancePlugin.GetGovernanceStore().UpdateVirtualKeyInMemory(virtualKey)
-	s.MCPServerHandler.SyncVKMCPServer(virtualKey)
+	governancePlugin.GetGovernanceStore().UpdateVirtualKeyInMemory(virtualKey)
+	if s.MCPServerHandler != nil {
+		s.MCPServerHandler.SyncVKMCPServer(virtualKey)
+	}
 	return virtualKey, nil
 }
@@
 func (s *BifrostHTTPServer) RemoveVirtualKey(ctx context.Context, id string) error {
@@
-	governancePlugin.GetGovernanceStore().DeleteVirtualKeyInMemory(id)
-	s.MCPServerHandler.DeleteVKMCPServer(preloadedVk.Value)
+	governancePlugin.GetGovernanceStore().DeleteVirtualKeyInMemory(id)
+	if s.MCPServerHandler != nil {
+		s.MCPServerHandler.DeleteVKMCPServer(preloadedVk.Value)
+	}
 	return nil
 }
Also applies to: 519-572, 591-610, 627-646

499-517: Reuse getGovernancePlugin in GetGovernanceData to centralize lookup logic

You’ve added getGovernancePlugin() with proper locking and BaseGovernancePlugin assertion, but GetGovernanceData still does its own FindPluginByName + type assertion.

For consistency and to keep governance plugin lookup logic in one place, consider:

calling getGovernancePlugin() from GetGovernanceData, and

returning nil on error.

This avoids duplicating the locking and assertion pattern and makes future changes to governance plugin resolution easier to maintain.

Also applies to: 648-661

212-321: Clarify expectations around governance plugin name overrides and loading

The EnterpriseOverrides integration is a nice addition, but there’s a subtle constraint around governance plugin naming:
LoadPlugins delegates governance loading to EnterpriseOverrides.LoadGovernancePlugin, and uses EnterpriseOverrides.GetGovernancePluginName() for status names and for skipping the built‑in governance entry in config.PluginConfigs.
LoadPlugin keys the governance branch on EnterpriseOverrides.GetGovernancePluginName(), but BifrostHTTPServer.LoadGovernancePlugin always calls:
LoadPlugin[*governance.GovernancePlugin](ctx, governance.PluginName, nil, ..., config, s)
If OSSToEnterprisePluginNameOverrides[governance.PluginName] is ever set to a different string, LoadPlugin’s governance case won’t match the hard‑coded governance.PluginName, and LoadGovernancePlugin/LoadPlugins will start failing with “plugin not found” errors.

If you plan to support renaming or swapping the governance plugin via OSSToEnterprisePluginNameOverrides, consider tightening this contract, for example by:

either having LoadGovernancePlugin pass s.GetGovernancePluginName() as the name argument and adjusting the type assertion accordingly (e.g., to an interface implemented by both OSS and enterprise governance plugins), or

constraining OSSToEnterprisePluginNameOverrides so that, for the OSS server implementation, GetGovernancePluginName() must always equal governance.PluginName, and documenting that expectation.

Right now the override map is effectively unsafe to change for governance without additional wiring.

Also applies to: 325-404, 1061-1087

📜 Review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between bc1bdce and 5ca488c.

⛔ Files ignored due to path filters (11)

core/go.sum is excluded by !**/*.sum
framework/go.sum is excluded by !**/*.sum
plugins/governance/go.sum is excluded by !**/*.sum
plugins/jsonparser/go.sum is excluded by !**/*.sum
plugins/logging/go.sum is excluded by !**/*.sum
plugins/maxim/go.sum is excluded by !**/*.sum
plugins/mocker/go.sum is excluded by !**/*.sum
plugins/otel/go.sum is excluded by !**/*.sum
plugins/semanticcache/go.sum is excluded by !**/*.sum
plugins/telemetry/go.sum is excluded by !**/*.sum
transports/go.sum is excluded by !**/*.sum

📒 Files selected for processing (35)

.editorconfig
core/go.mod
core/mcp.go
core/providers/nebius/nebius.go
docs/changelogs/v1.3.47.mdx
docs/features/governance/virtual-keys.mdx
framework/configstore/rdb.go
framework/configstore/store.go
framework/configstore/tables/budget.go
framework/configstore/tables/mcp.go
framework/configstore/tables/ratelimit.go
framework/go.mod
framework/modelcatalog/main.go
framework/modelcatalog/sync.go
plugins/governance/go.mod
plugins/governance/main.go
plugins/governance/resolver.go
plugins/governance/store.go
plugins/governance/tracker.go
plugins/governance/utils.go
plugins/jsonparser/go.mod
plugins/logging/go.mod
plugins/maxim/go.mod
plugins/mocker/go.mod
plugins/otel/go.mod
plugins/semanticcache/go.mod
plugins/telemetry/go.mod
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/lib/lib.go
transports/bifrost-http/server/server.go
transports/go.mod
ui/lib/constants/logs.ts

💤 Files with no reviewable changes (1)

framework/configstore/tables/mcp.go

🚧 Files skipped from review as they are similar to previous changes (16)

core/providers/nebius/nebius.go
framework/configstore/rdb.go
docs/changelogs/v1.3.47.mdx
core/mcp.go
transports/bifrost-http/lib/lib.go
framework/go.mod
framework/configstore/store.go
ui/lib/constants/logs.ts
.editorconfig
plugins/mocker/go.mod
framework/modelcatalog/sync.go
plugins/jsonparser/go.mod
plugins/governance/utils.go
plugins/maxim/go.mod
plugins/otel/go.mod
transports/bifrost-http/lib/config_test.go

🧰 Additional context used

📓 Path-based instructions (1)

**

⚙️ CodeRabbit configuration file

always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)

Files:

docs/features/governance/virtual-keys.mdx
framework/modelcatalog/main.go
framework/configstore/tables/ratelimit.go
transports/bifrost-http/handlers/middlewares.go
core/go.mod
plugins/telemetry/go.mod
plugins/governance/resolver.go
plugins/governance/main.go
plugins/governance/tracker.go
transports/bifrost-http/lib/config.go
plugins/semanticcache/go.mod
plugins/governance/go.mod
transports/go.mod
transports/bifrost-http/handlers/governance.go
plugins/logging/go.mod
transports/bifrost-http/server/server.go
framework/configstore/tables/budget.go
plugins/governance/store.go

🧠 Learnings (3)

📚 Learning: 2025-12-09T17:07:42.007Z

Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.

Applied to files:

framework/modelcatalog/main.go
framework/configstore/tables/ratelimit.go
transports/bifrost-http/handlers/middlewares.go
plugins/governance/resolver.go
plugins/governance/main.go
plugins/governance/tracker.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/server/server.go
framework/configstore/tables/budget.go
plugins/governance/store.go

📚 Learning: 2025-12-12T08:25:02.629Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.

Applied to files:

transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/server/server.go

📚 Learning: 2025-12-22T10:50:40.990Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.

Applied to files:

plugins/governance/resolver.go
plugins/governance/main.go
plugins/governance/tracker.go
plugins/governance/store.go

🧬 Code graph analysis (3)

transports/bifrost-http/handlers/middlewares.go (1)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

transports/bifrost-http/lib/config.go (1)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

framework/configstore/tables/budget.go (1)

framework/configstore/tables/utils.go (1)

ParseDuration (9-43)

🔇 Additional comments (35)

docs/features/governance/virtual-keys.mdx (1)

618-618: Error message documentation is accurate.

The documented error message "Budget exceeded: VK budget exceeded: 105.50 > 100.00 dollars" at line 618 matches the actual implementation in plugins/governance/resolver.go, which formats budget exceeded errors as fmt.Sprintf("Budget exceeded: %s", err.Error()). The removal of any redundant prefixes improves clarity without introducing breaking changes, as clients should rely on the error.type field rather than message text parsing.

plugins/telemetry/go.mod (1)

9-9: LGTM! Dependency updates look good.

The dependency bumps are consistent with the broader repository-wide update. The changes include minor/patch version updates to AWS SDK v2 components, ByteDance Sonic, go-openapi packages, OpenTelemetry, and related transitive dependencies. No breaking changes or security concerns are evident.

Also applies to: 17-18, 28-30, 38-39, 45-65, 75-76, 78-79, 88-89, 94-95, 98-113

plugins/governance/go.mod (1)

17-18: LGTM! Dependency updates support governance refactor.

The dependency version bumps align with the governance plugin refactor and follow the same pattern as other modules in the repository. All updates are minor/patch versions with no breaking changes indicated.

Also applies to: 28-30, 37-38, 44-64, 74-75, 77-78, 83-84, 89-91, 94-109

plugins/semanticcache/go.mod (1)

18-19: LGTM! Dependency updates are consistent.

The dependency version bumps follow the same pattern as other modules. The removal of the govalidator indirect dependency and the updates to AWS SDK v2, Sonic, OpenAPI, and OpenTelemetry packages are all minor/patch version changes with no breaking changes indicated.

Also applies to: 29-31, 38-39, 44-64, 74-75, 77-78, 83-84, 89-91, 94-109

plugins/logging/go.mod (1)

6-6: LGTM! Logging plugin dependencies updated.

The direct upgrade of ByteDance Sonic to v1.14.2 and related indirect dependency updates are consistent with the repository-wide dependency refresh. All changes are minor/patch versions.

Also applies to: 16-17, 27-29, 36-36, 42-62, 72-73, 75-76, 81-82, 87-89, 92-107

plugins/governance/store.go (8)

41-72: Well-documented interface contract.

The GovernanceStore interface has clear error semantics documentation that distinguishes between policy violations and infrastructure failures. This is crucial for callers to handle errors appropriately without retrying on policy denials.

96-149: GetGovernanceData efficiently aggregates in-memory state.

The implementation correctly iterates all sync.Maps and builds a snapshot of the current governance data. The nil checks and type assertions are defensive and appropriate for concurrent access patterns.

1052-1080: CreateVirtualKeyInMemory properly initializes all relationships.

The method correctly stores associated budgets, rate limits, and provider config relationships before storing the VK itself. This ensures all referenced entities are available when the VK is accessed.

1083-1194: UpdateVirtualKeyInMemory handles complex state transitions correctly.

The implementation properly:

Uses checkAndUpdateBudget and checkAndUpdateRateLimit to preserve usage unless limits change

Handles removals by deleting from memory when references are cleared

Updates both the entity's embedded reference AND the main sync.Map store

Falls back to CreateVirtualKeyInMemory if the VK doesn't exist

1197-1238: DeleteVirtualKeyInMemory correctly cascades deletions.

The method properly deletes associated budgets, rate limits, and provider config resources when removing a virtual key. Based on the learnings, the 1:1 relationship means this cascade is safe.

Based on learnings: budgets and rate limits maintain 1:1 relationships with their parent entities.

1241-1293: Team and Customer CRUD methods follow consistent patterns.

The CreateTeamInMemory, UpdateTeamInMemory, DeleteTeamInMemory, CreateCustomerInMemory, UpdateCustomerInMemory, and DeleteCustomerInMemory methods all:

Use the same clone-then-store pattern for safe concurrent access

Apply checkAndUpdateBudget for preserving usage state

Handle cascade deletions and reference nullification appropriately

Maintain consistency with VK CRUD patterns

Also applies to: 1296-1328, 1331-1343, 1345-1434

1439-1534: Reference update helpers maintain consistency across relationships.

The updateBudgetReferences and updateRateLimitReferences methods ensure that when a budget or rate limit is reset, all entities referencing it are updated to point to the new cloned instance. This maintains consistency across the in-memory graph.

1536-1616: Usage preservation logic handles limit changes correctly.

The checkAndUpdateBudget and checkAndUpdateRateLimit helpers implement the correct policy:

When max limits decrease below current usage + baseline, reset to 0

Otherwise preserve existing usage

Apply changes to reset duration and max limits

This prevents users from exceeding new lower limits while allowing usage to continue under increased limits.

framework/modelcatalog/main.go (2)

81-85: Well-documented optional callback type.

The ShouldSyncPricingFunc type is clearly documented as optional and explains its purpose. The signature is simple and appropriate for a conditional sync gate.

32-32: Clean integration of optional sync control.

The shouldSyncPricingFunc field is properly:

Added to the ModelCatalog struct

Accepted as a parameter in Init

Initialized in the struct initialization

Documented as optional (can be nil)

This enables conditional pricing sync behavior without breaking existing callers.

Also applies to: 88-108

framework/configstore/tables/budget.go (3)

25-26: Virtual field for runtime usage tracking.

The LastDBUsage field follows the established pattern for tracking last-known database values. The gorm:"-" and json:"-" tags correctly mark it as non-persistent and non-serialized.

33-46: Enhanced validation logic in BeforeSave.

The validation correctly:

Checks reset duration format and positivity

Ensures MaxLimit is not negative

Returns clear error messages

The validation will prevent invalid budget configurations from being saved.

48-52: AfterFind hook populates runtime field.

The AfterFind hook correctly initializes LastDBUsage from CurrentUsage after a database read. This enables tracking of the original database value while the in-memory value may be updated.

framework/configstore/tables/ratelimit.go (2)

33-36: Virtual fields for token and request usage tracking.

The LastDBTokenUsage and LastDBRequestUsage fields follow the same pattern as budget.go's LastDBUsage. The gorm:"-" and json:"-" tags correctly mark them as runtime-only fields.

83-88: AfterFind hook populates both virtual fields.

The AfterFind hook correctly initializes both LastDBTokenUsage and LastDBRequestUsage from their corresponding persisted usage fields. This enables tracking of original database values for both token and request limits.

core/go.mod (2)

6-22: Dependency version bumps look good.

These are minor/patch version bumps for AWS SDK v2 components, sonic, mcp-go, fasthttp, and golang.org/x packages. No breaking changes expected.

3-3: No action required. Go 1.25.5 was released on December 2, 2025 and is a valid stable release.

transports/bifrost-http/handlers/governance.go (5)

41-44: Good addition of nil validation for manager parameter.

This defensive check prevents runtime panics from nil manager usage. Consistent with the existing configStore nil check.

698-711: Good documentation of nil vs empty slice semantics for KeyIDs.

The comment at lines 698-701 clearly explains the tri-state behavior: nil (omitted) preserves existing keys, empty slice clears keys, values update keys. This prevents accidental key deletion.

567-569: Good addition of validation for existing rate limit updates.

Previously, validation was only applied when creating new rate limits. This ensures updated rate limits are also validated before persistence.

178-181: New budget and rate-limit GET endpoints registered.

These endpoints follow the established pattern and are properly chained with middlewares. Consistent with the existing governance API surface.

942-964: In-memory team filtering with customer_id is correctly implemented.

The filtering logic properly handles both filtered and unfiltered cases, maintaining consistent response structure with the database path.

transports/go.mod (3)

3-3: Same Go version concern as core/go.mod.

Go 1.25.5 appears here as well. Please verify this is intentional across all module files.

5-25: Direct dependency updates look good.

Addition of mcp-go v0.43.2 and bumps to sonic, fasthttp align with the broader PR changes.

59-79: go-openapi/swag subpackage additions reflect upstream restructuring.

The new subpackages (cmdutils, conv, fileutils, etc.) are part of the go-openapi/swag v0.25.x refactor that split the monolithic package into focused submodules.

plugins/governance/resolver.go (3)

64-75: Good refactor to use interface type for store field.

Changing from *GovernanceStore to GovernanceStore follows the Go idiom of accepting interfaces. This enables easier testing and alternative implementations.

219-233: Simplified checkBudgetHierarchy cleanly delegates to store.

The refactored method has clear responsibility: delegate to store.CheckBudget and convert errors to EvaluationResult. The logging provides useful debug context.

129-137: Hierarchical rate-limit and budget checks properly integrated.

The evaluation flow now correctly checks rate limits at the hierarchy level (provider → VK) before checking budgets (VK → Team → Customer). The method names clearly communicate the hierarchical nature of these checks.

transports/bifrost-http/lib/config.go (1)

47-50: EnterpriseOverrides wiring for pricing manager looks coherent

Threading EnterpriseOverrides through LoadConfig, loadConfigFromFile, loadConfigFromDefaults, initFrameworkConfigFromFile, and initDefaultFrameworkConfig is consistent and preserves the previous behavior for OSS (defaulting to modelcatalog.Init when no override is supplied). The selective field update strategy for startup rate‑limit resets also looks correct and avoids clobbering config fields.

Also applies to: 275-299, 357-408, 1505-1546, 1575-1629, 1866-1927

plugins/governance/main.go (1)

40-48: BaseGovernancePlugin and Init/InitFromStore wiring look solid

The introduction of BaseGovernancePlugin, the switch to a non‑pointer GovernanceStore field, and the split between Init (local store construction) and InitFromStore (custom store injection) form a clean lifecycle. Startup reset behavior gated on configStore and the updated Init documentation are consistent with the new design.

Also applies to: 71-163, 166-224, 673-675

transports/bifrost-http/handlers/governance.go

coderabbitai

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

transports/bifrost-http/server/server.go (1)
212-322: Fix governance plugin loading when name overrides are used

LoadPlugin dispatches the governance branch using:
case EnterpriseOverrides.GetGovernancePluginName():
    // ...
but BifrostHTTPServer.LoadGovernancePlugin always calls:
governancePlugin, err := LoadPlugin[*governance.GovernancePlugin](ctx, governance.PluginName, ...)
If OSSToEnterprisePluginNameOverrides maps "governance" to a different name (e.g. "enterprise-governance"), then:

GetGovernancePluginName() returns the overridden name.

The case EnterpriseOverrides.GetGovernancePluginName() in LoadPlugin expects the overridden name.

LoadGovernancePlugin still passes the literal "governance", so the governance case is never hit and the call falls through to "plugin governance not found".

This will break governance initialization in any configuration where overrides are in use.

A straightforward fix is to route LoadGovernancePlugin through GetGovernancePluginName() so the name used for dispatch matches the governance case:
Proposed fix
-func (s *BifrostHTTPServer) LoadGovernancePlugin(ctx context.Context, config *lib.Config) (schemas.Plugin, error) {
-	governancePlugin, err := LoadPlugin[*governance.GovernancePlugin](ctx, governance.PluginName, nil, &governance.Config{
-		IsVkMandatory: &config.ClientConfig.EnforceGovernanceHeader,
-	}, config, s)
+func (s *BifrostHTTPServer) LoadGovernancePlugin(ctx context.Context, config *lib.Config) (schemas.Plugin, error) {
+	name := s.GetGovernancePluginName()
+	governancePlugin, err := LoadPlugin[*governance.GovernancePlugin](ctx, name, nil, &governance.Config{
+		IsVkMandatory: &config.ClientConfig.EnforceGovernanceHeader,
+	}, config, s)
You may also want to use GetGovernancePluginName() in the case arm of LoadPlugin (already done) and in any status/logging paths (already using EnterpriseOverrides.GetGovernancePluginName() in LoadPlugins) to keep the naming consistent end-to-end.

Also applies to: 325-400, 1061-1087

♻️ Duplicate comments (2)

transports/bifrost-http/handlers/governance.go (1)

188-201: Nil check for GetGovernanceData() addresses previous review concern.

The implementation correctly returns a 500 error when governance data is unavailable, preventing nil pointer dereferences.

plugins/governance/tracker.go (1)

47-49: Consider making workerInterval and full dumps tunable for large deployments

Right now resetWorker runs every 10s and resetExpiredCounters unconditionally calls DumpRateLimits and DumpBudgets. On large installs with many VKs, this can create noticeable DB write load even when there are no meaningful changes.

Consider either:

Making workerInterval configurable (or at least bumping the default), and/or

Gating the Dump* calls behind a “dirty”/changed flag or running them less frequently than the reset scan.

This keeps the design but gives you a lever to dial back DB pressure when needed.

Also applies to: 107-147

🧹 Nitpick comments (5)

framework/configstore/tables/budget.go (1)
24-26: Consider documenting the purpose of LastDBUsage.

The runtime field LastDBUsage is used to track the database state after fetch, enabling dirty-state detection. Adding a brief comment explaining this pattern would improve maintainability.
📝 Suggested documentation
 	ConfigHash string `gorm:"type:varchar(255);null" json:"config_hash"`
 
 	CreatedAt time.Time `gorm:"index;not null" json:"created_at"`
 	UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"`
 
-	// Virtual fields for runtime use (not stored in DB)
+	// LastDBUsage captures the CurrentUsage value from the DB at load time,
+	// enabling in-memory change detection for governance tracking.
 	LastDBUsage float64 `gorm:"-" json:"-"`
transports/bifrost-http/lib/config_test.go (2)
194-196: Tighten MockConfigStore behavior for new methods (RetryOnNotFound/GetRateLimits).

The new methods correctly satisfy the expanded ConfigStore interface, but a slightly richer mock would better mirror real behavior:

RetryOnNotFound currently ignores maxRetries and retryDelay and just calls fn once. That’s fine for most unit tests, but if any caller under test asserts retry semantics (e.g., multiple reads on eventual consistency), this mock won’t catch regressions.

GetRateLimits always returns an empty slice, ignoring CreateRateLimit and governanceConfig.RateLimits. That can mask bugs in code paths that rely on previously created rate limits when using this mock.

Consider wiring GetRateLimits through m.governanceConfig so tests can observe created rate limits, and (optionally) tracking how many times RetryOnNotFound calls fn if you later need to assert retry behavior:
Possible refinement for GetRateLimits
 func (m *MockConfigStore) GetRateLimits(ctx context.Context) ([]tables.TableRateLimit, error) {
-	return []tables.TableRateLimit{}, nil
+	if m.governanceConfig == nil {
+		return nil, nil
+	}
+	return m.governanceConfig.RateLimits, nil
 }
Also applies to: 317-319

627-633: Optional: make testLogger surface errors/fatals to aid debugging.

The no-op testLogger keeps tests quiet, which is good, but makes diagnosing failing SQLite/governance tests harder. As an optional improvement, you could have Error/Fatal at least print to os.Stderr or testing.T.Logf (via a closure) so serious issues leave breadcrumbs in CI logs, while keeping Debug/Info muted.
plugins/governance/main.go (1)
40-47: Factor shared wiring between Init and InitFromStore into a single helper

The introduction of BaseGovernancePlugin and InitFromStore is solid, but Init and InitFromStore now duplicate the same wiring (isVkMandatory handling, resolver/tracker construction, startup resets, plugin struct initialization).

To avoid future divergence, consider an unexported helper along the lines of:
Possible internal helper sketch
func newGovernancePlugin(
    ctx context.Context,
    config *Config,
    logger schemas.Logger,
    governanceStore GovernanceStore,
    configStore configstore.ConfigStore,
    modelCatalog *modelcatalog.ModelCatalog,
    inMemoryStore InMemoryStore,
) (*GovernancePlugin, error) {
    if configStore == nil {
        logger.Warn("governance plugin requires config store to persist data, running in memory only mode")
    }
    if modelCatalog == nil {
        logger.Warn("governance plugin requires model catalog to calculate cost, all cost calculations will be skipped.")
    }

    var isVkMandatory *bool
    if config != nil {
        isVkMandatory = config.IsVkMandatory
    }

    resolver := NewBudgetResolver(governanceStore, logger)
    tracker := NewUsageTracker(ctx, governanceStore, resolver, configStore, logger)

    if configStore != nil {
        if err := tracker.PerformStartupResets(ctx); err != nil {
            logger.Warn("startup reset failed: %v", err)
        }
    }

    pluginCtx, cancel := context.WithCancel(ctx)
    return &GovernancePlugin{
        ctx:           pluginCtx,
        cancelFunc:    cancel,
        store:         governanceStore,
        resolver:      resolver,
        tracker:       tracker,
        configStore:   configStore,
        modelCatalog:  modelCatalog,
        logger:        logger,
        inMemoryStore: inMemoryStore,
        isVkMandatory: isVkMandatory,
    }, nil
}
Init would just build a LocalGovernanceStore and call this helper; InitFromStore would pass the injected store, keeping all behavior in one place.

Also applies to: 71-118, 166-224, 673-675
transports/bifrost-http/server/server.go (1)
499-517: Reuse getGovernancePlugin inside GetGovernanceData for consistency

You’ve added getGovernancePlugin() which correctly:

Protects s.Plugins with PluginsMutex.RLock.

Uses GetGovernancePluginName() for name resolution.

Asserts the plugin implements BaseGovernancePlugin.

GetGovernanceData reimplements a slightly different lookup (also under PluginsMutex.RLock) and repeats the BaseGovernancePlugin assertion. To avoid drift and keep the locking/type-assertion logic in one place, consider delegating:
Proposed refactor
-func (s *BifrostHTTPServer) GetGovernanceData() *governance.GovernanceData {
-	s.PluginsMutex.RLock()
-	governancePlugin, err := FindPluginByName[schemas.Plugin](s.Plugins, s.GetGovernancePluginName())
-	s.PluginsMutex.RUnlock()
-	if err != nil {
-		return nil
-	}
-	// Check if GetGovernanceStore method is implemented
-	if governancePlugin, ok := governancePlugin.(governance.BaseGovernancePlugin); ok {
-		return governancePlugin.GetGovernanceStore().GetGovernanceData()
-	}
-	return nil
-}
+func (s *BifrostHTTPServer) GetGovernanceData() *governance.GovernanceData {
+	governancePlugin, err := s.getGovernancePlugin()
+	if err != nil {
+		return nil
+	}
+	return governancePlugin.GetGovernanceStore().GetGovernanceData()
+}
This keeps all governance-plugin lookup rules centralized and easier to evolve.

Also applies to: 648-661

📜 Review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 5ca488c and e3f7191.

⛔ Files ignored due to path filters (11)

core/go.sum is excluded by !**/*.sum
framework/go.sum is excluded by !**/*.sum
plugins/governance/go.sum is excluded by !**/*.sum
plugins/jsonparser/go.sum is excluded by !**/*.sum
plugins/logging/go.sum is excluded by !**/*.sum
plugins/maxim/go.sum is excluded by !**/*.sum
plugins/mocker/go.sum is excluded by !**/*.sum
plugins/otel/go.sum is excluded by !**/*.sum
plugins/semanticcache/go.sum is excluded by !**/*.sum
plugins/telemetry/go.sum is excluded by !**/*.sum
transports/go.sum is excluded by !**/*.sum

📒 Files selected for processing (35)

.editorconfig
core/go.mod
core/mcp.go
core/providers/nebius/nebius.go
docs/changelogs/v1.3.47.mdx
docs/features/governance/virtual-keys.mdx
framework/configstore/rdb.go
framework/configstore/store.go
framework/configstore/tables/budget.go
framework/configstore/tables/mcp.go
framework/configstore/tables/ratelimit.go
framework/go.mod
framework/modelcatalog/main.go
framework/modelcatalog/sync.go
plugins/governance/go.mod
plugins/governance/main.go
plugins/governance/resolver.go
plugins/governance/store.go
plugins/governance/tracker.go
plugins/governance/utils.go
plugins/jsonparser/go.mod
plugins/logging/go.mod
plugins/maxim/go.mod
plugins/mocker/go.mod
plugins/otel/go.mod
plugins/semanticcache/go.mod
plugins/telemetry/go.mod
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/lib/lib.go
transports/bifrost-http/server/server.go
transports/go.mod
ui/lib/constants/logs.ts

💤 Files with no reviewable changes (1)

framework/configstore/tables/mcp.go

🚧 Files skipped from review as they are similar to previous changes (9)

framework/modelcatalog/sync.go
docs/changelogs/v1.3.47.mdx
framework/configstore/tables/ratelimit.go
transports/bifrost-http/lib/lib.go
.editorconfig
plugins/mocker/go.mod
docs/features/governance/virtual-keys.mdx
framework/modelcatalog/main.go
ui/lib/constants/logs.ts

🧰 Additional context used

📓 Path-based instructions (1)

**

⚙️ CodeRabbit configuration file

always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)

Files:

framework/configstore/store.go
transports/bifrost-http/handlers/middlewares.go
core/mcp.go
core/providers/nebius/nebius.go
transports/bifrost-http/handlers/governance.go
framework/go.mod
plugins/governance/tracker.go
plugins/otel/go.mod
plugins/governance/utils.go
core/go.mod
framework/configstore/rdb.go
plugins/telemetry/go.mod
plugins/governance/resolver.go
transports/bifrost-http/server/server.go
plugins/governance/main.go
plugins/semanticcache/go.mod
transports/bifrost-http/lib/config.go
transports/go.mod
plugins/governance/go.mod
plugins/governance/store.go
plugins/maxim/go.mod
plugins/logging/go.mod
transports/bifrost-http/lib/config_test.go
plugins/jsonparser/go.mod
framework/configstore/tables/budget.go

🧠 Learnings (4)

📚 Learning: 2025-12-09T17:07:42.007Z

Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.

Applied to files:

framework/configstore/store.go
transports/bifrost-http/handlers/middlewares.go
core/mcp.go
core/providers/nebius/nebius.go
transports/bifrost-http/handlers/governance.go
plugins/governance/tracker.go
plugins/governance/utils.go
framework/configstore/rdb.go
plugins/governance/resolver.go
transports/bifrost-http/server/server.go
plugins/governance/main.go
transports/bifrost-http/lib/config.go
plugins/governance/store.go
transports/bifrost-http/lib/config_test.go
framework/configstore/tables/budget.go

📚 Learning: 2025-12-12T08:25:02.629Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.

Applied to files:

transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/server/server.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go

📚 Learning: 2025-12-19T09:26:54.961Z

Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/providers/utils/utils.go:1050-1051
Timestamp: 2025-12-19T09:26:54.961Z
Learning: Update streaming end-marker handling so HuggingFace is treated as a non-[DONE] provider for backends that do not emit a DONE marker (e.g., meta llama on novita). In core/providers/utils/utils.go, adjust ProviderSendsDoneMarker() (or related logic) to detect providers that may not emit DONE and avoid relying on DONE as the sole end signal. Add tests to cover both DONE-emitting and non-DONE backends, with clear documentation in code comments explaining the rationale and any fallback behavior.

Applied to files:

core/providers/nebius/nebius.go

📚 Learning: 2025-12-22T10:50:40.990Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.

Applied to files:

plugins/governance/tracker.go
plugins/governance/utils.go
plugins/governance/resolver.go
plugins/governance/main.go
plugins/governance/store.go

🧬 Code graph analysis (10)

framework/configstore/store.go (1)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

transports/bifrost-http/handlers/middlewares.go (2)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

transports/bifrost-http/lib/middleware.go (1)

BifrostHTTPMiddleware (7-7)

plugins/governance/tracker.go (3)

plugins/governance/store.go (1)

GovernanceStore (50-72)

plugins/governance/resolver.go (1)

BudgetResolver (64-67)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

framework/configstore/rdb.go (5)

framework/configstore/tables/virtualkey.go (6)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

TableVirtualKeyMCPConfig (81-87)

TableVirtualKeyMCPConfig (90-92)

framework/configstore/tables/budget.go (2)

TableBudget (11-27)

TableBudget (30-30)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

framework/configstore/tables/team.go (2)

TableTeam (12-38)

TableTeam (41-41)

framework/configstore/tables/customer.go (2)

TableCustomer (6-22)

TableCustomer (25-25)

plugins/governance/resolver.go (3)

plugins/governance/store.go (1)

GovernanceStore (50-72)

framework/configstore/tables/virtualkey.go (2)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

core/schemas/bifrost.go (1)

ModelProvider (32-32)

transports/bifrost-http/server/server.go (3)

plugins/governance/store.go (1)

GovernanceData (33-39)

plugins/governance/main.go (4)

Config (32-34)

PluginName (21-21)

BaseGovernancePlugin (40-47)

GovernancePlugin (50-69)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

plugins/governance/main.go (2)

plugins/governance/store.go (2)

GovernanceStore (50-72)

NewLocalGovernanceStore (75-94)

plugins/governance/resolver.go (2)

BudgetResolver (64-67)

NewBudgetResolver (70-75)

transports/bifrost-http/lib/config.go (1)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

transports/bifrost-http/lib/config_test.go (1)

transports/bifrost-http/lib/config.go (1)

LoadConfig (275-355)

framework/configstore/tables/budget.go (1)

framework/configstore/tables/utils.go (1)

ParseDuration (9-43)

🔇 Additional comments (34)

core/providers/nebius/nebius.go (1)

52-52: LGTM! Field initialization is correct and consistent.

The sendBackRawRequest field is now properly initialized from the config, matching the pattern already established for sendBackRawResponse on line 53. This ensures the configuration value is threaded through to the provider instance and available for all method calls.

plugins/otel/go.mod (3)

8-9: Acknowledge broad dependency harmonization as part of the stack refactor.

The updates to gRPC (v1.77.0), protobuf (v1.36.11), golang.org/x modules (crypto, net, oauth2, sync, text, arch, sys), and genproto packages reflect a wider pattern of transitive dependency alignment across the repository (as noted in the AI summary). Ensure these coordinated bumps are tested together in the context of the full governance plugin refactor stack to avoid subtle compatibility issues or runtime deadlocks.

Since this PR is part of a Graphite stack targeting v1.4.0, confirm that integration tests for the entire stack pass before merging.

Also applies to: 96-102

92-94: The OTel v1.39.0 update is stable and backward compatible.

Most OpenTelemetry minor version changes are backward-compatible. The bump from v1.38.x to v1.39.0 follows semantic versioning and introduces no breaking changes. Since these are marked as indirect dependencies, the main concern is whether any direct plugin code uses affected APIs—verify the plugin implementation if it directly imports OTel trace/metric packages.

41-61: go-openapi package expansion is expected behavior from v0.25.x refactoring

The expansion of go-openapi subpackages (cmdutils, conv, fileutils, jsonname, jsonutils, loading, mangling, netutils, stringutils, typeutils, yamlutils) at v0.25.4 is part of an intentional mono-repo refactoring in go-openapi/swag, not a dependency bloat issue. Each subpackage has independent versioning and testing. Current security scans show no known vulnerabilities in these packages.

plugins/telemetry/go.mod (1)

1-118: Ensure this PR's dependency updates pass CI before merging.

The telemetry module's dependency version bumps (40+ updates across fasthttp, AWS SDK, bytedance/sonic, go-openapi modules, OpenTelemetry, MongoDB driver, qdrant, redis, weaviate, and golang.org packages) are already aligned with framework and governance modules. Key dependencies are synchronized: OpenTelemetry (v1.39.0), gRPC (v1.77.0), protobuf (v1.36.11), and bytedance/sonic (v1.14.2) match across telemetry, framework, and governance go.mod files. Core doesn't directly depend on these transitive packages, so no alignment issues there.

The only remaining validation needed is CI confirmation that the full build and integration tests pass with these versions across the entire module stack.

plugins/semanticcache/go.mod (2)

38-39: Dependency versions are coordinated across the stack; verify no breaking changes in service definitions.

The dependency updates in plugins/semanticcache/go.mod are consistent with the core, framework, and plugin modules:

Sonic v1.14.2 and loader v0.4.0 are synchronized across core, framework, and plugins/otel/go.mod

Redis v9.17.2 matches framework/vectorstore usage

OpenTelemetry v1.39.0 (otel, metric, trace) is uniform across all modules (core, framework, transports, and all plugins)

gRPC v1.77.0 and Protobuf v1.36.11 are pinned in plugins/otel/go.mod for service definitions

golang.org/x updates (crypto v0.46.0, net v0.48.0, etc.) are consistent across the stack

No version misalignment detected. However, confirm that gRPC v1.77.0 and Protobuf v1.36.11 do not introduce breaking changes in the otel plugin's service definitions.

44-64: The go-openapi dependencies in lines 44–64 are transitive (indirect) and not directly used by the semanticcache plugin, which focuses on semantic caching and vectorstore operations. No spec parsing, validation, or code generation logic exists in this plugin that would be affected by go-openapi version updates. The referenced asaskevich/govalidator package does not appear in the codebase.

Likely an incorrect or invalid review comment.

plugins/logging/go.mod (4)

5-9: Verify direct dependency versions are compatible with indirect dependency updates.

The direct dependencies (bifrost/core v1.2.40, bifrost/framework v1.1.50) remain unchanged while many indirect dependencies have been updated. Ensure these direct dependencies are compatible with the new versions of transitive dependencies, especially the major opensearch/OpenTelemetry and gRPC upgrades.

Since this is part of a coordinated stack across PRs (#1121, #1120, #1119, #1041, #1154), verify that all modules in the stack have compatible dependency versions. You may want to check the go.mod files of bifrost/core and bifrost/framework to confirm they've been updated in parallel.

3-3: No action needed. Go 1.25.5 is a valid, stable release released December 2, 2025, and is current for this context dated December 22, 2025.

27-27: The S3 SDK is an indirect dependency for this module and v1.94.0 introduces only backward-compatible changes (new optional field for LifecycleExpirationDate). No verification is required.

42-62: The go-openapi packages listed are all indirect dependencies (transitive) and are not directly imported or used in the codebase. No breaking API calls to go-openapi modules were found in the governance plugin refactor or elsewhere. This concern can be safely removed as breaking changes in transitive dependencies pose minimal risk when there are no direct usages.

plugins/governance/utils.go (1)

18-25: LGTM: Clean pointer equality helper.

The generic equalPtr function correctly handles nil-safe pointer comparison with appropriate use of the comparable constraint.

framework/configstore/tables/budget.go (2)

40-43: LGTM: Appropriate validation for budget limits.

The validation ensures that MaxLimit is non-negative, which is a sensible constraint for budget configuration.

48-52: LGTM: AfterFind hook correctly populates runtime state.

The hook captures the database state into LastDBUsage after loading, supporting change detection as part of the broader in-memory governance refactor.

core/mcp.go (1)

1140-1140: Remove unnecessary type assertion by making InProcessServer type-safe.

The code change correctly aligns with the library API, but the reasoning was inverted. Rather than the library accepting a "more general interface type," the change actually makes the code more type-safe by specifying InProcessServer as *server.MCPServer in the schema instead of interface{}. This eliminates the runtime type assertion that was previously required and improves compile-time type safety. The library's NewInProcessClient has always required *server.MCPServer, so this change simply makes the caller's code more explicit about the expected type.

Likely an incorrect or invalid review comment.

plugins/maxim/go.mod (2)

14-114: Dependency version bumps look consistent.

The indirect dependency updates across AWS SDK, OpenAPI tooling, Redis, MongoDB, OpenTelemetry, and other modules appear to be routine maintenance updates aligned with the broader PR's dependency refresh.

3-3: No issues found. Go 1.25.5 was released on December 2, 2025 and is a valid Go version.

framework/configstore/store.go (2)

104-104: Interface addition for GetRateLimits is consistent with existing patterns.

The new method follows the same pattern as GetBudgets and other retrieval methods in the interface.

146-148: RetryOnNotFound interface addition looks good.

The method signature is well-designed with configurable retry parameters. This provides a clean abstraction for retryable operations across the config store implementations.

transports/bifrost-http/handlers/middlewares.go (1)

49-73: Nil check for enterpriseOverrides properly addresses previous review concern.

The implementation correctly handles the nil case by bypassing interception when enterpriseOverrides is not provided. The governance plugin name is now dynamically retrieved from the enterprise overrides, improving extensibility.

framework/configstore/rdb.go (4)

1399-1464: Transactional delete for DeleteVirtualKey is well-structured.

The cascading cleanup of provider configs, join table entries, budgets, rate limits, and MCP configs within a transaction ensures data consistency. The preload of ProviderConfigs ensures all associated IDs are available before cleanup.

One minor observation: the raw SQL at line 1412 uses a hardcoded table name. Consider using GORM's association methods for consistency, though this approach is acceptable for performance.

1665-1699: Transactional delete for DeleteTeam correctly nullifies foreign keys.

The approach of nullifying team_id on associated virtual keys before deleting the team prevents FK constraint violations while preserving the virtual keys. The budget cleanup is also handled correctly.

1754-1792: Transactional delete for DeleteCustomer follows the same pattern as DeleteTeam.

Consistently nullifies customer_id on virtual keys and teams, then cleans up the budget before deleting the customer.

2154-2179: RetryOnNotFound implementation is clean and well-designed.

The implementation correctly:

Respects context cancellation between retries

Uses Go 1.22+ range-over-int syntax

Avoids sleeping after the final attempt

Returns the last error when all retries are exhausted

transports/bifrost-http/handlers/governance.go (6)

24-26: Interface addition for GetGovernanceData() aligns with in-memory governance architecture.

The return type *governance.GovernanceData provides access to the in-memory governance state for efficient lookups.

42-44: Constructor nil check for manager is appropriate.

Returning an error when the manager is nil prevents potential nil pointer dereferences later in the handler methods.

310-329: Key validation in createVirtualKey ensures data integrity.

The code validates that all requested key IDs exist and returns a clear error message if some keys are not found. This prevents creating provider configs with invalid key references.

635-654: Key handling for new provider configs in updateVirtualKey follows the same pattern.

Consistent with the create path, ensuring key validation for new provider configurations during updates.

706-719: Conditional key update preserves existing keys when KeyIDs is omitted.

The nil check on pc.KeyIDs correctly distinguishes between:

Field omitted (nil): preserve existing keys

Field present but empty ([]): clear keys

Field with values: update keys

This is a good design choice for partial updates.

1445-1501: New getBudgets and getRateLimits handlers follow established patterns.

Both handlers:

Support the from_memory query parameter

Include nil checks for governance data

Fall back to database queries when not using in-memory mode

Return consistent JSON response structure with count
transports/bifrost-http/lib/config_test.go (1)
1832-1841: Key comparison change correctly avoids duplicating modified keys.

The merge logic in TestProviderHashComparison_DifferentHash now treats DB and file keys as the “same logical key” when either the content hash matches or the names match:
if dbKeyHash == fileKeyHash || fileKey.Name == dbKey.Name {
    found = true
}
This prevents a file-updated key with the same name but different content from being misclassified as “DB-only” and appended as a second, duplicate key. That aligns the test with the intended semantics of “dashboard-added vs file-owned” keys and looks like the right fix.
plugins/governance/resolver.go (1)

194-233: Hierarchical rate-limit and budget wrappers look consistent with store contract

The new checkRateLimitHierarchy / checkBudgetHierarchy helpers correctly delegate to GovernanceStore and translate non-nil errors into explicit EvaluationResults, while preserving Decision from CheckRateLimit. This matches the documented “error == policy violation, not infra failure” contract and keeps EvaluateRequest clean.

transports/bifrost-http/lib/config.go (1)

47-50: EnterpriseOverrides wiring for pricing manager is coherent

Threading EnterpriseOverrides through config loading and delegating pricing manager initialization to LoadPricingManager (with a sane fallback to modelcatalog.Init) is clean and keeps OSS vs enterprise concerns separated without changing existing behavior when overrides are nil.

Also applies to: 1505-1546, 1866-1927

plugins/governance/store.go (1)

1568-1616: No action required. The equalPtr generic function is properly defined in plugins/governance/utils.go (lines 20-27) within the same governance package and is correctly used in the checkAndUpdateRateLimit function. The code will compile successfully.

Likely an incorrect or invalid review comment.

framework/configstore/rdb.go

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (9)

framework/modelcatalog/sync.go (1)
57-58: Consider updating the function comment.

The function comment doesn't mention the conditional early exit behavior introduced by shouldSyncPricingFunc. Consider adding a brief note about this to improve discoverability.
📝 Suggested documentation improvement
-// syncPricing syncs pricing data from URL to database and updates cache
+// syncPricing syncs pricing data from URL to database and updates cache.
+// The sync can be cancelled early if shouldSyncPricingFunc returns false.
 func (mc *ModelCatalog) syncPricing(ctx context.Context) error {
plugins/governance/resolver.go (1)
263-306: Consider simplifying the nested rate limit violation logic.

The current implementation has empty if blocks (lines 269-270, 289-290) that exist only to allow the else branches to execute. While this works correctly, it reduces readability. Consider inverting the conditions.
🔎 Suggested simplification for token limit check
 // Check token limits
 if config.RateLimit.TokenMaxLimit != nil && config.RateLimit.TokenCurrentUsage >= *config.RateLimit.TokenMaxLimit {
     // Check if token limit needs reset
     if config.RateLimit.TokenResetDuration != nil {
         if duration, err := configstoreTables.ParseDuration(*config.RateLimit.TokenResetDuration); err == nil {
-            if time.Since(config.RateLimit.TokenLastReset).Round(time.Millisecond) >= duration {
-                // Token limit expired but hasn't been reset yet - not violated
-            } else {
-                // Token limit exceeded and not expired
+            if time.Since(config.RateLimit.TokenLastReset).Round(time.Millisecond) < duration {
+                // Token limit exceeded and not expired - violated
                 return true
             }
+            // Token limit expired but hasn't been reset yet - not violated
         } else {
             // Parse error - assume violated
             return true
         }
     } else {
         // No reset duration - violated
         return true
     }
 }
plugins/governance/tracker.go (1)
47-49: Consider making workerInterval configurable or increasing default value for scalability.

The 10-second interval for dumping rate limits and budgets to the database may create noticeable load at scale. Consider either making this value configurable (e.g., via a package-level variable or config parameter) or increasing the default to 30-60 seconds for large deployments.
🔎 Option 1: Make configurable via constructor parameter

Add a field to UsageTracker:
 type UsageTracker struct {
 	store       GovernanceStore
 	resolver    *BudgetResolver
 	configStore configstore.ConfigStore
 	logger      schemas.Logger
+	workerInterval time.Duration

 	// Background workers
 	trackerCtx    context.Context
Update constructor:
-func NewUsageTracker(ctx context.Context, store GovernanceStore, resolver *BudgetResolver, configStore configstore.ConfigStore, logger schemas.Logger) *UsageTracker {
+func NewUsageTracker(ctx context.Context, store GovernanceStore, resolver *BudgetResolver, configStore configstore.ConfigStore, logger schemas.Logger, workerInterval time.Duration) *UsageTracker {
+	if workerInterval <= 0 {
+		workerInterval = 30 * time.Second // safer default
+	}
 	tracker := &UsageTracker{
 		store:       store,
 		resolver:    resolver,
 		configStore: configStore,
 		logger:      logger,
+		workerInterval: workerInterval,
 		done:        make(chan struct{}),
 	}
Use instance field instead of constant:
 func (t *UsageTracker) startWorkers(ctx context.Context) {
 	// Counter reset manager (business logic)
-	t.resetTicker = time.NewTicker(workerInterval)
+	t.resetTicker = time.NewTicker(t.workerInterval)
 	t.wg.Add(1)
 	go t.resetWorker(ctx)
 }
🔎 Option 2: Increase default to reduce DB load
 const (
-	workerInterval = 10 * time.Second
+	workerInterval = 60 * time.Second  // Reduced DB write frequency for production scale
 )
plugins/governance/main.go (1)
110-164: Consider extracting common initialization logic to reduce duplication.

Init and InitFromStore share approximately 45 lines of nearly identical code (nil checks, resolver/tracker initialization, startup resets, plugin construction). This duplication makes maintenance error-prone—changes must be applied to both functions.
🔎 Proposed refactor to extract shared logic

Extract a common helper:
// initializePlugin performs common initialization logic for both Init and InitFromStore
func initializePlugin(
	ctx context.Context,
	config *Config,
	logger schemas.Logger,
	governanceStore GovernanceStore,
	configStore configstore.ConfigStore,
	modelCatalog *modelcatalog.ModelCatalog,
	inMemoryStore InMemoryStore,
) (*GovernancePlugin, error) {
	if configStore == nil {
		logger.Warn("governance plugin requires config store to persist data, running in memory only mode")
	}
	if modelCatalog == nil {
		logger.Warn("governance plugin requires model catalog to calculate cost, all cost calculations will be skipped.")
	}

	// Handle nil config - use safe default for IsVkMandatory
	var isVkMandatory *bool
	if config != nil {
		isVkMandatory = config.IsVkMandatory
	}

	// Initialize components in dependency order
	resolver := NewBudgetResolver(governanceStore, logger)
	tracker := NewUsageTracker(ctx, governanceStore, resolver, configStore, logger)

	// Perform startup reset check for any expired limits from downtime
	if configStore != nil {
		if err := tracker.PerformStartupResets(ctx); err != nil {
			logger.Warn("startup reset failed: %v", err)
			// Continue initialization even if startup reset fails (non-critical)
		}
	}

	ctx, cancelFunc := context.WithCancel(ctx)
	plugin := &GovernancePlugin{
		ctx:           ctx,
		cancelFunc:    cancelFunc,
		store:         governanceStore,
		resolver:      resolver,
		tracker:       tracker,
		configStore:   configStore,
		modelCatalog:  modelCatalog,
		logger:        logger,
		isVkMandatory: isVkMandatory,
		inMemoryStore: inMemoryStore,
	}
	return plugin, nil
}
Then simplify both constructors:
func Init(
	ctx context.Context,
	config *Config,
	logger schemas.Logger,
	configStore configstore.ConfigStore,
	governanceConfig *configstore.GovernanceConfig,
	modelCatalog *modelcatalog.ModelCatalog,
	inMemoryStore InMemoryStore,
) (*GovernancePlugin, error) {
	governanceStore, err := NewLocalGovernanceStore(ctx, logger, configStore, governanceConfig)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize governance store: %w", err)
	}
	return initializePlugin(ctx, config, logger, governanceStore, configStore, modelCatalog, inMemoryStore)
}

func InitFromStore(
	ctx context.Context,
	config *Config,
	logger schemas.Logger,
	governanceStore GovernanceStore,
	configStore configstore.ConfigStore,
	modelCatalog *modelcatalog.ModelCatalog,
	inMemoryStore InMemoryStore,
) (*GovernancePlugin, error) {
	if governanceStore == nil {
		return nil, fmt.Errorf("governance store is nil")
	}
	return initializePlugin(ctx, config, logger, governanceStore, configStore, modelCatalog, inMemoryStore)
}
Also applies to: 178-224
transports/bifrost-http/lib/config_test.go (4)

194-196: Mock RetryOnNotFound ignores retries and delay parameters

RetryOnNotFound just invokes fn once and ignores maxRetries / retryDelay. That’s fine for current unit tests that don’t care about retry semantics, but it diverges from production behavior and may hide issues if tests start depending on retry logic.

Consider documenting this intentionally simplified behavior in a comment, or extending the mock later if you add tests around retry/backoff.

317-319: Align GetRateLimits’ empty result semantics with other getters

This mock returns an empty slice ([]tables.TableRateLimit{}) while other getters like GetBudgets, GetCustomers, and GetTeams return nil, nil when there’s no data. The difference is usually harmless, but it’s slightly inconsistent and could matter if callers distinguish nil vs non‑nil slices.

Consider returning nil, nil here as well for consistency with the rest of the mock store.

627-633: testLogger Fatal is a no‑op

Adding Fatal and SetLevel makes testLogger conform to the expanded logger interface, and keeping Fatal as a no‑op is reasonable for tests so they don’t call os.Exit.

If you ever add tests that depend on fatal conditions aborting work, you may want Fatal to record that it was called (e.g., flip a flag) so assertions can verify it.

6675-6676: LoadConfig callsites now pass nil for EnterpriseOverrides

Updating tests to call LoadConfig(ctx, tempDir, nil) matches the new signature and keeps them on the default behavior path (no overrides applied).

At some point it may be worth adding one or two tests that pass a non‑nil EnterpriseOverrides to exercise the new override plumbing end‑to‑end, but the current change is correct as-is.

plugins/governance/store.go (1)

1088-1193: Note: Zero baselines for VK updates.

The update logic uses baseline = 0.0 for budgets (line 1107) and baseline = 0, 0 for rate limits (line 1125) when calling the checkAndUpdate* helpers. This means remote node usage data is not considered during VK config updates—only local in-memory state is preserved or reset. The comments note "can be made configurable later."

If VK updates should account for distributed usage (e.g., sum across nodes before deciding to reset), consider making baselines configurable or fetching them from gossip aggregation.

📜 Review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e3f7191 and 7ae83e4.

⛔ Files ignored due to path filters (11)

core/go.sum is excluded by !**/*.sum
framework/go.sum is excluded by !**/*.sum
plugins/governance/go.sum is excluded by !**/*.sum
plugins/jsonparser/go.sum is excluded by !**/*.sum
plugins/logging/go.sum is excluded by !**/*.sum
plugins/maxim/go.sum is excluded by !**/*.sum
plugins/mocker/go.sum is excluded by !**/*.sum
plugins/otel/go.sum is excluded by !**/*.sum
plugins/semanticcache/go.sum is excluded by !**/*.sum
plugins/telemetry/go.sum is excluded by !**/*.sum
transports/go.sum is excluded by !**/*.sum

📒 Files selected for processing (35)

.editorconfig
core/go.mod
core/mcp.go
core/providers/nebius/nebius.go
docs/changelogs/v1.3.47.mdx
docs/features/governance/virtual-keys.mdx
framework/configstore/rdb.go
framework/configstore/store.go
framework/configstore/tables/budget.go
framework/configstore/tables/mcp.go
framework/configstore/tables/ratelimit.go
framework/go.mod
framework/modelcatalog/main.go
framework/modelcatalog/sync.go
plugins/governance/go.mod
plugins/governance/main.go
plugins/governance/resolver.go
plugins/governance/store.go
plugins/governance/tracker.go
plugins/governance/utils.go
plugins/jsonparser/go.mod
plugins/logging/go.mod
plugins/maxim/go.mod
plugins/mocker/go.mod
plugins/otel/go.mod
plugins/semanticcache/go.mod
plugins/telemetry/go.mod
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/lib/lib.go
transports/bifrost-http/server/server.go
transports/go.mod
ui/lib/constants/logs.ts

💤 Files with no reviewable changes (1)

framework/configstore/tables/mcp.go

✅ Files skipped from review due to trivial changes (2)

.editorconfig
docs/changelogs/v1.3.47.mdx

🚧 Files skipped from review as they are similar to previous changes (10)

ui/lib/constants/logs.ts
framework/modelcatalog/main.go
core/mcp.go
docs/features/governance/virtual-keys.mdx
core/providers/nebius/nebius.go
framework/go.mod
plugins/mocker/go.mod
plugins/governance/utils.go
framework/configstore/store.go
plugins/semanticcache/go.mod

🧰 Additional context used

📓 Path-based instructions (1)

**

⚙️ CodeRabbit configuration file

always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)

Files:

framework/configstore/tables/ratelimit.go
transports/bifrost-http/handlers/middlewares.go
framework/configstore/tables/budget.go
plugins/logging/go.mod
plugins/telemetry/go.mod
transports/bifrost-http/handlers/governance.go
plugins/jsonparser/go.mod
framework/configstore/rdb.go
plugins/governance/tracker.go
plugins/governance/go.mod
transports/bifrost-http/server/server.go
plugins/governance/main.go
core/go.mod
transports/bifrost-http/lib/lib.go
plugins/governance/resolver.go
transports/go.mod
plugins/otel/go.mod
plugins/maxim/go.mod
transports/bifrost-http/lib/config.go
framework/modelcatalog/sync.go
transports/bifrost-http/lib/config_test.go
plugins/governance/store.go

🧠 Learnings (3)

📚 Learning: 2025-12-09T17:07:42.007Z

Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.

Applied to files:

framework/configstore/tables/ratelimit.go
transports/bifrost-http/handlers/middlewares.go
framework/configstore/tables/budget.go
transports/bifrost-http/handlers/governance.go
framework/configstore/rdb.go
plugins/governance/tracker.go
transports/bifrost-http/server/server.go
plugins/governance/main.go
transports/bifrost-http/lib/lib.go
plugins/governance/resolver.go
transports/bifrost-http/lib/config.go
framework/modelcatalog/sync.go
transports/bifrost-http/lib/config_test.go
plugins/governance/store.go

📚 Learning: 2025-12-12T08:25:02.629Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.

Applied to files:

transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/server/server.go
transports/bifrost-http/lib/lib.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go

📚 Learning: 2025-12-22T10:50:40.990Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.

Applied to files:

plugins/governance/tracker.go
plugins/governance/main.go
plugins/governance/resolver.go
plugins/governance/store.go

🧬 Code graph analysis (9)

transports/bifrost-http/handlers/middlewares.go (1)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

framework/configstore/tables/budget.go (1)

framework/configstore/tables/utils.go (1)

ParseDuration (9-43)

framework/configstore/rdb.go (5)

framework/configstore/tables/virtualkey.go (4)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

framework/configstore/tables/budget.go (2)

TableBudget (11-27)

TableBudget (30-30)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

framework/configstore/tables/team.go (2)

TableTeam (12-38)

TableTeam (41-41)

framework/configstore/tables/customer.go (2)

TableCustomer (6-22)

TableCustomer (25-25)

plugins/governance/tracker.go (2)

plugins/governance/store.go (1)

GovernanceStore (50-72)

plugins/governance/resolver.go (1)

BudgetResolver (64-67)

transports/bifrost-http/server/server.go (3)

plugins/governance/store.go (1)

GovernanceData (33-39)

plugins/governance/main.go (4)

Config (32-34)

PluginName (21-21)

BaseGovernancePlugin (40-47)

GovernancePlugin (50-69)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

plugins/governance/main.go (3)

plugins/governance/store.go (2)

GovernanceStore (50-72)

NewLocalGovernanceStore (75-94)

plugins/governance/resolver.go (2)

BudgetResolver (64-67)

NewBudgetResolver (70-75)

plugins/governance/tracker.go (2)

UsageTracker (33-45)

NewUsageTracker (52-66)

transports/bifrost-http/lib/config.go (1)

framework/modelcatalog/main.go (2)

ModelCatalog (23-46)

Init (88-138)

transports/bifrost-http/lib/config_test.go (4)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

core/schemas/logger.go (1)

LogLevel (6-6)

core/schemas/account.go (1)

Key (8-20)

transports/bifrost-http/lib/config.go (1)

LoadConfig (275-355)

plugins/governance/store.go (7)

framework/configstore/tables/virtualkey.go (4)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

framework/configstore/tables/budget.go (2)

TableBudget (11-27)

TableBudget (30-30)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

plugins/governance/resolver.go (1)

EvaluationRequest (30-35)

ui/lib/types/config.ts (1)

ModelProvider (198-201)

framework/configstore/clientconfig.go (1)

GovernanceConfig (720-727)

ui/lib/types/governance.ts (4)

Budget (5-11)

RateLimit (13-25)

Team (27-35)

Customer (37-44)

🔇 Additional comments (49)

framework/modelcatalog/sync.go (1)

61-66: Early exit logic is correct and thread-safe.

The predicate-based cancellation logic is clean and efficiently placed before expensive operations. The nil-return on cancellation correctly treats it as normal flow rather than an error.

mc.shouldSyncPricingFunc is accessed without synchronization because it is immutable after initialization. The field is set once during ModelCatalog.Init() and never modified thereafter, making concurrent reads safe.

transports/go.mod (10)

1-25: Consider verifying dependency versions against the broader PR stack context.

This PR is part of a stack (#1121, #1120, #1119, #1041, #1154). Ensure that:

Dependency versions in transports/go.mod are consistent with versions used in related PR modules (core, framework, plugins)

The bifrost internal module versions (lines 12–19) match corresponding versions in those PRs

No dependency conflicts or version skew exists across the stack

As per coding guidelines, review changes in context of the entire stack.

Are the dependency versions in this transports/go.mod consistent with the versions declared in related stack PRs (#1121, #1120, #1119, #1041)? Specifically:

github.com/maximhq/bifrost/core v1.2.40

github.com/maximhq/bifrost/framework v1.1.50

github.com/maximhq/bifrost/plugins/governance v1.3.51 (and other plugins)

1-25: Verify dependency consistency across the PR stack.

This PR is part of a larger stack (#1121, #1120, #1119, #1041, #1154). Given the coding guidelines requirement to review changes in context of the entire stack, confirm:

The bifrost internal module versions (lines 12–19) are consistent across all stack PRs

No dependency version conflicts or skew exist between transports/go.mod and related module dependencies in core, framework, and plugin PRs

The mcp-go dependency (line 11) and openapi updates (lines 59–79) are used consistently across the stack

Have the dependency versions declared in transports/go.mod (particularly bifrost/core v1.2.40, bifrost/framework v1.1.50, bifrost/plugins/governance v1.3.51, and new mcp-go v0.43.2) been validated as consistent with the corresponding versions in related stack PRs #1121, #1120, #1119, #1041?

123-133: Verify compatibility with google.golang.org/grpc v1.77.0 API changes.

Lines 123–133 include updates to golang.org/x/* (arch, crypto, net, oauth2, sync, sys, text) and google.golang.org/* (genproto, grpc, protobuf) packages. The golang.org/x/crypto and golang.org/x/net updates are routine with no documented breaking changes. However, google.golang.org/grpc v1.77.0 includes API changes: the mem.Reader interface was replaced with a struct, and old pick_first balancer policy support was removed.

Confirm:

No code in transports or internal bifrost modules directly uses mem.Reader interface or legacy GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST environment variable

gRPC v1.77.0 changes are compatible with existing load balancer and transport layer usage

7-7: Minor dependency updates are safe; no breaking changes detected.

Sonic v1.14.2 includes fixes for validation behavior, EscapeHTML pooling, and assembly improvements, while fasthttp v1.68.0 fixes named return bugs and updates dependencies. Both updates are routine patch/minor version bumps with no breaking changes or compatibility concerns.

11-11: Verify the exact mcp-go version in line 11 of transports/go.mod.

Public release history for mark3labs/mcp-go does not show v0.43.2. Confirm whether this is the intended version or if a typo exists (e.g., v0.32.x). If v0.43.2 is intentional, document why this unreleased or pre-release version is required for Files API support.

123-133: No major breaking changes or TLS/cryptographic incompatibilities identified in these versions.

The latest golang.org/x/net available is v0.46.0 (published Nov 11, 2025), not v0.48.0 as referenced. golang.org/x/crypto v0.46.0 was published Dec 8, 2025 with no documented breaking changes affecting core cryptographic APIs or TLS implementations. google.golang.org/grpc v1.77.0 removes support for the old pick_first load balancer policy environment variable; the new pick_first has been the default since v1.71.0, representing experimental feature removal rather than breaking API changes.

Verify the correct version constraint for golang.org/x/net in the go.mod file (v0.46.0 vs. v0.48.0).

105-106: Address Weaviate v5.6.0 breaking change in vector type definition.

Weaviate v5.6.0 carries the breaking change where models.Vector changed from []float32 to interface{}. Verify that governance logic handling vector data—particularly EnterpriseOverrides and data persistence patterns for policies—is compatible with this interface{} vector type. For redis/go-redis v9.17.2 and qdrant/go-client v1.16.2, no critical API breaking changes were identified at the patch level. Confirm the governance implementation does not rely on vector handling code that expects the legacy []float32 slice type.

11-11: Verify the new MCP dependency version and its role in Files API support.

The addition of github.com/mark3labs/mcp-go v0.43.2 requires verification before merging. This version does not appear in publicly available release channels. Please confirm:

This version exists and is accessible to the build system

It aligns with the governance refactoring and Files API objectives

It is compatible with bifrost's internal modules (core v1.2.40, framework v1.1.50, governance v1.3.51)

59-79: Verify go-openapi v0.25.x updates for compatibility and necessity.

Lines 59–79 update the go-openapi module ecosystem to v0.25.x versions. While the traditional swag API is still supported and code that imports swag will still compile and work the same, specific breaking change documentation for the other modules (analysis, errors, jsonpointer, jsonreference, loads, runtime, spec, strfmt, validate) is not publicly available through standard search.

Confirm:

These specific versions work harmoniously together without introducing runtime or code generation breaks

These updates are necessary for Files API support (POST /v1/files endpoint)

59-79: The go-openapi ecosystem versions in lines 59–79 show no changes in this PR. These versions (runtime v0.29.2, swag v0.25.4, strfmt v0.25.0, etc.) were already present in the codebase and have not been updated. Review the actual changes being made before commenting on dependency updates.

Likely an incorrect or invalid review comment.

plugins/logging/go.mod (1)

42-64: OpenAPI swag module restructuring noted.

The go-openapi/swag v0.25.4 restructuring introduces several sub-packages (cmdutils, conv, fileutils, jsonname, etc.). This is consistent with the upstream library changes and aligns with updates across other modules in this PR stack.

core/go.mod (2)

5-23: Core module direct dependency updates are appropriate.

The updates to AWS SDK v2, Sonic, MCP-Go, fasthttp, and golang.org/x modules are consistent with the broader dependency refresh across the repository.

27-27: New Masterminds/semver indirect dependency.

The addition of Masterminds/semver/v3 v3.3.1 as an indirect dependency likely comes from a new transitive requirement. This is a well-maintained semantic versioning library.

plugins/governance/go.mod (1)

1-113: Governance plugin dependency updates align with the refactoring.

The dependency updates are consistent with the broader PR stack and support the governance store refactoring introduced in this PR.

plugins/otel/go.mod (2)

8-9: gRPC and protobuf direct dependency updates.

The updates to gRPC v1.77.0 and protobuf v1.36.11 are appropriate for the OTEL plugin to maintain compatibility with the latest OpenTelemetry ecosystem.

92-94: OpenTelemetry components updated to v1.39.0.

The OTEL core, metric, and trace packages are updated consistently to v1.39.0, maintaining API compatibility across the telemetry stack.

plugins/maxim/go.mod (1)

1-114: Maxim plugin dependency updates are consistent with the PR stack.

The indirect dependency updates align with the coordinated refresh across all modules.

plugins/governance/resolver.go (3)

64-75: Correct change from pointer-to-interface to interface type.

Changing store *GovernanceStore to store GovernanceStore is the idiomatic Go approach. Interfaces in Go already contain a pointer to the underlying data, so passing a pointer to an interface (*interface{}) is almost never necessary and adds unnecessary indirection.

194-217: Rate limit hierarchy check correctly prioritizes provider-level limits.

The implementation appropriately checks provider-level rate limits first by calling CheckRateLimit, then enriches the result with the relevant rateLimitInfo by checking provider configs before falling back to VK-level limits. This aligns with the hierarchical governance model.

219-233: Budget hierarchy check implementation is clean and consistent.

The atomic budget checking via CheckBudget prevents race conditions as noted in the comment. The error handling and logging are appropriate, and the implementation aligns with the 1:1 relationship invariant for budgets and their parent entities (per project learnings).

plugins/jsonparser/go.mod (1)

3-56: Dependency updates require verification against the full PR stack.

The go.mod updates show minor/patch version increments for AWS SDK, Sonic, fasthttp, and golang.org/x modules. However, the original claim that these "align with coordinated updates across core, framework, and other plugin modules" cannot be verified from this file alone. Per the coding guidelines, review must assess changes in the context of the entire PR stack, not as isolated updates. Confirm that all dependency versions match those used across related modules in this stack before approving.

plugins/governance/tracker.go (3)

34-34: LGTM: Clean refactor to interface-based store with in-memory tracking.

The change from *GovernanceStore to GovernanceStore interface is correct and enables the new in-memory/persistent store architecture. The usage of UpdateRateLimitUsageInMemory and UpdateBudgetUsageInMemory properly separates in-memory state updates from persistence, which is handled by the background worker.

Also applies to: 52-52, 89-89, 98-98

228-255: LGTM: Proper use of selective updates in transaction to preserve config fields.

The transactional startup reset correctly uses selective field updates (updates map) to only modify token_current_usage, token_last_reset, request_current_usage, and request_last_reset. This prevents overwriting max_limit or reset_duration that may have been changed during startup. The use of SkipHooks: true is appropriate for reset operations that don't require validation.

13-13: LGTM: Appropriate GORM import for transactional updates.

The addition of gorm.io/gorm import is necessary for the transactional startup reset logic and is used appropriately in ExecuteTransaction at line 230.

Also applies to: 230-230

plugins/governance/main.go (2)

40-47: LGTM: Well-designed interface for governance plugin abstraction.

The BaseGovernancePlugin interface provides a clean contract for governance plugin implementations, enabling testing and alternate implementations. All methods align with the existing GovernancePlugin implementation, and the return of GovernanceStore (interface) maintains the abstraction layer.

56-56: LGTM: Consistent interface-based store type usage.

The change from *GovernanceStore to GovernanceStore interface at line 56 and the matching return type at line 674 are consistent with the broader refactor to interface-based architecture.

Also applies to: 674-676

framework/configstore/tables/ratelimit.go (2)

33-35: LGTM: Properly tagged virtual fields for runtime tracking.

The LastDBTokenUsage and LastDBRequestUsage fields are correctly tagged with gorm:"-" and json:"-" to prevent persistence and serialization. The naming clearly indicates these are runtime-only baseline values from the last DB read.

83-88: LGTM: Clean AfterFind hook for baseline tracking.

The AfterFind hook correctly populates the virtual fields with current DB values after fetch. This enables the in-memory store to calculate deltas for efficient database updates. The implementation is simple and correct.

framework/configstore/tables/budget.go (1)

25-26: LGTM: Consistent virtual field pattern with rate limits.

The LastDBUsage virtual field and AfterFind hook follow the same pattern as TableRateLimit, providing a consistent approach for tracking DB baseline values across governance tables. The implementation correctly captures the CurrentUsage value after fetch for delta calculations.

Also applies to: 48-52

transports/bifrost-http/lib/lib.go (1)

18-22: LGTM: Well-designed abstraction for enterprise customization.

The EnterpriseOverrides interface provides a clean contract for customizing governance plugin names, loading custom governance implementations, and managing pricing. The interface is focused and enables dependency injection for testing and enterprise deployments.

transports/bifrost-http/handlers/middlewares.go (2)

49-49: LGTM: Proper nil guard for enterprise overrides.

The addition of the enterpriseOverrides parameter and the nil check at lines 58-61 correctly address the previous review concern. The middleware safely bypasses interception when enterprise overrides are not provided, preventing nil pointer dereference at line 65.

Also applies to: 58-61

65-65: LGTM: Clean integration with enterprise overrides for plugin detection.

Using enterpriseOverrides.GetGovernancePluginName() instead of a hardcoded constant enables enterprise customization of governance plugin names. The call is safely guarded by the nil check at lines 58-61.

plugins/telemetry/go.mod (1)

9-118: LGTM: Routine dependency updates aligned with core/framework.

The dependency updates across fasthttp, sonic, AWS SDK, OpenAPI modules, OpenTelemetry, and golang.org/x packages are routine maintenance to keep the telemetry plugin current. All updates appear to be minor or patch versions with no breaking changes. These align with the broader dependency updates in core and framework modules across the PR.

transports/bifrost-http/lib/config_test.go (1)

11483-11515: Improved hash parity failure messages look good

The updated t.Errorf messages that include “Before save” and “After load” hashes (and flags like StdioConfig populated) in the runtime‑vs‑migration parity tests are well‑structured and use matching format verbs/arguments, which will make diagnosing parity regressions much easier.

transports/bifrost-http/server/server.go (3)

499-517: LGTM: Proper locking pattern for governance plugin access.

The getGovernancePlugin helper correctly acquires PluginsMutex.RLock() before calling FindPluginByName, releases the lock immediately, then performs type assertion without holding the lock. This matches the pattern in GetGovernanceData and addresses the data-race concerns from previous reviews.

1061-1087: LGTM: Clean EnterpriseOverrides implementation.

The three methods implement the lib.EnterpriseOverrides interface, allowing the OSS build to use default governance/pricing while enabling enterprise builds to override with custom implementations. The self-reference pattern (passing s as EnterpriseOverrides to LoadPlugin) is correct and allows recursive resolution of overrides.

522-528: Constants verified: lib.DBLookupMaxRetries and lib.DBLookupDelay are properly defined.

The retry constants are defined in transports/bifrost-http/lib/config.go with reasonable values: DBLookupMaxRetries = 5 and DBLookupDelay = 1 * time.Second. These settings are appropriate for handling eventual consistency in database lookups during virtual key reloading.

transports/bifrost-http/lib/config.go (2)

275-275: LGTM: EnterpriseOverrides integration.

The EnterpriseOverrides parameter is correctly threaded through the configuration loading pipeline (LoadConfig → loadConfigFromFile/loadConfigFromDefaults → initFrameworkConfigFromFile/initDefaultFrameworkConfig). The conditional pricing manager initialization properly delegates to enterprise overrides when provided and falls back to default modelcatalog.Init when nil.

Also applies to: 298-298, 354-354, 400-400, 1505-1505, 1576-1576, 1618-1618, 1866-1866

1529-1544: PricingManager initialization failures result in silent feature degradation.

Both initFrameworkConfigFromFile and initDefaultFrameworkConfig log pricing manager initialization errors as warnings rather than propagating them, allowing config initialization to succeed with config.PricingManager set to nil. This is by design—when pricing manager fails to initialize, plugins (telemetry, logging, governance, otel) detect the nil value, log warnings, and continue with cost calculations and pricing features disabled.

All downstream code properly handles nil PricingManager through guard clauses and conditional checks. Verify that initialization failures are clearly visible in logs and that monitoring is in place if pricing/cost tracking is critical.

plugins/governance/store.go (9)

234-308: LGTM: Independent token and request limit enforcement.

The refactored logic correctly determines tokenExpired and requestExpired independently (lines 236-259), then conditionally applies each check (lines 272-279, 283-289). This ensures that when one limit window expires, the other limit is still enforced if active. This addresses the previous critical concern.

✅ Confirmed as addressed per previous review feedback.

405-437: LGTM: Clone-then-store pattern eliminates data race.

The method now clones the budget struct (line 420), mutates the copy (lines 421-424), and atomically replaces the sync.Map entry via Store (line 427). This prevents torn reads that could occur with in-place mutations. The updateBudgetReferences call (line 431) propagates the reset budget to all referencing entities.

✅ Confirmed as addressed per previous review feedback.

447-508: LGTM: Clone-then-store pattern for rate limits.

Similar to budget resets, this method uses the clone-then-store pattern (lines 474, 498) to avoid data races. The updateRateLimitReferences call (line 502) ensures all VKs and provider configs referencing this rate limit are updated with the reset copy.

✅ Confirmed as addressed per previous review feedback.

311-349: LGTM: Clone-then-store in budget usage updates.

The method creates a clone (line 324), performs all mutations on the clone (lines 325-342), and atomically stores the updated copy (line 340). This ensures thread-safe updates.

1211-1231: LGTM: Cascade deletion of budgets and rate limits.

The deletion correctly removes associated budgets and rate limits from the VK and its provider configs. Per the learning, budgets and rate limits maintain 1:1 relationships with their parent entities, so no shared-reference checks are needed.

Based on learnings: budgets and rate limits have 1:1 relationships with VKs, teams, and customers—no cross-entity sharing occurs.

661-673: LGTM: Deadlock-aware persistence for multi-node deployments.

Both DumpRateLimits and DumpBudgets detect database deadlocks (via error string matching for PostgreSQL SQLSTATE 40P01 and MySQL error 1213) and treat them as non-errors. This is appropriate for multi-node setups where concurrent writes are expected—the current node's data will propagate via gossip and persist in the next dump cycle.

Also applies to: 728-740

1536-1566: LGTM: Consistent budget and rate limit update logic.

The helpers implement a clear policy:

If currentUsage + baseline >= newMaxLimit, reset usage to 0 (handles limit reductions).

Otherwise, preserve existing usage and accept the new limit/duration.

This ensures that reducing a limit doesn't leave usage in an over-limit state. The checkAndUpdateRateLimit helper correctly applies this logic independently for token and request limits.

Also applies to: 1571-1616

41-93: LGTM: Clear interface contract and constructor.

The GovernanceStore interface explicitly documents error semantics (lines 44-49): non-nil errors from CheckBudget and CheckRateLimit indicate policy violations, not infrastructure failures. This prevents callers from retrying on explicit denials. The NewLocalGovernanceStore constructor cleanly handles both DB-backed and in-memory initialization paths.

1580-1585: The equalPtr generic utility function is correctly defined in plugins/governance/utils.go at line 20 and properly handles nil pointer comparisons with an explicit nil check before dereferencing values.

transports/go.mod

coderabbitai

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

plugins/telemetry/go.mod (1)

1-118: Review the actual impact of dependency updates against the codebase instead of assuming breaking changes.

The original concerns are overstated. AWS SDK S3 v1.94.0 adds support for LifecycleExpirationDate in S3 Inventory configurations—a feature addition, not a breaking change. The go.opentelemetry.io/otel/trace and metric packages are separate modules and remain so; there is no API split occurring. Recent OpenTelemetry-Go releases (v1.37.0–v1.38.0) contain feature additions and semantic convention upgrades, not breaking changes.

Instead of speculative verification requests, verify compatibility by:

Confirm the telemetry plugin code uses only stable OpenTelemetry APIs.

Run tests to ensure no import errors or runtime failures.

Check if any of the indirect dependencies have actual advisories via go list -m all | xargs -I {} go get -d {}@latest followed by checking the vulnerability database.

Drop the request for stack-wide version alignment unless there is evidence of actual version conflicts across PRs.

♻️ Duplicate comments (1)

framework/configstore/rdb.go (1)
1782-1789: Dead code: Find does not return gorm.ErrRecordNotFound for empty results.

GORM's Find method returns an empty slice (not an error) when no records match. The ErrRecordNotFound check at lines 1785-1787 is dead code and will never trigger.
🔎 Proposed fix
 func (s *RDBConfigStore) GetRateLimits(ctx context.Context) ([]tables.TableRateLimit, error) {
 	var rateLimits []tables.TableRateLimit
 	if err := s.db.WithContext(ctx).Find(&rateLimits).Error; err != nil {
-		if errors.Is(err, gorm.ErrRecordNotFound) {
-			return nil, ErrNotFound
-		}
 		return nil, err
 	}
 	return rateLimits, nil
 }

🧹 Nitpick comments (10)

docs/changelogs/v1.3.47.mdx (1)

20-25: Consider moving formatting changes to a separate PR.

These formatting changes to the v1.3.47 changelog are unrelated to the governance plugin refactor. While low-risk, separating them into a dedicated formatting/documentation cleanup PR would improve clarity and make the governance changes easier to review in this PR.

Also applies to: 29-34, 38-41
transports/bifrost-http/lib/config_test.go (3)
194-196: Mock RetryOnNotFound ignores retry semantics; consider simulating retries or clarifying intent

MockConfigStore.RetryOnNotFound just calls fn(ctx) once and returns, ignoring maxRetries and retryDelay. That’s fine if tests never care about retry behavior, but it diverges from the production contract and can hide bugs in code that depends on multiple attempts or idempotency under retries.

If you want a closer mock without adding real sleeps, you could loop up to maxRetries without delay and stop on first non‑not‑found error or success, or at least add a comment noting that this mock intentionally executes the callback only once so future tests don’t over‑assume its behavior.

317-319: Wire MockConfigStore.GetRateLimits to stored governance state and align return shape

Right now GetRateLimits always returns an empty slice and ignores m.governanceConfig.RateLimits. That makes the mock less useful for tests that create rate limits via CreateRateLimit and then want to verify them via the interface, and it’s also inconsistent with nearby getters like GetBudgets, GetCustomers, GetTeams, etc., which all return nil, nil when unset.

Consider implementing this similarly to other getters, e.g.:
Proposed GetRateLimits implementation
-func (m *MockConfigStore) GetRateLimits(ctx context.Context) ([]tables.TableRateLimit, error) {
-	return []tables.TableRateLimit{}, nil
-}
+func (m *MockConfigStore) GetRateLimits(ctx context.Context) ([]tables.TableRateLimit, error) {
+	if m.governanceConfig == nil || len(m.governanceConfig.RateLimits) == 0 {
+		return nil, nil
+	}
+	return m.governanceConfig.RateLimits, nil
+}
This keeps behavior consistent with the rest of the mock and better mirrors the real store.

6675-6775: LoadConfig(ctx, tempDir, nil) usage matches updated signature; consider helper to reduce repetition

The updates to call LoadConfig(ctx, tempDir, nil) across the SQLite integration tests correctly satisfy the new EnterpriseOverrides parameter while indicating “no overrides” via nil. Given how many tests follow this pattern, you might eventually want a small helper like loadTestConfig(ctx, dir) that wraps LoadConfig with a nil overrides argument to keep call sites DRY, but the current explicit calls are functionally sound.
plugins/governance/tracker.go (2)

104-125: Periodic reset + dump loop may warrant configurability for large deployments.

Using a fixed workerInterval of 10s to:

compute expired rate limits/budgets in memory,

persist resets, and

call DumpRateLimits(ctx, nil, nil) and DumpBudgets(ctx, nil)

gives tight consistency but can generate frequent DB I/O as state grows. Consider:

Making workerInterval configurable (e.g., via client/governance config) so operators can tune it by scale.

Optionally short‑circuiting dumps when nothing changed or using a less frequent dump cadence than the reset cadence.

Functionally this is correct; this is a scale/perf tuning knob. Based on learnings about hierarchical governance state being fully in-memory with periodic persistence.

Also applies to: 127-148

153-262: Startup rate-limit reset logic is sound but only covers VK-level limits; budgets use store path.

The startup flow:

Iterates all VKs and selectively zeroes token/request usage when the configured reset duration has elapsed, collecting only changed TableRateLimits.

Persists those changes via a GORM transaction that updates only the usage + last_reset columns, preserving max_limit / reset_duration.

Independently calls ResetExpiredBudgetsInMemory + ResetExpiredBudgets on the store for budgets.

Behavior looks correct and matches the GovernanceStore contract, but note:

Only VK-level rate limits are adjusted here; any future team/customer/provider‑level rate limits rely solely on the in-memory store’s periodic reset path.

If you expect startup to normalize all hierarchical rate limits, you might want a symmetric ResetExpiredRateLimitsInMemory‑based path for consistency with budgets, or at least document that startup only normalizes VK‑attached limits. Based on learnings about 1:1 relationships and hierarchical governance data.

plugins/governance/resolver.go (1)

129-137: Rate-limit hierarchy helper correctly centralizes decision logic but only surfaces VK/provider info.

checkRateLimitHierarchy:

Delegates to store.CheckRateLimit and treats any non-nil error as a policy violation, returning the Decision plus a formatted reason.

Locates a TableRateLimit to return in RateLimitInfo by checking provider‑level configs first, then VK-level vk.RateLimit.

This matches the documented contract that CheckRateLimit errors represent governance violations rather than infra failures and keeps EvaluationRequest logic lean. If you later add team/customer‑level rate limits, consider extending the RateLimitInfo selection (or adding a separate field) so callers can see which level in the hierarchy actually triggered the violation; today they only see provider/VK context. Based on learnings about hierarchical governance data and error semantics.

Also applies to: 194-217

transports/bifrost-http/server/server.go (1)

648-661: Consider reusing getGovernancePlugin inside GetGovernanceData to avoid duplicated lookup logic.

GetGovernanceData reimplements the governance plugin lookup with PluginsMutex.RLock + FindPluginByName + BaseGovernancePlugin assertion. For consistency and to avoid future drift, you could:

Delegate to getGovernancePlugin() and just return store.GetGovernanceData() when it succeeds.

This would also keep all governance plugin error handling and type assumptions in one place.

transports/bifrost-http/lib/config.go (1)

47-50: DBLookupMaxRetries/DBLookupDelay constants provide a clear home for DB retry tuning.

Defining DBLookupMaxRetries and DBLookupDelay here (and reusing them from the server for RetryOnNotFound) centralizes DB lookup retry policy and makes it easy to adjust if you see spurious ErrNotFound around governance reloads. Values (5 retries, 1s delay) are conservative but reasonable defaults.

plugins/governance/store.go (1)

96-149: Returned pointers expose internal state - ensure read-only access.

GetGovernanceData returns direct pointers to objects in the sync.Map. While acceptable for read-only API exposure, callers must not modify the returned data or it will corrupt the in-memory state. Consider documenting this constraint or returning clones if mutation by callers is a risk.

📜 Review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 7ae83e4 and 864eae5.

⛔ Files ignored due to path filters (11)

core/go.sum is excluded by !**/*.sum
framework/go.sum is excluded by !**/*.sum
plugins/governance/go.sum is excluded by !**/*.sum
plugins/jsonparser/go.sum is excluded by !**/*.sum
plugins/logging/go.sum is excluded by !**/*.sum
plugins/maxim/go.sum is excluded by !**/*.sum
plugins/mocker/go.sum is excluded by !**/*.sum
plugins/otel/go.sum is excluded by !**/*.sum
plugins/semanticcache/go.sum is excluded by !**/*.sum
plugins/telemetry/go.sum is excluded by !**/*.sum
transports/go.sum is excluded by !**/*.sum

📒 Files selected for processing (35)

.editorconfig
core/go.mod
core/mcp.go
core/providers/nebius/nebius.go
docs/changelogs/v1.3.47.mdx
docs/features/governance/virtual-keys.mdx
framework/configstore/rdb.go
framework/configstore/store.go
framework/configstore/tables/budget.go
framework/configstore/tables/mcp.go
framework/configstore/tables/ratelimit.go
framework/go.mod
framework/modelcatalog/main.go
framework/modelcatalog/sync.go
plugins/governance/go.mod
plugins/governance/main.go
plugins/governance/resolver.go
plugins/governance/store.go
plugins/governance/tracker.go
plugins/governance/utils.go
plugins/jsonparser/go.mod
plugins/logging/go.mod
plugins/maxim/go.mod
plugins/mocker/go.mod
plugins/otel/go.mod
plugins/semanticcache/go.mod
plugins/telemetry/go.mod
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/lib/lib.go
transports/bifrost-http/server/server.go
transports/go.mod
ui/lib/constants/logs.ts

💤 Files with no reviewable changes (1)

framework/configstore/tables/mcp.go

✅ Files skipped from review due to trivial changes (1)

ui/lib/constants/logs.ts

🚧 Files skipped from review as they are similar to previous changes (13)

core/providers/nebius/nebius.go
framework/configstore/tables/budget.go
plugins/semanticcache/go.mod
docs/features/governance/virtual-keys.mdx
framework/go.mod
.editorconfig
framework/modelcatalog/sync.go
plugins/logging/go.mod
framework/modelcatalog/main.go
framework/configstore/store.go
plugins/governance/main.go
plugins/governance/utils.go
core/go.mod

🧰 Additional context used

📓 Path-based instructions (1)

**

⚙️ CodeRabbit configuration file

always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)

Files:

core/mcp.go
framework/configstore/tables/ratelimit.go
plugins/mocker/go.mod
framework/configstore/rdb.go
docs/changelogs/v1.3.47.mdx
transports/bifrost-http/lib/lib.go
plugins/governance/resolver.go
plugins/governance/tracker.go
transports/bifrost-http/lib/config_test.go
plugins/governance/go.mod
transports/bifrost-http/handlers/governance.go
plugins/jsonparser/go.mod
transports/bifrost-http/server/server.go
plugins/otel/go.mod
transports/bifrost-http/lib/config.go
transports/bifrost-http/handlers/middlewares.go
plugins/maxim/go.mod
plugins/telemetry/go.mod
transports/go.mod
plugins/governance/store.go

🧠 Learnings (3)

📚 Learning: 2025-12-09T17:07:42.007Z

Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.

Applied to files:

core/mcp.go
framework/configstore/tables/ratelimit.go
framework/configstore/rdb.go
transports/bifrost-http/lib/lib.go
plugins/governance/resolver.go
plugins/governance/tracker.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/server/server.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/handlers/middlewares.go
plugins/governance/store.go

📚 Learning: 2025-12-12T08:25:02.629Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.

Applied to files:

transports/bifrost-http/lib/lib.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/server/server.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/handlers/middlewares.go

📚 Learning: 2025-12-22T10:50:40.990Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.

Applied to files:

plugins/governance/resolver.go
plugins/governance/tracker.go
plugins/governance/store.go

🧬 Code graph analysis (10)

framework/configstore/rdb.go (5)

framework/configstore/tables/virtualkey.go (6)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

TableVirtualKeyMCPConfig (81-87)

TableVirtualKeyMCPConfig (90-92)

framework/configstore/tables/budget.go (2)

TableBudget (11-27)

TableBudget (30-30)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

framework/configstore/tables/team.go (2)

TableTeam (12-38)

TableTeam (41-41)

framework/configstore/tables/customer.go (2)

TableCustomer (6-22)

TableCustomer (25-25)

transports/bifrost-http/lib/lib.go (3)

core/schemas/plugin.go (1)

Plugin (61-87)

framework/configstore/store.go (1)

ConfigStore (18-157)

framework/modelcatalog/main.go (1)

ModelCatalog (23-46)

plugins/governance/resolver.go (4)

plugins/governance/store.go (1)

GovernanceStore (50-72)

framework/configstore/tables/virtualkey.go (2)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

core/schemas/bifrost.go (1)

ModelProvider (32-32)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

plugins/governance/tracker.go (1)

plugins/governance/resolver.go (1)

BudgetResolver (64-67)

transports/bifrost-http/lib/config_test.go (3)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

core/schemas/logger.go (1)

LogLevel (6-6)

transports/bifrost-http/lib/config.go (1)

LoadConfig (275-355)

transports/bifrost-http/handlers/governance.go (3)

plugins/governance/store.go (1)

GovernanceData (33-39)

framework/configstore/tables/virtualkey.go (2)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

ui/lib/types/governance.ts (1)

Budget (5-11)

transports/bifrost-http/server/server.go (3)

plugins/governance/store.go (1)

GovernanceData (33-39)

plugins/governance/main.go (4)

Config (32-34)

PluginName (21-21)

BaseGovernancePlugin (40-47)

GovernancePlugin (50-69)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

transports/bifrost-http/lib/config.go (1)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

transports/bifrost-http/handlers/middlewares.go (2)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

transports/bifrost-http/lib/middleware.go (1)

BifrostHTTPMiddleware (7-7)

plugins/governance/store.go (8)

framework/configstore/tables/virtualkey.go (4)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

framework/configstore/tables/team.go (2)

TableTeam (12-38)

TableTeam (41-41)

framework/configstore/tables/customer.go (2)

TableCustomer (6-22)

TableCustomer (25-25)

framework/configstore/tables/budget.go (2)

TableBudget (11-27)

TableBudget (30-30)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

framework/configstore/clientconfig.go (1)

GovernanceConfig (720-727)

core/schemas/models.go (1)

Model (109-129)

ui/lib/types/governance.ts (4)

Budget (5-11)

RateLimit (13-25)

Team (27-35)

Customer (37-44)

🔇 Additional comments (43)

plugins/mocker/go.mod (2)

15-57: All indirect dependencies are at secure versions with no reported CVEs.

The updated versions of AWS SDK v2 (v1.32.6), Sonic (v1.14.2), Fasthttp (v1.68.0), and mcp-go (v0.43.2) do not have publicly reported security vulnerabilities. These are patch-level maintenance releases with routine bug fixes and no known security issues as of December 2025.

43-43: No actionable verification needed for mcp-go 0.43.2 compatibility.

The version bump from 0.41.1 to 0.43.2 reflects incremental updates consisting solely of bug fixes (notification handling, schema marshaling). No breaking changes are documented for versions 0.42.x or 0.43.x. This dependency update is aligned with ecosystem standards for MCP client libraries and poses no risk to the governance plugin refactor.

plugins/maxim/go.mod (1)

18-109: Verify compatibility of 0.x version updates for pre-release dependencies.

Most transitive dependency updates appear to be patch-level maintenance releases. redis/go-redis v9.17.2 contains only bugfixes for connection pool race conditions and context timeout handling, and go-openapi packages are consolidated to v0.25.x for good cohesion. However, pre-release packages with 0.x versioning should be reviewed:

mark3labs/mcp-go: v0.41.1 → v0.43.2 (0.x package; under active development with some advanced capabilities still in progress)

bytedance/sonic/loader: v0.3.0 → v0.4.0 (0.x minor bump)

Confirm these specific updates don't introduce breaking changes that would impact plugin functionality.

plugins/otel/go.mod (3)

41-61: The coordinated go-openapi package versions (v0.25.x series and related versions) maintain backward compatibility. The traditional swag API is still supported and code that imports swag will still compile and work the same.

8-9: Ensure updated dependencies remain compatible with OpenTelemetry integration.

The gRPC v1.77.0 and protobuf v1.36.11 updates are secure; both versions are well beyond known CVE fixes. Weaviate v1.34.5 is a stable recent release with documented fixes and no API breaking changes. Verify that gRPC/protobuf coordinated updates don't introduce compatibility issues with any generated code in the OpenTelemetry plugins, particularly for gRPC server configurations. Check integration tests pass with the updated ecosystem versions.

92-94: Verify OpenTelemetry v1.39.0 availability and compatibility before approval.

The opentelemetry-go main branch go.mod references v1.39.0 for otel/metric and otel/trace, but the latest documented stable release is v1.38.0, which is the last to support Go 1.23, with the next release requiring at least Go 1.24. Cannot confirm v1.39.0 as officially released or verify backward compatibility claims without access to official release notes. Requires verification that v1.39.0 has been publicly released and compatibility with existing otel code before merging.

core/mcp.go (1)

1134-1151: No changes needed - the type is already enforced by schema definition.

The field InProcessServer in schemas.MCPClientConfig is defined as *server.MCPServer, which means the type is already statically enforced at the schema level. The removed assertion was redundant and not providing additional safety. The current implementation correctly passes the already-typed field to client.NewInProcessClient, which expects the same type.

Likely an incorrect or invalid review comment.

framework/configstore/tables/ratelimit.go (1)

32-36: LGTM! Runtime fields for tracking DB state are well-implemented.

The virtual fields LastDBTokenUsage and LastDBRequestUsage correctly capture the DB state after each fetch via the AfterFind hook, enabling in-memory comparison to detect resets or changes. The use of gorm:"-" and json:"-" ensures these fields are not persisted or serialized.

Also applies to: 83-88

framework/configstore/rdb.go (5)

669-669: Minor formatting cleanup.

Removed stray whitespace in Bedrock ARN processing block.

1657-1690: LGTM! Transactional team deletion with FK nullification.

The transactional approach correctly nullifies team_id on related VirtualKey records and deletes the team's budget before removing the team itself, ensuring referential integrity.

1743-1780: LGTM! Transactional customer deletion with FK nullification.

The implementation correctly nullifies customer_id on related VirtualKey and Team records, then deletes the customer's budget before removing the customer, maintaining referential integrity within a transaction.

2136-2161: LGTM! RetryOnNotFound helper correctly implements retry logic.

The implementation properly retries on ErrNotFound or gorm.ErrRecordNotFound up to maxRetries times with configurable delays, respecting context cancellation. The use of time.After with select on ctx.Done() ensures the retry loop can be interrupted.

1394-1458: GORM foreign key constraints are properly configured and aligned with manual cascading logic.

Verification confirms that TableVirtualKey and TableVirtualKeyProviderConfig both have onDelete:CASCADE configured for their Budget and RateLimit relationships. The many-to-many relationship for provider config keys also includes constraint:OnDelete:CASCADE. The manual deletion sequence in the transaction correctly mirrors these constraints, ensuring consistency between the schema and runtime behavior.

transports/bifrost-http/lib/lib.go (1)

4-9: LGTM! Clean abstraction for enterprise feature injection.

The EnterpriseOverrides interface provides a clear contract for governance plugin and pricing manager integration, enabling flexible enterprise feature loading without tight coupling to specific implementations.

Also applies to: 18-22

transports/bifrost-http/handlers/middlewares.go (1)

49-73: LGTM! Nil-check for enterpriseOverrides properly addresses previous concern.

The middleware now safely handles nil enterpriseOverrides by bypassing interception, preventing any nil pointer dereferences when calling GetGovernancePluginName(). The control flow correctly skips interception if either enterpriseOverrides is nil or the governance plugin is not loaded.

transports/bifrost-http/handlers/governance.go (6)

25-25: LGTM! Interface extension and nil-check for governance manager.

Adding GetGovernanceData() to the GovernanceManager interface enables in-memory data access, and the nil-check in NewGovernanceHandler prevents panics when the manager is not provided.

Also applies to: 42-44

178-181: LGTM! New routes for budgets and rate-limits.

The new GET endpoints /api/governance/budgets and /api/governance/rate-limits provide read-only access to governance data, aligning with the in-memory data exposure pattern.

188-201: LGTM! All from_memory handlers include nil-checks for GetGovernanceData().

Each handler that accesses in-memory governance data via GetGovernanceData() correctly checks if the returned data is nil before dereferencing its fields, preventing nil pointer panics. This addresses the concerns raised in previous reviews.

Also applies to: 422-440, 950-976, 1059-1076, 1209-1222, 1301-1318, 1449-1473, 1476-1501

310-329: LGTM! Key association logic validates key count and handles errors.

The code correctly retrieves keys by IDs using GetKeysByIDs, validates that all requested keys were found by comparing lengths, and returns descriptive errors if any keys are missing. This ensures provider configs only reference valid keys.

Also applies to: 635-655, 706-719

575-577: LGTM! Rate limit validation calls ensure data integrity.

Adding validateRateLimit() calls after updating rate limit fields ensures that all constraints (non-negative/non-zero limits, valid durations, required pairings) are enforced before persisting changes.

Also applies to: 788-790

936-936: LGTM! Error logging improves observability.

Adding error logging when virtual key deletion fails provides valuable debugging information for operators.

plugins/jsonparser/go.mod (1)

12-13: Verify dependency versions for validity and security advisories.

All indirect dependency versions exist and are valid published releases. fasthttp versions before 1.34.0 have a Directory Traversal vulnerability, but the project uses v1.68.0, which is not affected. No known security vulnerabilities were identified for AWS SDK v2, Sonic, Klauspost compress, or golang.org/x modules in their current versions.

transports/bifrost-http/lib/config_test.go (3)

624-633: testLogger stub correctly satisfies logger interface

The no‑op testLogger now implements Debug/Info/Warn/Error/Fatal and SetLevel(level schemas.LogLevel), matching the logger interface without side effects, which is appropriate for these tests.

2199-2207: Nice explicit collision check coverage for ProviderConfig hash variants

The hashes map enumerating no_optional, with_network, with_proxy, with_conc, with_custom, with_raw, and all_fields and then checking for collisions gives good, localized assurance that each optional field combination produces a distinct provider hash. The use of a seen map keyed by hash is a clear and inexpensive way to enforce this.

1833-1840: Expanded key match condition in DifferentHash path looks logically correct

The updated condition if dbKeyHash == fileKeyHash || fileKey.Name == dbKey.Name when preserving dashboard‑added keys in the provider hash mismatch path makes “sameness” depend on either identical content hash or shared name. That avoids treating a renamed/edited file key with the same logical name as a separate dashboard key, while still ensuring genuine DB‑only keys are appended. This aligns with the intended semantics and keeps the merge behavior predictable.

transports/go.mod (1)

7-8: Dependency bumps and new mcp-go dep look consistent; ensure full test run across transports/governance paths.

Version updates for Sonic, fasthttp, AWS SDK v2 components, Redis, Qdrant, Weaviate (core + client), OTEL, and OpenAPI tooling plus the new github.com/mark3labs/mcp-go v0.43.2 dependency are all reasonable on paper, but compatibility with existing transport/governance code (vector stores, HTTP handlers, MCP integration) ultimately depends on runtime tests. Please confirm go test ./... (and any governance/MCP integration tests) pass under this new dependency set.

Also applies to: 11-12, 22-25, 31-32, 42-45, 52-53, 59-79, 90-90, 105-107, 112-113, 116-120, 123-133

plugins/governance/tracker.go (3)

34-45: GovernanceStore as interface field + tracker wiring look coherent.

Switching store to the GovernanceStore interface and threading it via NewUsageTracker aligns with the rest of the governance stack (e.g., resolver using the same interface) and keeps the tracker decoupled from concrete store implementations. Background worker setup (trackerCtx/trackerCancel, done, wg) remains consistent with prior design.

Also applies to: 51-66

87-101: Streaming gating + in-memory usage updates are aligned with hierarchical store semantics.

The shouldUpdateTokens/Requests/Budget conditions correctly avoid double-counting streaming chunks, and switching to UpdateRateLimitUsageInMemory / UpdateBudgetUsageInMemory keeps the tracker consistent with the new in-memory GovernanceStore contract (fast-path updates with periodic dumps). No obvious race or logic issues here assuming the store implementation is concurrency-safe.

264-279: Cleanup assumes single invocation; safe under current plugin lifecycle.

Cleanup unconditionally closes t.done and stops the ticker after canceling the tracker context. Given the plugin lifecycle guarantees (each governance plugin instance cleaned up exactly once on shutdown/reload), this is acceptable and keeps the worker shutdown simple. Just ensure tests and any future call sites also maintain that single‑call invariant.

plugins/governance/resolver.go (2)

63-75: Resolver now depends directly on GovernanceStore interface; constructor matches.

Using GovernanceStore (non-pointer) in BudgetResolver and NewBudgetResolver keeps resolver aligned with the store abstraction used by tracker and avoids leaking a concrete store type. This should simplify testing and alternate store implementations.

219-233: Atomic budget check helper cleanly mirrors CheckBudget semantics.

checkBudgetHierarchy’s use of store.CheckBudget(ctx, vk, request, nil) and mapping any error to DecisionBudgetExceeded with debug logging is consistent with the interface’s “error == violation” contract and the new atomic/hierarchical budget behavior. No issues from a correctness standpoint.

transports/bifrost-http/server/server.go (4)

201-210: Governance plugin loading and plugin state wiring are consistent with EnterpriseOverrides abstraction.

GovernanceInMemoryStore.GetConfiguredProviders correctly uses Config.Mu.RLock to expose providers to the governance plugin without extra copies.

LoadPlugin/LoadPlugins now thread EnterpriseOverrides, allowing the same hook to control governance plugin and pricing manager behavior while keeping other built-ins unchanged.

SyncLoadedPlugin and RemovePlugin still use CAS over Config.Plugins and protect BifrostHTTPServer.Plugins with PluginsMutex, so the new uses of LoadPlugin in ReloadPlugin remain race-free.

This keeps the server’s plugin model coherent with the new governance stack and EnterpriseOverrides interface.

Also applies to: 269-273, 324-439, 773-812

499-517: New getGovernancePlugin helper fixes prior races and centralizes governance plugin access.

getGovernancePlugin now wraps governance plugin lookup with PluginsMutex.RLock and performs a BaseGovernancePlugin assertion once, returning a strongly-typed instance.

ReloadVirtualKey, RemoveVirtualKey, ReloadTeam/RemoveTeam, and ReloadCustomer/RemoveCustomer all use this helper, eliminating the prior unprotected reads of s.Plugins and duplicated type assertions.

The use of ConfigStore.RetryOnNotFound with DBLookupMaxRetries/DBLookupDelay in ReloadVirtualKey is a reasonable way to handle eventual consistency around VK creation.

Overall, these changes fix the earlier concurrency concerns around plugin access and make governance reload helpers more robust.

Also applies to: 519-572, 574-610, 612-646, 648-661

943-1006: EnterpriseOverrides threading through RegisterAPIRoutes and TransportInterceptorMiddleware is coherent.

RegisterAPIRoutes now receives EnterpriseOverrides and uses GetGovernancePluginName() to conditionally initialize the governance handler, which keeps HTTP wiring agnostic to the actual plugin name.

Bootstrap passes s as both ServerCallbacks and EnterpriseOverrides, and inference middlewares are prefixed with TransportInterceptorMiddleware(s.Config, s), so the interceptor can also honor EnterpriseOverrides (e.g., custom governance plugin or pricing behavior).

This wiring cleanly integrates the new EnterpriseOverrides interface into the HTTP layer without changing existing handler contracts.

Also applies to: 1225-1235

1061-1068: The review comment cannot be verified because the code it references does not exist in the repository. Specifically:

GetGovernancePluginName() function is not found anywhere in the codebase

LoadGovernancePlugin() function is not found anywhere in the codebase

OSSToEnterprisePluginNameOverrides field does not exist in the BifrostHTTPServer struct

The actual code at lines 1061-1068 contains the Bootstrap() function, not GetGovernancePluginName(). The current implementation uses governance.PluginName as a constant (line 263, 378, 482, 508, 555, 573, 608, 626, 937) without any override mechanism.

To verify this review comment, please provide the specific code changes, branch, or pull request being reviewed, as the referenced implementation is not present in the current main branch.

transports/bifrost-http/lib/config.go (1)

275-299: EnterpriseOverrides integration keeps default behavior while enabling custom pricing/governance wiring.

LoadConfig now accepts an EnterpriseOverrides and passes it into both loadConfigFromFile and loadConfigFromDefaults, which in turn route into initFrameworkConfigFromFile / initDefaultFrameworkConfig.

Both framework init helpers:

Build a modelcatalog.Config from store/file data exactly as before.

If EnterpriseOverrides != nil, delegate pricing manager creation to EnterpriseOverrides.LoadPricingManager.

Otherwise, fall back to the previous modelcatalog.Init path.

loadConfigFromDefaults similarly passes EnterpriseOverrides into initDefaultFrameworkConfig, so the default “no config.json” path also supports custom pricing initialization.

This preserves earlier behavior when callers pass nil (e.g., tests calling LoadConfig(ctx, dir, nil)) while letting the HTTP server (or future enterprise entrypoints) take control of how pricing is wired.

Also applies to: 355-408, 1505-1546, 1574-1629, 1866-1928

plugins/governance/store.go (7)

41-72: Well-documented interface contract.

The error semantics contract in the doc comment clearly establishes that CheckRateLimit and CheckBudget return errors for policy violations, not infrastructure failures. This prevents callers from incorrectly retrying on governance denials.

234-308: Previous logic error properly addressed.

The rate limit checking now correctly handles token and request expiration independently. Each limit type is checked based on its own expiration state, preventing scenarios where an expired token window would incorrectly bypass an active request limit check (and vice versa).

418-437: Data race issue properly addressed with clone-then-store pattern.

The budget reset now correctly creates a copy (copiedBudget := *budget), modifies the copy, and atomically stores it back via gs.budgets.Store(key, &copiedBudget). This prevents readers from observing torn/inconsistent state.

1082-1222: Complex update logic correctly handles baselines and clone-then-store.

The method properly:

Retrieves existing budget/rate limit from the authoritative gs.budgets/gs.rateLimits maps rather than potentially stale VK references

Uses the checkAndUpdateBudget/checkAndUpdateRateLimit helpers to preserve usage unless limits change

Creates a map of existing provider configs for efficient lookup when detecting removals

Falls back to CreateVirtualKeyInMemory if the VK doesn't exist

1238-1266: Cascade deletion correctly handles 1:1 relationship.

Based on the established design that budgets and rate limits maintain 1:1 relationships with parent entities, the cascade deletion logic correctly removes associated resources. The VK reference cleanup for teams and customers properly uses the clone-then-store pattern.

661-674: Deadlock handling follows eventual consistency design.

The approach of treating deadlocks as non-errors in the multi-node setup is reasonable given that usage data syncs via gossip and will be written in the next dump cycle. The error string matching for deadlock detection (checking "deadlock", PostgreSQL "40P01", MySQL "1213") is a common pattern.

1478-1574: Reference update helpers correctly propagate reset state.

These helpers ensure that when a budget or rate limit is reset, all entities holding a reference to it receive the updated state. The O(n) iteration is acceptable since these run during background reset operations, not on the critical request path.

plugins/governance/store.go

coderabbitai

Actionable comments posted: 2

♻️ Duplicate comments (2)

transports/go.mod (1)
105-113: Verify weaviate v5.6.0 breaking change and cross-module version alignment in governance stack.

The weaviate-go-client v5.6.0 update (line 113) introduces a breaking change where models.Vector changed from []float32 to interface{} for multi-dimensional vector support. This is unresolved from a previous review comment and is critical for the governance plugin refactor, as the new LocalGovernanceStore and vector-based enforcement patterns depend on correct Vector type handling.

Additionally, verify that:

All weaviate Vector type usages throughout the codebase handle the interface{} type correctly (type assertions, nil checks, etc.)

redis/go-redis/v9 v9.17.2 (line 106) is fully compatible with the hierarchical rate-limit/budget structures in LocalGovernanceStore

qdrant/go-client v1.16.2 (line 105) vector APIs align with governance enforcement patterns

This go.mod's dependency versions match those in the governance plugin module (the PR stack includes plugin updates; ensure consistency across stack modules)

To verify weaviate Vector usage patterns and stack module alignment, run:
#!/bin/bash
# Description: Verify weaviate Vector type handling and cross-module dependency alignment

# Search for weaviate Vector type usages in codebase
echo "=== Weaviate Vector type usages ==="
rg -nP --type=go '\bVector\b' --context=3 | head -100

# Check if governance plugin module's go.mod exists and compare versions
echo -e "\n=== Governance plugin module dependency versions ==="
if [ -f "plugins/governance/go.mod" ]; then
  rg -n "weaviate|redis|qdrant" plugins/governance/go.mod
fi

# Check core module alignment
echo -e "\n=== Core module dependency versions ==="
if [ -f "core/go.mod" ]; then
  rg -n "weaviate|redis|qdrant" core/go.mod
fi
plugins/governance/store.go (1)
1612-1614: Doc comment inconsistency: > vs >= comparison.

The doc comment states "If currentUsage+baseline > newMaxLimit" but the actual code at lines 1634 and 1652 uses >= comparison. This should be consistent with the implementation.
🔎 Proposed fix
 // checkAndUpdateRateLimit checks and updates a rate limit with usage reset logic
-// If currentUsage+baseline > newMaxLimit, reset usage to 0
+// If currentUsage+baseline >= newMaxLimit, reset usage to 0
 // Otherwise preserve existing usage and accept reset duration and max limit changes
Note: This issue was flagged in a previous review but appears to not have been addressed yet.

🧹 Nitpick comments (5)

transports/bifrost-http/lib/config_test.go (1)

130-137: Mock RetryOnNotFound is a simple passthrough (acceptable for tests)

The mock implementation of RetryOnNotFound just calls fn(ctx) once and ignores maxRetries/retryDelay. That’s fine if tests using the mock only care about the happy path and not retry timing, but it means you’re not exercising any retry/backoff behavior through this mock.

If, in future, you add tests around retry semantics, you may want a slightly richer mock that at least tracks invocation count or simulates a transient-not-found before succeeding, but it’s not required for this PR.

Also applies to: 194-196

plugins/governance/tracker.go (1)

47-49: Make worker interval / dump cadence configurable for large deployments

workerInterval is hard‑coded to 10s while resetExpiredCounters now unconditionally calls DumpRateLimits and DumpBudgets each tick. On installations with many rate limits/budgets this can create avoidable DB load.

Consider either:

making workerInterval configurable (or at least bumping its default), and/or

decoupling “expiry reset” from “full dump” so dumps can run less frequently.

Also applies to: 141-147

plugins/governance/resolver.go (1)

194-217: Consider surfacing team/customer rate‑limit info in violations

checkRateLimitHierarchy picks RateLimitInfo only from provider‑level configs or the VK’s own RateLimit. If a rate‑limit decision comes from a team/customer‑level limit (which exist per the 1:1 invariants), RateLimitInfo will stay nil even though Decision indicates a violation.

If API consumers are expected to see the concrete limit that triggered the block, consider either:

extending GovernanceStore.CheckRateLimit to return the concrete *TableRateLimit, or

augmenting this helper to also check any team/customer RateLimit attached to the hierarchy.

Based on learnings, budgets/rate limits are 1:1 with VKs/teams/customers.

plugins/governance/main.go (1)

71-82: Deduplicate common Init / InitFromStore wiring (optional)

Init and InitFromStore largely duplicate the same steps (derive isVkMandatory, build resolver and tracker, run startup resets, create plugin struct). To reduce drift on future changes, consider extracting a small internal helper that takes a GovernanceStore and shared deps, and have both constructors call it.

Also applies to: 166-224

transports/bifrost-http/server/server.go (1)

499-517: getGovernancePlugin helper correctly centralizes locked lookup and type assertion

The new getGovernancePlugin method wraps FindPluginByName with PluginsMutex.RLock/RUnlock and enforces that the plugin implements governance.BaseGovernancePlugin. Using this helper in ReloadVirtualKey, RemoveVirtualKey, ReloadTeam, RemoveTeam, ReloadCustomer, and RemoveCustomer fixes the earlier data‑race risk on s.Plugins and makes the in‑memory governance updates safer.

You might also consider reusing this helper in GetGovernanceData to avoid duplicating the lookup and assertion logic.

Also applies to: 552-572, 575-610, 613-646

📜 Review details

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 864eae5 and d877912.

⛔ Files ignored due to path filters (11)

core/go.sum is excluded by !**/*.sum
framework/go.sum is excluded by !**/*.sum
plugins/governance/go.sum is excluded by !**/*.sum
plugins/jsonparser/go.sum is excluded by !**/*.sum
plugins/logging/go.sum is excluded by !**/*.sum
plugins/maxim/go.sum is excluded by !**/*.sum
plugins/mocker/go.sum is excluded by !**/*.sum
plugins/otel/go.sum is excluded by !**/*.sum
plugins/semanticcache/go.sum is excluded by !**/*.sum
plugins/telemetry/go.sum is excluded by !**/*.sum
transports/go.sum is excluded by !**/*.sum

📒 Files selected for processing (35)

.editorconfig
core/go.mod
core/mcp.go
core/providers/nebius/nebius.go
docs/changelogs/v1.3.47.mdx
docs/features/governance/virtual-keys.mdx
framework/configstore/rdb.go
framework/configstore/store.go
framework/configstore/tables/budget.go
framework/configstore/tables/mcp.go
framework/configstore/tables/ratelimit.go
framework/go.mod
framework/modelcatalog/main.go
framework/modelcatalog/sync.go
plugins/governance/go.mod
plugins/governance/main.go
plugins/governance/resolver.go
plugins/governance/store.go
plugins/governance/tracker.go
plugins/governance/utils.go
plugins/jsonparser/go.mod
plugins/logging/go.mod
plugins/maxim/go.mod
plugins/mocker/go.mod
plugins/otel/go.mod
plugins/semanticcache/go.mod
plugins/telemetry/go.mod
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/lib/lib.go
transports/bifrost-http/server/server.go
transports/go.mod
ui/lib/constants/logs.ts

💤 Files with no reviewable changes (1)

framework/configstore/tables/mcp.go

✅ Files skipped from review due to trivial changes (2)

.editorconfig
docs/changelogs/v1.3.47.mdx

🚧 Files skipped from review as they are similar to previous changes (12)

core/mcp.go
core/providers/nebius/nebius.go
ui/lib/constants/logs.ts
docs/features/governance/virtual-keys.mdx
framework/modelcatalog/main.go
framework/modelcatalog/sync.go
plugins/governance/utils.go
framework/configstore/tables/budget.go
framework/configstore/tables/ratelimit.go
transports/bifrost-http/lib/lib.go
plugins/otel/go.mod
plugins/maxim/go.mod

🧰 Additional context used

📓 Path-based instructions (1)

**

⚙️ CodeRabbit configuration file

always check the stack if there is one for the current PR. do not give localized reviews for the PR, always see all changes in the light of the whole stack of PRs (if there is a stack, if there is no stack you can continue to make localized suggestions/reviews)

Files:

transports/bifrost-http/handlers/middlewares.go
core/go.mod
framework/configstore/rdb.go
plugins/semanticcache/go.mod
plugins/logging/go.mod
transports/bifrost-http/handlers/governance.go
plugins/governance/tracker.go
plugins/jsonparser/go.mod
plugins/governance/main.go
framework/configstore/store.go
plugins/mocker/go.mod
plugins/governance/resolver.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
plugins/governance/go.mod
transports/go.mod
transports/bifrost-http/server/server.go
framework/go.mod
plugins/telemetry/go.mod
plugins/governance/store.go

🧠 Learnings (3)

📚 Learning: 2025-12-09T17:07:42.007Z

Learnt from: qwerty-dvorak
Repo: maximhq/bifrost PR: 1006
File: core/schemas/account.go:9-18
Timestamp: 2025-12-09T17:07:42.007Z
Learning: In core/schemas/account.go, the HuggingFaceKeyConfig field within the Key struct is currently unused and reserved for future Hugging Face inference endpoint deployments. Do not flag this field as missing from OpenAPI documentation or require its presence in the API spec until the feature is actively implemented and used. When the feature is added, update the OpenAPI docs accordingly; otherwise, treat this field as non-breaking and not part of the current API surface.

Applied to files:

transports/bifrost-http/handlers/middlewares.go
framework/configstore/rdb.go
transports/bifrost-http/handlers/governance.go
plugins/governance/tracker.go
plugins/governance/main.go
framework/configstore/store.go
plugins/governance/resolver.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/server/server.go
plugins/governance/store.go

📚 Learning: 2025-12-12T08:25:02.629Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1000
File: transports/bifrost-http/integrations/router.go:709-712
Timestamp: 2025-12-12T08:25:02.629Z
Learning: In transports/bifrost-http/**/*.go, update streaming response handling to align with OpenAI Responses API: use typed SSE events such as response.created, response.output_text.delta, response.done, etc., and do not rely on the legacy data: [DONE] termination marker. Note that data: [DONE] is only used by the older Chat Completions and Text Completions streaming APIs. Ensure parsers, writers, and tests distinguish SSE events from the [DONE] sentinel and handle each event type accordingly for correct stream termination and progress updates.

Applied to files:

transports/bifrost-http/handlers/middlewares.go
transports/bifrost-http/handlers/governance.go
transports/bifrost-http/lib/config.go
transports/bifrost-http/lib/config_test.go
transports/bifrost-http/server/server.go

📚 Learning: 2025-12-22T10:50:40.990Z

Learnt from: Pratham-Mishra04
Repo: maximhq/bifrost PR: 1154
File: plugins/governance/store.go:1165-1186
Timestamp: 2025-12-22T10:50:40.990Z
Learning: In the Bifrost governance plugin, budgets and rate limits have 1:1 relationships with their parent entities (virtual keys, teams, customers). Do not assume sharing; ensure cascade deletion logic only deletes budgets/rate limits when there are no shared references. Enforce invariants in code and add tests to verify no cross-entity sharing and that cascade deletes only remove the specific child of the parent. If a counterexample arises, adjust data model or add guards.

Applied to files:

plugins/governance/tracker.go
plugins/governance/main.go
plugins/governance/resolver.go
plugins/governance/store.go

🧬 Code graph analysis (9)

transports/bifrost-http/handlers/middlewares.go (2)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

transports/bifrost-http/lib/middleware.go (1)

BifrostHTTPMiddleware (7-7)

framework/configstore/rdb.go (5)

framework/configstore/tables/virtualkey.go (6)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

TableVirtualKeyMCPConfig (81-87)

TableVirtualKeyMCPConfig (90-92)

framework/configstore/tables/budget.go (2)

TableBudget (11-27)

TableBudget (30-30)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

framework/configstore/tables/team.go (2)

TableTeam (12-38)

TableTeam (41-41)

framework/configstore/tables/customer.go (2)

TableCustomer (6-22)

TableCustomer (25-25)

transports/bifrost-http/handlers/governance.go (4)

plugins/governance/store.go (1)

GovernanceData (33-39)

framework/configstore/tables/key.go (2)

TableKey (13-58)

TableKey (61-61)

framework/configstore/tables/virtualkey.go (2)

TableVirtualKeyProviderConfig (23-36)

TableVirtualKeyProviderConfig (39-41)

ui/lib/types/governance.ts (1)

Budget (5-11)

plugins/governance/tracker.go (3)

plugins/governance/store.go (1)

GovernanceStore (50-72)

plugins/governance/resolver.go (1)

BudgetResolver (64-67)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

plugins/governance/main.go (3)

plugins/governance/store.go (1)

GovernanceStore (50-72)

plugins/governance/resolver.go (2)

BudgetResolver (64-67)

NewBudgetResolver (70-75)

plugins/governance/tracker.go (2)

UsageTracker (33-45)

NewUsageTracker (52-66)

framework/configstore/store.go (1)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

plugins/governance/resolver.go (4)

plugins/governance/store.go (1)

GovernanceStore (50-72)

framework/configstore/tables/virtualkey.go (2)

TableVirtualKey (95-122)

TableVirtualKey (125-125)

core/schemas/bifrost.go (1)

ModelProvider (32-32)

framework/configstore/tables/ratelimit.go (2)

TableRateLimit (11-36)

TableRateLimit (39-39)

transports/bifrost-http/lib/config.go (2)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

framework/modelcatalog/main.go (2)

ModelCatalog (23-46)

Init (88-138)

transports/bifrost-http/server/server.go (3)

plugins/governance/store.go (1)

GovernanceData (33-39)

plugins/governance/main.go (4)

Config (32-34)

PluginName (21-21)

BaseGovernancePlugin (40-47)

GovernancePlugin (50-69)

transports/bifrost-http/lib/lib.go (1)

EnterpriseOverrides (18-22)

🔇 Additional comments (48)

transports/go.mod (3)

59-79: OpenAPI tooling updates align with HTTP endpoint additions for governance.

The comprehensive go-openapi suite update (swag and subpackages to v0.25.4, analysis to v0.24.2, errors to v0.22.5, validate to v0.25.1) is well-coordinated and aligns with the PR objective of adding HTTP endpoints for budgets and rate-limits. No version conflicts detected across the go-openapi ecosystem.

31-32: Verify AWS SDK v2 compatibility with S3 file operations in governance context.

AWS SDK v2 components are bumped (config v1.32.6, credentials v1.19.6, S3 v1.94.0, SSO v1.30.8). Given the PR's objective to add File APIs support for providers (enabling RAG and file upload workflows), verify that:

S3 v1.94.0 changes don't affect the file upload/management API surface

Credential handling remains secure across the new code paths

Also applies to: 42-42, 44-44

118-120: Verify foundational library version compatibility with concurrency and encryption patterns.

OpenTelemetry (v1.39.0), golang.org/x packages (crypto v0.46.0, net v0.48.0, sync v0.19.0, sys v0.39.0), and google.golang.org packages are bumped. These are foundational for the governance infrastructure's:

Rate-limit enforcement (sync primitives)

Secure configuration handling (crypto)

Distributed tracing of governance decisions (otel)

Ensure integration tests cover concurrent rate-limit scenarios and distributed environments.

Also applies to: 123-133

plugins/jsonparser/go.mod (1)

12-13: Dependency updates are consistently aligned across modules.

All changes are indirect dependency version bumps (minor/patch versions). The updates align with broader ecosystem refreshes across the PR stack.

Also applies to: 23-23, 25-25, 32-33, 37-37, 40-40, 47-47, 50-54

plugins/mocker/go.mod (1)

15-16: Dependency updates are safe and consistent.

Direct dependencies remain unchanged. Indirect updates follow the coordinated refresh pattern across the module stack.

Also applies to: 26-26, 28-28, 35-36, 40-40, 43-43, 50-50, 53-57

plugins/logging/go.mod (2)

42-62: Verify OpenAPI and Weaviate client compatibility with the governance refactor.

The go-openapi library family (analysis, errors, jsonpointer, etc.) is updated to v0.24.x/v0.25.x versions. Weaviate (v1.34.5) and Qdrant (v1.16.2) clients are also bumped. Ensure these updates don't introduce breaking changes in the logging plugin's usage.

Also applies to: 81-82, 88-89

6-6: Remaining dependency updates are consistent and safe.

AWS SDK, Bytedance Sonic, MongoDB driver, OpenTelemetry, and golang.org/x module updates follow the coordinated refresh pattern across the stack with appropriate patch/minor version bumps.

Also applies to: 16-17, 27-27, 29-29, 72-72, 75-75, 94-107

core/go.mod (2)

27-27: Clarify the purpose of the new Masterminds/semver/v3 indirect dependency.

A new indirect dependency on Masterminds/semver/v3 v3.3.1 appears in this file. This likely relates to the governance refactor's versioning logic. Confirm this is necessary and document its usage if not already clear from the implementation.

8-22: Core module dependency updates are appropriately scoped.

AWS SDK v2, Bytedance Sonic, MCP-Go (v0.43.2), fasthttp, OAuth2, and golang.org/x updates follow the stack's dependency refresh. mark3labs/mcp-go bump is significant but appears coordinated with framework updates.

Also applies to: 45-45, 50-50, 52-52, 64-66

plugins/semanticcache/go.mod (2)

44-64: Verify vector database and cache client compatibility.

The semantic cache plugin depends on Weaviate (v1.34.5), Qdrant (v1.16.2), and Redis (v9.17.2) clients. These updates should be validated to ensure no breaking changes in vector operations or cache interactions. Additionally, go-openapi library updates need verification for OpenAPI schema generation.

Also applies to: 90-91, 83-84

18-19: Remaining updates are appropriately coordinated.

AWS SDK, Sonic, MongoDB driver, OpenTelemetry, gRPC, protobuf, and golang.org/x module updates align with stack-wide dependency refresh.

Also applies to: 29-29, 31-31, 38-39, 74-74, 77-77, 94-94, 96-109

plugins/governance/go.mod (3)

5-5: Confirm GORM v1.31.1 compatibility with transactional governance features.

The governance plugin depends on GORM v1.31.1 for persistence. The AI summary mentions "transactional cascade deletions and retry helpers in ConfigStore." Verify that this GORM version supports the transaction patterns and cascade delete logic required by the governance refactor.

44-64: OpenAPI and gRPC updates align with HTTP endpoints feature.

The go-openapi library updates (v0.24.x/v0.25.x) and gRPC/protobuf bumps (v1.77.0/v1.36.11) support the "HTTP endpoints for budgets and rate-limits" mentioned in the PR objective. Ensure these versions are compatible with the governance API surface.

Also applies to: 108-109

17-18: Coordinated dependency updates support governance refactor scope.

AWS SDK, Sonic, compression, MCP-Go, vector database clients, MongoDB, OpenTelemetry, and golang.org/x updates follow the stack's synchronized refresh and appropriately support the governance plugin's expanded responsibilities.

Also applies to: 28-28, 30-30, 37-38, 74-74, 77-77, 83-84, 90-91, 94-94, 96-109

framework/go.mod (2)

8-12: Framework's direct dependency updates are stable and consistent.

Qdrant (v1.16.2), Redis (v9.17.2), Weaviate (v1.34.5), and crypto (v0.46.0) updates align with plugin-level usage. OpenTelemetry and sync module updates are appropriate.

Also applies to: 45-47, 49-49

23-23: Indirect dependency refresh maintains framework stability.

AWS SDK, Sonic, MongoDB, fasthttp, gRPC, and protobuf updates are coordinated across the stack and appropriate for a foundational framework module.

Also applies to: 56-57, 64-64, 70-71, 89-89, 92-92, 103-114

plugins/telemetry/go.mod (2)

9-9: Verify OpenTelemetry v1.39.0 integration with governance observability.

Telemetry plugin directly requires fasthttp (v1.68.0) and indirectly depends on OpenTelemetry components (v1.39.0). The governance refactor likely introduces new metrics/tracing for rate-limits and budgets. Confirm that OTEL v1.39.0 provides the necessary instrumentation APIs.

Also applies to: 100-102

17-18: Telemetry observability updates are appropriately coordinated.

All dependency updates follow the stack's synchronized refresh. AWS SDK, Sonic, go-openapi, data-store clients, and OTEL bumps are aligned and support observability across the governance refactor.

Also applies to: 28-28, 30-30, 38-39, 45-65, 75-75, 78-78, 88-89, 94-95, 98-98, 103-113

framework/configstore/store.go (2)

104-104: LGTM!

The GetRateLimits method addition is consistent with the existing GetBudgets method pattern and completes the public API surface for retrieving governance entities.

146-148: LGTM!

The RetryOnNotFound method is a useful utility for handling eventual consistency scenarios. The signature is flexible with configurable maxRetries and retryDelay parameters, allowing callers to tune retry behavior based on their specific needs.

framework/configstore/rdb.go (6)

1394-1468: Well-structured transactional cascade deletion.

The DeleteVirtualKey implementation properly handles cascade deletions within a transaction:

Fetches the virtual key with preloaded ProviderConfigs

Deletes join table entries for provider config keys

Collects budget/rate limit IDs before deleting provider configs

Deletes provider configs, then budgets and rate limits

Deletes MCP configs

Finally deletes the virtual key and its direct budget/rate limit

The ordering correctly respects foreign key constraints by deleting child records before parent records.

1549-1577: LGTM! Proper cascade deletion for provider config.

The updated DeleteVirtualKeyProviderConfig correctly fetches the record first to capture BudgetID and RateLimitID before deletion, then cleans up the associated budget and rate limit records. This ensures no orphaned governance records remain.

1694-1729: LGTM! Team deletion with proper cleanup.

The transactional DeleteTeam correctly:

Nullifies team_id on associated virtual keys (preventing orphan references)

Stores the budget ID before team deletion

Deletes the team, then the associated budget

This maintains referential integrity while cleaning up related governance data.

1782-1821: LGTM! Customer deletion with comprehensive cleanup.

The transactional DeleteCustomer properly handles:

Nullifying customer_id on associated virtual keys

Nullifying customer_id on associated teams

Deleting the customer and its budget

This prevents orphan references and cleans up governance data atomically.

1823-1830: LGTM! Clean implementation following past review feedback.

The GetRateLimits implementation correctly uses GORM's Find without the unnecessary ErrRecordNotFound check, as addressed from previous review feedback. This is consistent with GetBudgets at line 1889-1895.

2177-2202: LGTM! Well-implemented retry logic with proper context handling.

The RetryOnNotFound implementation:

Correctly checks for both ErrNotFound and gorm.ErrRecordNotFound

Properly handles context cancellation via select with ctx.Done()

Avoids sleeping after the last attempt

Returns the last error when all retries are exhausted

The use of range maxRetries is a clean Go 1.22+ idiom for iteration.

transports/bifrost-http/handlers/middlewares.go (1)

49-73: LGTM! Clean enterprise override integration.

The TransportInterceptorMiddleware changes properly:

Accept enterpriseOverrides as a parameter for flexibility

Add a nil-guard at lines 58-61 to safely skip interception when no enterprise overrides are provided (addressing prior review feedback)

Use enterpriseOverrides.GetGovernancePluginName() to decouple from the governance plugin package

This enables pluggable governance implementations while maintaining backward compatibility.

transports/bifrost-http/handlers/governance.go (8)

24-26: LGTM!

The GetGovernanceData() method addition to the GovernanceManager interface enables in-memory data retrieval for governance entities, supporting the new from_memory query parameter feature.

41-44: LGTM! Proper nil validation.

Adding the nil-check for manager ensures the handler cannot be created in an invalid state, failing fast during initialization rather than at runtime.

178-181: LGTM! New GET endpoints for governance entities.

The new /api/governance/budgets and /api/governance/rate-limits endpoints complete the read API surface for governance entities, consistent with the existing CRUD patterns.

188-201: LGTM! Comprehensive from_memory support with proper nil-checks.

All from_memory handlers now include nil-checks for GetGovernanceData() and return appropriate 500 errors when governance data is unavailable. This addresses the concerns from previous reviews about potential nil pointer dereferences.

Also applies to: 420-440, 948-976, 1057-1076, 1207-1222, 1299-1318, 1447-1501

196-199: Note: Response format differs between from_memory and DB paths.

The from_memory path returns data.VirtualKeys which is a map[string]*TableVirtualKey, while the DB path returns a slice []TableVirtualKey. This is a conscious design choice for efficient in-memory lookups vs. ordered DB results, but API consumers should be aware of this difference.

Consider documenting this behavioral difference in API documentation or ensuring consumers handle both formats appropriately.

310-329: LGTM! Proper key validation and association.

The key handling for provider configs correctly:

Fetches keys by IDs when specified

Validates that all requested keys were found

Associates keys with the provider config

The error messages are descriptive and include the provider name for debugging context.

575-577: Good addition of rate limit validation.

Adding validateRateLimit calls during updates ensures consistency with the create path and prevents invalid rate limit configurations from being persisted.

Also applies to: 788-790

695-719: LGTM! Thoughtful key update semantics.

The comment at lines 706-709 clearly documents the update behavior:

nil KeyIDs → leave existing keys unchanged

Empty slice → clear all keys

Values present → update to new keys

This follows the standard PATCH semantics where omitted fields preserve existing values.

transports/bifrost-http/lib/config_test.go (2)

624-633: No-op testLogger implementation is appropriate for these tests

The testLogger methods are all no‑ops, which is ideal for these SQLite integration tests: it satisfies the logger interface without polluting test output or affecting behavior (including Fatal, which should not terminate tests).

2199-2207: Hash uniqueness and ordering-stability tests are thorough and well-structured

The additional assertions around:

unique hashes for different optional-field combinations (hashes map in TestProviderHashComparison_OptionalFieldsPresence),

key hash stability across Models permutations,

virtual key hash stability across ProviderConfigs, AllowedModels, KeyIDs, MCPConfigs, and ToolsToExecute permutations,

give strong coverage that the hashing logic is both sensitive to real config changes and insensitive to ordering artifacts. These tests should catch most subtle regressions in future refactors.

Also applies to: 9528-9579, 9583-9712

plugins/governance/tracker.go (1)

221-255: Selective rate‑limit reset persistence looks correct

The startup reset path only updates token_current_usage, token_last_reset, request_current_usage, and request_last_reset inside a transaction, explicitly skipping hooks and leaving *_max_limit and *_reset_duration untouched. This matches the intended “fix usage only, don’t trample config” semantics and should be safe with the virtual LastDB*Usage fields.

plugins/governance/resolver.go (2)

63-75: Resolver wiring to GovernanceStore looks consistent

Using GovernanceStore (non‑pointer) in BudgetResolver and delegating from EvaluateRequest into checkRateLimitHierarchy / checkBudgetHierarchy keeps the resolver cleanly aligned with the new interface; virtual‑key existence, provider/model filters, and decision mapping remain intact.

Also applies to: 129-137

220-233: Atomic budget check via GovernanceStore is aligned with new contract

checkBudgetHierarchy’s use of store.CheckBudget(ctx, vk, request, nil) and mapping any non‑nil error to DecisionBudgetExceeded matches the documented “error == violation” semantics and keeps budget logic centralized in the store.

plugins/governance/main.go (1)

40-47: BaseGovernancePlugin + GetGovernanceStore provide a clean introspection contract

Exposing GetGovernanceStore() via BaseGovernancePlugin and having GovernancePlugin implement it gives the HTTP layer a stable way to reach governance state without overloading the generic plugin interface. The GetGovernanceStore implementation simply returns the wired GovernanceStore, which is appropriate.

Also applies to: 673-675

transports/bifrost-http/server/server.go (4)

201-210: GovernanceInMemoryStore locking pattern is sound for read‑only access

Switching GovernanceInMemoryStore to hold *lib.Config and using Config.Mu.RLock around Config.Providers reads gives a clear, low‑overhead pattern for the governance plugin to inspect configured providers without copying.

213-322: EnterpriseOverrides‑aware LoadPlugin/LoadPlugins design looks consistent

Plumbing EnterpriseOverrides through LoadPlugin/LoadPlugins and using EnterpriseOverrides.GetGovernancePluginName() in the governance case lets OSS and enterprise builds agree on which plugin name is “the governance plugin” while still sharing the same loading logic. The governance case wiring (constructing GovernanceInMemoryStore and calling governance.Init) matches the plugin’s new constructor.

519-550: RetryOnNotFound + type assertion in ReloadVirtualKey

Wrapping GetVirtualKey with RetryOnNotFound and the shared DBLookupMaxRetries/DBLookupDelay gives the reload path some resilience to transient “not found” races right after creation. The subsequent type assertion to *tables.TableVirtualKey and in‑memory update via UpdateVirtualKeyInMemory are straightforward.

Just ensure RetryOnNotFound only retries on ErrNotFound and returns nil, nil or a distinct error when the VK truly doesn’t exist so this path doesn’t mask persistent failures.

1127-1166: EnterpriseOverrides propagation through Bootstrap and middleware looks good

Passing s into lib.LoadConfig and LoadPlugins, and then threading it into RegisterAPIRoutes and TransportInterceptorMiddleware(s.Config, s), cleanly wires the new EnterpriseOverrides hooks (governance plugin and pricing manager) into server startup without changing the public server API.

Also applies to: 1225-1235

transports/bifrost-http/lib/config.go (3)

47-50: DB lookup retry constants are fine as shared defaults

Defining DBLookupMaxRetries = 5 and DBLookupDelay = 1 * time.Second here provides a clear, centrally documented default for RetryOnNotFound callers and keeps magic numbers out of the server code.

261-299: LoadConfig’s EnterpriseOverrides plumbing preserves existing behavior

Extending LoadConfig and loadConfigFromFile to accept an EnterpriseOverrides argument, while forwarding it only into the framework/pricing initialization paths, keeps all other config loading semantics unchanged. Callers that don’t care about enterprise overrides can safely pass nil.

1505-1546: EnterpriseOverrides‑aware pricing manager initialization is well‑factored

Both initFrameworkConfigFromFile and initDefaultFrameworkConfig now:

build a modelcatalog.Config from store/file/defaults,

set config.FrameworkConfig.Pricing from it, and

prefer EnterpriseOverrides.LoadPricingManager when provided, falling back to modelcatalog.Init otherwise.

This cleanly opens an extension point for custom pricing managers without regressing OSS behavior; failures are logged and leave config.PricingManager nil, which existing call sites already tolerate.

Please ensure any EnterpriseOverrides implementation you add respects the same contract as modelcatalog.Init (idempotent initialization, safe to call once at bootstrap, and non‑blocking on hot paths).

Also applies to: 1866-1927

transports/bifrost-http/lib/config_test.go

transports/bifrost-http/server/server.go

akshaydeo · 2025-12-24T12:11:27Z

Merge activity

Dec 24, 12:11 PM UTC: A user started a stack merge that includes this pull request via Graphite.
Dec 24, 12:11 PM UTC: @akshaydeo merged this pull request with Graphite.

feat: governance plugin refactor

270cc82

Pratham-Mishra04 requested review from akshaydeo and danpiths December 22, 2025 10:22

This was referenced Dec 22, 2025

chore: governance tests added #1041

Merged

feat: add Nix development environment with flake support #1119

Open

feat: add support for per-model and per-provider level budgeting and rate limiting in governance plugin #1120

Open

Pratham-Mishra04 mentioned this pull request Dec 22, 2025

chore: adds test for governance plugin's model configs and providers budgeting and rate limiting #1121

Open

17 tasks

coderabbitai bot requested changes Dec 22, 2025

View reviewed changes

Pratham-Mishra04 force-pushed the 12-08-feat_governance_plugin_refactors branch from bc1bdce to 5ca488c Compare December 22, 2025 11:14

coderabbitai bot requested changes Dec 22, 2025

View reviewed changes

transports/bifrost-http/handlers/governance.go Show resolved Hide resolved

transports/bifrost-http/handlers/governance.go Show resolved Hide resolved

Pratham-Mishra04 force-pushed the 12-08-feat_governance_plugin_refactors branch from 5ca488c to e3f7191 Compare December 22, 2025 11:30

coderabbitai bot requested changes Dec 22, 2025

View reviewed changes

framework/configstore/rdb.go Show resolved Hide resolved

Pratham-Mishra04 force-pushed the 12-08-feat_governance_plugin_refactors branch from e3f7191 to 7ae83e4 Compare December 22, 2025 12:13

coderabbitai bot requested changes Dec 22, 2025

View reviewed changes

transports/go.mod Show resolved Hide resolved

coderabbitai bot approved these changes Dec 22, 2025

View reviewed changes

Pratham-Mishra04 force-pushed the 12-08-feat_governance_plugin_refactors branch from 7ae83e4 to 864eae5 Compare December 22, 2025 15:28

coderabbitai bot requested changes Dec 22, 2025

View reviewed changes

plugins/governance/store.go Show resolved Hide resolved

coderabbitai bot approved these changes Dec 22, 2025

View reviewed changes

deadlock fix

d877912

Pratham-Mishra04 force-pushed the 12-08-feat_governance_plugin_refactors branch from 864eae5 to d877912 Compare December 23, 2025 14:42

coderabbitai bot requested changes Dec 23, 2025

View reviewed changes

transports/bifrost-http/lib/config_test.go Show resolved Hide resolved

transports/bifrost-http/server/server.go Show resolved Hide resolved

This was referenced Dec 23, 2025

feat: adds backend go functions for model/provider governance configuration #1169

Open

feat: adds governance frontend for models and providers #1170

Open

coderabbitai bot approved these changes Dec 24, 2025

View reviewed changes

akshaydeo merged commit 580f591 into v1.4.0 Dec 24, 2025
2 of 3 checks passed

akshaydeo deleted the 12-08-feat_governance_plugin_refactors branch December 24, 2025 12:11

feat: governance plugin refactor #1154

feat: governance plugin refactor #1154

Conversation

Pratham-Mishra04 commented Dec 22, 2025

Summary

Changes

Type of change

Affected areas

How to test

Screenshots/Recordings

Breaking changes

Related issues

Security considerations

Checklist

Uh oh!

Pratham-Mishra04 commented Dec 22, 2025 • edited by danpiths Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai bot commented Dec 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

Poem

Pre-merge checks and finishing touches

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

akshaydeo commented Dec 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge activity

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Pratham-Mishra04 commented Dec 22, 2025 •

edited by danpiths

Loading

coderabbitai bot commented Dec 22, 2025 •

edited

Loading

akshaydeo commented Dec 24, 2025 •

edited

Loading