feat: real-time governance updates via WebSocket and base model selection in UI#1532
Conversation
|
Caution Review failedThe pull request is closed. 📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds real-time handling for Changes
Sequence Diagram(s)sequenceDiagram
participant UI as UI Component
participant WS as WebSocket
participant RTK as RTK Query Cache
participant API as Backend API
UI->>RTK: Register governance queries (modelConfigs, providers, virtualKeys)
Note over UI,RTK: Explicit polling options removed
rect rgba(0,128,0,0.5)
WS->>RTK: governance_update (budgets & rate_limits)
RTK->>RTK: Merge diffs into getModelConfigs cache
RTK->>RTK: Merge diffs into getProviderGovernance cache
RTK->>RTK: Merge diffs into getVirtualKeys cache (and provider_configs)
RTK->>UI: Notify subscribers (components re-render)
end
rect rgba(0,0,128,0.5)
UI->>API: useLazyGetBaseModelsQuery -> GET /models/base?query=...
API-->>UI: List of base model names
UI->>RTK: Populate/select options / re-render
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
danpiths
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
6fecb69 to
499cf07
Compare
danpiths
force-pushed
the
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
branch
from
23aba08 to
86828b2
Compare
danpiths
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
499cf07 to
cb53486
Compare
danpiths
force-pushed
the
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
branch
from
86828b2 to
8210c71
Compare
danpiths
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
cb53486 to
7beaffe
Compare
danpiths
force-pushed
the
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
branch
from
8210c71 to
b72c861
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@ui/components/ui/modelMultiselect.tsx`:
- Around line 159-176: The refresh when provider is empty uses a lower limit (5)
causing the list to shrink; update the branch that calls getModels when
shouldLoadOnEmpty is true to use the same limit as initial load (20) by changing
the limit expression in that call (the getModels invocation inside the
shouldLoadOnEmpty branch) to use currentQuery ? 20 : 20 (i.e. 20) so it matches
the other paths (also keep getBaseModels unchanged); check the surrounding logic
in the handleChange/getModels/getBaseModels flow to ensure consistency with
provider, currentQuery, keys, and shouldUseBaseModels variables.
7beaffe to
7e505cd
Compare
danpiths
force-pushed
the
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
branch
from
b72c861 to
fc2a474
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@ui/components/sidebar.tsx`:
- Around line 427-440: The "Budgets & Limits" menu item is only RBAC-gated via
hasGovernanceAccess but also needs to be gated by the governance feature flag
(core_config.enable_governance) so it disappears when governance is disabled;
update SidebarItemView (where item.subItems are mapped) to filter out the
subitem with title "Budgets & Limits" (or url "/workspace/model-limits") when
isGovernanceEnabled (or core_config.enable_governance) is false — either pass
isGovernanceEnabled into SidebarItemView and filter item.subItems before
mapping, or apply the check inline during the map so that the Budgets & Limits
entry is not rendered unless both isGovernanceEnabled and hasGovernanceAccess
are true.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@ui/hooks/useStoreSync.tsx`:
- Around line 61-65: The code incorrectly iterates with
Object.values(draft.virtual_keys) even though draft.virtual_keys is a
VirtualKey[] per GetVirtualKeysResponse; change the loop in the updateQueryData
callback to iterate directly over the array (for (const vk of
draft.virtual_keys) { ... }) and update the preceding comment in
useStoreSync.tsx to describe that virtual_keys is an array (not a
map/from_memory) so the logic merges array elements into the virtual keys cache
correctly; keep the rest of the updateQueryData and dispatch usage unchanged.
🧹 Nitpick comments (1)
ui/components/ui/modelMultiselect.tsx (1)
171-176: Inconsistent limit inshouldLoadOnEmptyrefresh path.When
shouldLoadOnEmptyis true (but not base_models mode), the refresh useslimit: 5without a query (Line 175), while the initial load at Line 93 useslimit: 20. This can shrink the options list unexpectedly after selection.Suggested fix for consistency
} else if (shouldLoadOnEmpty) { getModels({ query: currentQuery || undefined, keys: keys && keys.length > 0 ? keys : undefined, - limit: currentQuery ? 20 : 5, + limit: 20, }); }
danpiths
force-pushed
the
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
branch
from
fc2a474 to
fed10d4
Compare
danpiths
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
7e505cd to
f9286c2
Compare
danpiths
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
f9286c2 to
96bf11b
Compare
danpiths
force-pushed
the
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
branch
from
fed10d4 to
74e5e7e
Compare
danpiths
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
96bf11b to
663bfd1
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@ui/components/sidebar.tsx`:
- Around line 431-444: The "Routing Rules" subitem isn't currently being hidden
when governance is disabled; update the subitem filtering logic in the sidebar
rendering (the code that inspects subItem.url) to treat
"/workspace/routing-rules" the same as "/workspace/model-limits": add a
conditional that returns/nulls the subitem when ((subItem.url ===
"/workspace/model-limits" || subItem.url === "/workspace/routing-rules") &&
!isGovernanceEnabled). Ensure you modify the branch that already checks
isGovernanceEnabled (referencing isGovernanceEnabled and subItem.url) so both
entries are gated consistently.
Pratham-Mishra04
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
663bfd1 to
d98d5dc
Compare
Pratham-Mishra04
force-pushed
the
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
branch
from
74e5e7e to
b4f2e3a
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@ui/app/workspace/providers/views/modelProviderConfig.tsx`:
- Around line 59-60: The code reads coreConfig.client_config.enable_governance
directly which can throw if client_config is undefined; update the
isGovernanceEnabled assignment (the variable isGovernanceEnabled that uses
useGetCoreConfigQuery) to use optional chaining like
coreConfig?.client_config?.enable_governance and fall back to false to mirror
the pattern used elsewhere (e.g., providerGovernanceTable.tsx) so it safely
handles missing client_config.
🧹 Nitpick comments (2)
core/changelog.md (1)
2-2: Clarify relationship between model namespace fix and base model selection.The model namespace preservation fix seems related to the base model selection feature mentioned in the PR summary. If preserving namespaces like
meta-llama/Llama-3.1-8Bwas a prerequisite for base model selection to work correctly, consider expanding this changelog entry to make that relationship explicit, e.g.:- fix: model names with namespaces (e.g., `meta-llama/Llama-3.1-8B`) are now correctly preserved instead of being incorrectly split as provider-prefixed models, enabling proper base model selectionThis would help readers understand the context and impact of the fix.
ui/components/ui/modelMultiselect.tsx (1)
166-170: Simplify redundant ternary expression.The expression
currentQuery ? 20 : 20always evaluates to20, making the ternary unnecessary.Suggested fix
} else if (shouldUseBaseModels) { getBaseModels({ query: currentQuery || undefined, - limit: currentQuery ? 20 : 20, + limit: 20, });
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx (1)
117-127:⚠️ Potential issue | 🟡 MinorConsider adding a null guard for the edit sheet similar to the detail sheet.
The detail sheet (line 127) correctly guards against a null
selectedVirtualKey, which handles the case where the item is deleted while the sheet is open. However,VirtualKeySheetdoesn't have an equivalent guard foreditingVirtualKeyin edit mode.If a virtual key is deleted via a WebSocket update while the edit sheet is open,
editingVirtualKeybecomesnull, and the sheet would unexpectedly switch to "add new" behavior.🛡️ Suggested fix for consistency
- {showVirtualKeySheet && ( + {showVirtualKeySheet && (editingVirtualKeyId === null || editingVirtualKey !== null) && ( <VirtualKeySheet virtualKey={editingVirtualKey} teams={teams} customers={customers} onSave={handleVirtualKeySaved} onCancel={() => setShowVirtualKeySheet(false)} /> )}This ensures the sheet only renders when either adding a new key (
editingVirtualKeyId === null) or when the editing key still exists in the cache.
🧹 Nitpick comments (3)
ui/app/workspace/providers/views/modelProviderConfig.tsx (1)
59-60: Consider aligning config query approach withproviderGovernanceTable.tsx.This component uses
useGetCoreConfigQuery({})withoutfromDB, whileproviderGovernanceTable.tsxusesuseLazyGetCoreConfigQuerywith{ fromDB: true }. This could lead to transient UI inconsistencies if the cached config differs from the DB value—the tab might briefly appear or disappear before the child component's own check resolves.If the intent is to use memory/cached data per the PR's
from_memory=truemigration, consider applying the same approach inproviderGovernanceTable.tsxfor consistency across the governance feature surface.ui/hooks/useStoreSync.tsx (1)
27-29: Consider handling partial data more defensively.If
datais undefined or null, accessingdata.datawill throw. Consider adding a guard:const unsubGovUpdate = subscribe("governance_update", (data) => { - const { budgets, rate_limits } = data.data || {}; + const { budgets, rate_limits } = data?.data || {}; if (!budgets && !rate_limits) return;ui/components/ui/modelMultiselect.tsx (1)
171-177: Inconsistent limit between initial load and refresh.When
shouldLoadOnEmptyis true (but not base_models mode), the initial load at Line 93 useslimit: 20, but the refresh after selection at Line 175 useslimit: currentQuery ? 20 : 5. This can cause the dropdown list to shrink after a selection.Suggested fix
} else if (shouldLoadOnEmpty) { getModels({ query: currentQuery || undefined, keys: keys && keys.length > 0 ? keys : undefined, - limit: currentQuery ? 20 : 5, + limit: 20, }); }
danpiths
force-pushed
the
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
branch
from
b4f2e3a to
8236164
Compare
danpiths
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
d98d5dc to
5543680
Compare
danpiths
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
5543680 to
9477348
Compare
danpiths
force-pushed
the
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
branch
from
8236164 to
75edfa0
Compare
Merge activity
|
Pratham-Mishra04
changed the base branch from
feature/02-03-feat_add_from_memory_governance_endpoints_websocket_event_broadcasting_and_base_models_api
to
graphite-base/1532
Pratham-Mishra04
force-pushed
the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
from
9477348 to
cc685c9
Compare
Pratham-Mishra04
deleted the
feature/02-03-feat_real-time_governance_updates_via_websocket_and_base_model_selection_in_ui
branch
2 participants
Summary
The frontend now receives real-time budget and rate-limit updates via WebSocket
instead of polling every 10 seconds. When a governance mutation occurs (e.g.
budget usage increases after an LLM call), the backend emits a minimal diff
(
{ budgets: { "id": {...} } }or{ rate_limits: { "id": {...} } }). Thefrontend merges this diff into the RTK Query cache using Immer draft mutations,
updating all affected views (model limits, provider governance, virtual keys)
immediately.
Also switches governance queries to
from_memory=truefor faster reads, removesall 10-second polling intervals, adds base model selection support in the model
limits form, and fixes stale data in edit/detail sheets.
Re-enable sidebar navigation — restore "Budgets & Limits" and "Routing
Rules" menu items in the sidebar that were temporarily hidden, and restore the
Networkicon import.Disabled state UX for governance-dependent items — when governance is
disabled, show governance-dependent sidebar items and provider config tabs as
disabled with tooltips explaining how to enable them, instead of hiding them
completely.
Changes
useStoreSync.tsx: Subscribe togovernance_updateWebSocket events andmerge budget/rate-limit diffs into
getModelConfigs,getProviderGovernance,and
getVirtualKeysRTK Query caches using Immer draft mutations (mutatedraft, don't return)
governanceApi.ts: SwitchgetVirtualKeys,getModelConfigs, andgetProviderGovernancequeries to usefrom_memory=truequery parameterbaseApi.ts: AddBaseModelstag typeprovidersApi.ts: AddgetBaseModelsquery endpoint forGET /api/models/basemodelMultiselect.tsx: SupportloadModelsOnEmptyProvider="base_models"mode — when no provider is selected, loads distinct base model names instead
of all models from all providers. Adds
useLazyGetBaseModelsQueryintegrationwith search/filter support.
modelLimitSheet.tsx: UseloadModelsOnEmptyProvider="base_models"forthe model selector (governance configs should use canonical base model names)
modelLimitsTable.tsx: StoreeditingModelConfigId(string) instead offull
ModelConfigobject in state; derive current object viauseMemofromprops so it stays reactive to RTK cache updates
modelLimitsView.tsx: RemovepollingInterval: 10000andskipPollingIfUnfocusedgovernanceFormFragment.tsx: RemovepollingInterval: 10000andskipPollingIfUnfocusedproviderGovernanceTable.tsx: RemovepollingInterval: 10000andskipPollingIfUnfocusedvirtualKeysTable.tsx: StoreeditingVirtualKeyIdandselectedVirtualKeyId(strings) instead of full objects; derive viauseMemofrom props for live updates
/workspace/model-limits)/workspace/routing-rules)Networkicon import fromlucide-reactsidebar.tsx: Add disabled state styling for governance-dependent sidebaritems (Budgets & Limits, Routing Rules) when governance is disabled — show
items with muted styling, not-allowed cursor, and tooltip explaining how to
enable governance. Remove
pointer-events-noneto allow hover effects.modelProviderConfig.tsx: Show "Governance" tab in provider configurationeven when governance is disabled, but display it as disabled with tooltip.
Previously the tab was completely hidden.
Type of change
Affected areas
How to test
Manual testing:
Real-time budget updates: Open the Model Limits page. Make an LLM request
through Bifrost. Observe that "Current Usage" values update immediately
without page refresh.
Real-time rate limit updates: Open the Provider Governance page for a
provider with rate limits. Make an LLM request. Observe token/request usage
bars update in real-time.
Virtual keys live updates: Open a virtual key's detail sheet. Make
requests using that virtual key. Budget and rate limit displays should update
live.
Edit sheet live data: Open the edit sheet for a model config. While the
sheet is open, make LLM requests that affect that model's budget. The
"Current Usage" in the edit sheet should reflect the latest value.
Base model selection: Go to Model Limits → Add Model Limit. Leave the
provider dropdown empty. The model selector should show distinct base model
names (e.g.
gpt-4o,claude-3-5-sonnet) instead of provider-specificvariants.
No polling: Open browser DevTools Network tab. Verify there are no
repeated GET requests to
/api/governance/*every 10 seconds. Updates shouldonly arrive via WebSocket.
Disabled governance sidebar items: Disable governance in Config →
Governance. Verify that "Budgets & Limits" and "Routing Rules" sidebar items
appear with muted/disabled styling. Hover over them to see tooltip explaining
how to enable. Clicking should do nothing.
Disabled governance tab in provider config: With governance disabled, go
to Model Providers → select a provider. In "Provider level configuration",
verify the "Governance" tab is visible but disabled/grayed out with a tooltip.
Screenshots/Recordings
N/A
Breaking changes
Related issues
N/A
Security considerations
No security implications. The WebSocket subscription consumes the same
governance data already accessible through authenticated GET endpoints. No new
data is exposed.
Checklist
docs/contributing/README.mdand followed the guidelines