feat: assign virtual keys on teams page directly#2757
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughTeam queries now include a correlated subquery providing Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant API as API Server
participant Store as ConfigStore (GORM)
participant DB as Database
participant Plugin as GovernancePlugin
Client->>API: Request teams / governance
API->>Store: Query teams (Select includes correlated COUNT of virtual_keys)
Store->>DB: SELECT teams + correlated subquery AS virtual_key_count
DB-->>Store: Rows with virtual_key_count
Store-->>API: Teams (preloaded relations, virtual_key_count)
API->>Plugin: Build governance snapshot (teams, virtualKeys, customers)
Plugin->>Plugin: reset team.VirtualKeyCount = 0
Plugin->>Plugin: iterate virtualKeys -> increment matching team.VirtualKeyCount
Plugin->>Plugin: wire teams into customers (now with updated counts)
Plugin-->>API: governance data (teams with virtual_key_count)
API-->>Client: Response
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 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 docstrings
🧪 Generate unit tests (beta)
Comment |
Confidence Score: 5/5Safe to merge; the only finding is a documentation gap in the governance schema, not a runtime defect. All code changes (GORM tag, correlated subquery, in-memory store reset/recompute) are correct. The sole issue is that docs/openapi/schemas/management/governance.yaml and docs/openapi/openapi.json — the canonical Team schema needs the virtual_key_count field added. Important Files Changed
Reviews (6): Last reviewed commit: "feat: assign virtual keys on teams page ..." | Re-trigger Greptile |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@framework/configstore/rdb.go`:
- Around line 2525-2527: The in-memory snapshot returned by
governanceManager.GetGovernanceData() / GetGovernanceConfig still leaves
Team.VirtualKeyCount at 0 because teams are loaded with a plain Find(&teams);
update the snapshot build to populate VirtualKeyCount by running the same
subquery used in the DB path (SELECT COUNT(*) FROM governance_virtual_keys WHERE
team_id = governance_teams.id) or an equivalent aggregate query and assign the
result into each Team.VirtualKeyCount, referencing the Team struct field
VirtualKeyCount and the governance_virtual_keys table; alternatively, if you
prefer not to compute counts for in-memory snapshots, ensure the handlers omit
or nil out the VirtualKeyCount field when it was not computed so both paths
return consistent responses.
In `@ui/lib/types/governance.ts`:
- Around line 177-178: The TypeScript fields team_id and customer_id currently
allow null (team_id?: string | null; customer_id?: string | null), but sending
null will be indistinguishable from omitted on the Go side (the handler checks
if req.TeamID != nil), so change both fields to optional non-nullable strings
(team_id?: string; customer_id?: string;) so clients omit the field to leave
associations unchanged; if you need an explicit "clear" operation instead, add a
separate boolean flag or explicit clear token in the contract and update the
corresponding request types and client code to use that signal.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a08dcbce-5e4d-47fd-8210-35938771417e
📒 Files selected for processing (3)
framework/configstore/rdb.goframework/configstore/tables/team.goui/lib/types/governance.ts
f80256a to
4d4b7f8
Compare
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docs/openapi/schemas/management/governance.yaml (1)
361-438:⚠️ Potential issue | 🟠 MajorAdd
virtual_key_counttoTeamschema to match runtime responses.
Teamresponses now include a computedvirtual_key_count, but this schema still omits that property. This causes API contract drift for docs/codegen/validation.📘 Proposed schema patch
Team: type: object description: Team configuration properties: @@ virtual_keys: type: array nullable: true description: > Virtual keys assigned to this team. This field may be omitted or returned as null in some responses (for example, when a team is embedded inside a virtual-key response) to avoid nested `virtual_keys` recursion. items: $ref: '#/VirtualKey' + virtual_key_count: + type: integer + format: int64 + description: Number of virtual keys assigned to this team.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@docs/openapi/schemas/management/governance.yaml` around lines 361 - 438, The Team schema is missing the runtime-only computed property virtual_key_count; update the Team object definition to include a new property named virtual_key_count (type: integer, nullable: true or integer according to runtime behavior) alongside virtual_keys so API docs/codegen/validation match responses, and ensure CreateTeamRequest/UpdateTeamRequest/ListTeamsResponse need no changes; modify the Team schema block (properties: ...) to add virtual_key_count with an appropriate description like "Computed number of virtual keys for this team".
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@docs/openapi/schemas/management/governance.yaml`:
- Around line 361-438: The Team schema is missing the runtime-only computed
property virtual_key_count; update the Team object definition to include a new
property named virtual_key_count (type: integer, nullable: true or integer
according to runtime behavior) alongside virtual_keys so API
docs/codegen/validation match responses, and ensure
CreateTeamRequest/UpdateTeamRequest/ListTeamsResponse need no changes; modify
the Team schema block (properties: ...) to add virtual_key_count with an
appropriate description like "Computed number of virtual keys for this team".
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0c13f59f-471d-4d0b-bab0-b7e5ba8fd336
📒 Files selected for processing (6)
docs/openapi/openapi.jsondocs/openapi/schemas/management/governance.yamldocs/openapi/schemas/management/users.yamlframework/configstore/rdb.goframework/configstore/tables/team.goui/lib/types/governance.ts
✅ Files skipped from review due to trivial changes (1)
- docs/openapi/schemas/management/users.yaml
🚧 Files skipped from review as they are similar to previous changes (2)
- ui/lib/types/governance.ts
- framework/configstore/rdb.go
4d4b7f8 to
c5f134d
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
framework/configstore/rdb.go (1)
2525-2527: Extract thevirtual_key_countprojection into a shared team-query scope.The same
Select("governance_teams.*, (SELECT COUNT(*) ...) AS virtual_key_count")now exists in four places. Keeping that SQL inline everywhere makes it easy for one path to drift again the next time team loading changes.♻️ Suggested cleanup
+func withTeamVirtualKeyCount(db *gorm.DB) *gorm.DB { + return db.Select(` + governance_teams.*, + (SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS virtual_key_count + `) +} + func (s *RDBConfigStore) GetTeams(ctx context.Context, customerID string) ([]tables.TableTeam, error) { - query := s.db.WithContext(ctx). - Select("governance_teams.*, (SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS virtual_key_count"). + query := withTeamVirtualKeyCount(s.db.WithContext(ctx)). Preload("Customer").Preload("Budget").Preload("RateLimit")- if err := baseQuery. - Select("governance_teams.*, (SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS virtual_key_count"). + if err := withTeamVirtualKeyCount(baseQuery). Preload("Customer").Preload("Budget").Preload("RateLimit").- if err := s.db.WithContext(ctx). - Select("governance_teams.*, (SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS virtual_key_count"). + if err := withTeamVirtualKeyCount(s.db.WithContext(ctx)). Preload("Customer").Preload("Budget").Preload("RateLimit").- if err := s.db.WithContext(ctx). - Select("governance_teams.*, (SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS virtual_key_count"). + if err := withTeamVirtualKeyCount(s.db.WithContext(ctx)). Find(&teams).Error; err != nil {As per coding guidelines, "
**: 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)".Also applies to: 2569-2570, 2583-2586, 3434-3436
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` around lines 2525 - 2527, Extract the repeated Select projection "(SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS virtual_key_count" into a reusable GORM scope (e.g., WithVirtualKeyCount or TeamWithVirtualKeyCount) and replace the inline Select in every place where query := s.db.WithContext(ctx).Select("governance_teams.*, (SELECT COUNT(*) ... ) AS virtual_key_count") is used (including the other occurrences noted). Implement the scope so it returns a func(*gorm.DB) *gorm.DB that applies Select("governance_teams.*, (SELECT COUNT(*) ... ) AS virtual_key_count"), and update callers (the variables named query and any usages of Preload("Customer").Preload("Budget").Preload("RateLimit")) to call .Scopes(WithVirtualKeyCount) instead of duplicating the SQL.plugins/governance/store.go (1)
286-355: Add a focused regression test for count/wiring order.Please add a unit/integration test that asserts
GetGovernanceData()returns identicalVirtualKeyCountfor bothTeams[id]and the correspondingcustomer.Teams[]entry after VK create/update/delete flows. This will lock in the ordering contract introduced here.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@plugins/governance/store.go` around lines 286 - 355, Add a regression test that exercises VK create/update/delete flows and then calls GetGovernanceData(), asserting that the VirtualKeyCount on the top-level teams map (teams[id]) equals the VirtualKeyCount on the corresponding team object embedded inside its customer (customer.Teams[]). Locate the call path that populates teams/customers and wires VirtualKeyCount (the loop over virtualKeys and the subsequent teams→customers wiring in GetGovernanceData or equivalent method), simulate VK changes, invoke GetGovernanceData(), and assert equality for each matching team id to lock in the ordering contract.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/openapi/openapi.json`:
- Around line 33457-33467: The OpenAPI spec is missing a security scheme for the
x-goog-api-key header used by the /api/governance/virtual-keys/quota handler;
add a GoogleApiKeyAuth (or GoogleApiKey) security scheme defining an apiKey in:
header with name: x-goog-api-key under components.securitySchemes (or
securityDefinitions if this spec is v2), and then include a reference to that
scheme (e.g., "GoogleApiKeyAuth": []) in the quota operation's "security" array
alongside VirtualKeyAuth/BearerAuth/ApiKeyAuth so clients and docs reflect the
handler's accepted auth.
---
Nitpick comments:
In `@framework/configstore/rdb.go`:
- Around line 2525-2527: Extract the repeated Select projection "(SELECT
COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS
virtual_key_count" into a reusable GORM scope (e.g., WithVirtualKeyCount or
TeamWithVirtualKeyCount) and replace the inline Select in every place where
query := s.db.WithContext(ctx).Select("governance_teams.*, (SELECT COUNT(*) ...
) AS virtual_key_count") is used (including the other occurrences noted).
Implement the scope so it returns a func(*gorm.DB) *gorm.DB that applies
Select("governance_teams.*, (SELECT COUNT(*) ... ) AS virtual_key_count"), and
update callers (the variables named query and any usages of
Preload("Customer").Preload("Budget").Preload("RateLimit")) to call
.Scopes(WithVirtualKeyCount) instead of duplicating the SQL.
In `@plugins/governance/store.go`:
- Around line 286-355: Add a regression test that exercises VK
create/update/delete flows and then calls GetGovernanceData(), asserting that
the VirtualKeyCount on the top-level teams map (teams[id]) equals the
VirtualKeyCount on the corresponding team object embedded inside its customer
(customer.Teams[]). Locate the call path that populates teams/customers and
wires VirtualKeyCount (the loop over virtualKeys and the subsequent
teams→customers wiring in GetGovernanceData or equivalent method), simulate VK
changes, invoke GetGovernanceData(), and assert equality for each matching team
id to lock in the ordering contract.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e1e40c1e-96e3-49ef-983e-c6e2e61dbf81
📒 Files selected for processing (6)
docs/openapi/openapi.jsondocs/openapi/schemas/management/governance.yamldocs/openapi/schemas/management/users.yamlframework/configstore/rdb.goframework/configstore/tables/team.goplugins/governance/store.go
✅ Files skipped from review due to trivial changes (1)
- docs/openapi/schemas/management/users.yaml
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/openapi/schemas/management/governance.yaml
- framework/configstore/tables/team.go
c5f134d to
924e214
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
framework/configstore/rdb.go (1)
4502-4502: Minor: Consider extracting the repeated SELECT clause.The correlated subquery for
virtual_key_countis duplicated across 4 methods (GetTeams,GetTeamsPaginated,GetTeam,GetGovernanceConfig). Consider extracting it to a constant:const teamSelectWithVKCount = "governance_teams.*, (SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS virtual_key_count"This is optional - the current approach works correctly and the duplication is manageable.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` at line 4502, Extract the duplicated SELECT clause used to compute virtual_key_count into a single constant (e.g., teamSelectWithVKCount) and replace the repeated string in GetTeams, GetTeamsPaginated, GetTeam, and GetGovernanceConfig with that constant; ensure the constant value is exactly "governance_teams.*, (SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id = governance_teams.id) AS virtual_key_count" and update each query-building site to reference teamSelectWithVKCount so behavior remains identical.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/openapi/openapi.json`:
- Around line 51646-51668: The OpenAPI schema removed the request field
enable_litellm_fallbacks without a deprecation path; either restore it as a
deprecated property on the request schemas or explicitly bump the API version
and document the breaking change. Update the request schema definitions to
include "enable_litellm_fallbacks": { "type": "boolean", "deprecated": true,
"description": "Deprecated: use compat.convert_text_to_chat" } so existing
clients validate for one release, or alternatively add an API-version bump and
release note; ensure this change aligns with the DB migration function
migrationReplaceEnableLiteLLMWithCompatColumns and mirror the same decision for
the other affected schema instances (around the second range noted).
---
Nitpick comments:
In `@framework/configstore/rdb.go`:
- Line 4502: Extract the duplicated SELECT clause used to compute
virtual_key_count into a single constant (e.g., teamSelectWithVKCount) and
replace the repeated string in GetTeams, GetTeamsPaginated, GetTeam, and
GetGovernanceConfig with that constant; ensure the constant value is exactly
"governance_teams.*, (SELECT COUNT(*) FROM governance_virtual_keys WHERE team_id
= governance_teams.id) AS virtual_key_count" and update each query-building site
to reference teamSelectWithVKCount so behavior remains identical.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 207732d7-ecab-4999-88d4-1e1abbfc6747
📒 Files selected for processing (6)
docs/openapi/openapi.jsondocs/openapi/schemas/management/governance.yamldocs/openapi/schemas/management/users.yamlframework/configstore/rdb.goframework/configstore/tables/team.goplugins/governance/store.go
✅ Files skipped from review due to trivial changes (1)
- docs/openapi/schemas/management/users.yaml
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/openapi/schemas/management/governance.yaml
- framework/configstore/tables/team.go
924e214 to
bcc3d05
Compare
bcc3d05 to
5900a79
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docs/openapi/openapi.json`:
- Around line 41537-41541: The OpenAPI spec defines the GoogleApiKeyAuth
security scheme but never references it; update the operation objects that
should accept Google API keys (e.g., the /api/governance/virtual-keys/quota path
and any other relevant operations) to include GoogleApiKeyAuth in their security
arrays by adding an entry like {"GoogleApiKeyAuth": []} to each operation's
security array (or create a security array if missing), ensuring the
operation-level security lists both existing schemes (if any) and
GoogleApiKeyAuth so clients can discover and use the x-goog-api-key header.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0fadeafe-9447-4e36-ac24-68b7aa15b911
📒 Files selected for processing (6)
docs/openapi/openapi.jsondocs/openapi/schemas/management/governance.yamldocs/openapi/schemas/management/users.yamlframework/configstore/rdb.goframework/configstore/tables/team.goplugins/governance/store.go
✅ Files skipped from review due to trivial changes (2)
- docs/openapi/schemas/management/users.yaml
- plugins/governance/store.go
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/openapi/schemas/management/governance.yaml
- framework/configstore/tables/team.go
5900a79 to
ea7a2d2
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
framework/configstore/rdb.go (1)
2522-2522: Please add one regression test for every team read path.This field now depends on the same projection being present in four separate queries. A small integration test covering list, paginated list, single-team, and
GetGovernanceConfigwould make future drift much harder.Also applies to: 2527-2529, 2571-2571, 2585-2588, 3436-3438
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@framework/configstore/rdb.go` at line 2522, Add regression tests that assert the virtual_key_count projection introduced by teamSelectWithVKCount is present and correct across every team read path: the teams list (e.g., ListTeams), the paginated list (e.g., ListTeamsPaginated or equivalent), the single-team getter (e.g., GetTeamByID/GetTeam), and GetGovernanceConfig. For each test, seed a known set of governance_virtual_keys linked to a team, call the corresponding method, and assert the returned record(s) include virtual_key_count with the expected numeric value; include at least one test for empty/zero count and one for >0 to catch projection drift.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@framework/configstore/rdb.go`:
- Line 2522: Add regression tests that assert the virtual_key_count projection
introduced by teamSelectWithVKCount is present and correct across every team
read path: the teams list (e.g., ListTeams), the paginated list (e.g.,
ListTeamsPaginated or equivalent), the single-team getter (e.g.,
GetTeamByID/GetTeam), and GetGovernanceConfig. For each test, seed a known set
of governance_virtual_keys linked to a team, call the corresponding method, and
assert the returned record(s) include virtual_key_count with the expected
numeric value; include at least one test for empty/zero count and one for >0 to
catch projection drift.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a357f8c7-115c-420b-9d45-f9b86f132e8d
📒 Files selected for processing (6)
docs/openapi/openapi.jsondocs/openapi/schemas/management/governance.yamldocs/openapi/schemas/management/users.yamlframework/configstore/rdb.goframework/configstore/tables/team.goplugins/governance/store.go
✅ Files skipped from review due to trivial changes (2)
- docs/openapi/schemas/management/users.yaml
- plugins/governance/store.go
🚧 Files skipped from review as they are similar to previous changes (3)
- framework/configstore/tables/team.go
- docs/openapi/schemas/management/governance.yaml
- docs/openapi/openapi.json
Merge activity
|

Summary
Adds a virtual key count field to team queries to display the number of virtual keys associated with each team without requiring additional database queries or loading the full virtual keys collection.
Changes
VirtualKeyCountfield toTableTeamstruct as a computed field populated via correlated subqueryGetTeams,GetTeamsPaginated, andGetTeammethods to include virtual key count in SELECT queriesteam_idandcustomer_idinUpdateVirtualKeyRequestVirtualKeysfield in team table to be omitempty for cleaner JSON responsesType of change
Affected areas
How to test
Verify that team queries return the virtual key count and that virtual key updates handle nullable team/customer IDs:
Test API endpoints that retrieve teams to ensure
virtual_key_countfield is populated correctly.Screenshots/Recordings
N/A
Breaking changes
Related issues
N/A
Security considerations
The correlated subquery only counts records and doesn't expose sensitive virtual key data. No additional security implications.
Checklist
docs/contributing/README.mdand followed the guidelines