Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions framework/configstore/rdb.go
Original file line number Diff line number Diff line change
Expand Up @@ -2417,8 +2417,6 @@ func (s *RDBConfigStore) GetVirtualKeysPaginated(ctx context.Context, params Vir
baseQuery = baseQuery.Where("customer_id = ?", params.CustomerID)
} else if params.TeamID != "" {
baseQuery = baseQuery.Where("team_id = ?", params.TeamID)
} else if params.AccessProfileID > 0 {
baseQuery = baseQuery.Where("access_profile_id = ?", params.AccessProfileID)
}
if params.Search != "" {
search := "%" + strings.ToLower(params.Search) + "%"
Expand Down
1 change: 0 additions & 1 deletion framework/configstore/store.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ type VirtualKeyQueryParams struct {
Search string
CustomerID string
TeamID string
AccessProfileID uint // When set, return VKs attached to this access profile template
SortBy string // name, budget_spent, created_at, status (default: created_at)
Order string // asc, desc (default: asc)
Export bool // When true, skip default pagination limits (caller controls limit)
Expand Down
30 changes: 10 additions & 20 deletions framework/configstore/tables/virtualkey.go
Original file line number Diff line number Diff line change
Expand Up @@ -215,11 +215,10 @@ type TableVirtualKey struct {
ProviderConfigs []TableVirtualKeyProviderConfig `gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE" json:"provider_configs"` // Empty means no providers allowed (deny-by-default)
MCPConfigs []TableVirtualKeyMCPConfig `gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE" json:"mcp_configs"`

// Foreign key relationships (mutually exclusive: TeamID, CustomerID, or AccessProfileID)
TeamID *string `gorm:"type:varchar(255);index" json:"team_id,omitempty"`
CustomerID *string `gorm:"type:varchar(255);index" json:"customer_id,omitempty"`
AccessProfileID *uint `gorm:"index" json:"access_profile_id,omitempty"`
RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"`
// Foreign key relationships (mutually exclusive: either TeamID or CustomerID, not both)
TeamID *string `gorm:"type:varchar(255);index" json:"team_id,omitempty"`
CustomerID *string `gorm:"type:varchar(255);index" json:"customer_id,omitempty"`
RateLimitID *string `gorm:"type:varchar(255);index" json:"rate_limit_id,omitempty"`

CalendarAligned bool `gorm:"default:false" json:"calendar_aligned"`

Expand Down Expand Up @@ -254,22 +253,13 @@ func (vk *TableVirtualKey) IsActiveValue() bool {
return *vk.IsActive
}

// BeforeSave is a GORM hook that enforces mutual exclusion among team, customer, and
// access profile attachments, computes a SHA-256 hash of the plaintext value for indexed
// lookups, and encrypts the virtual key value before writing to the database.
// BeforeSave is a GORM hook that enforces mutual exclusion (team vs customer), computes
// a SHA-256 hash of the plaintext value for indexed lookups, and encrypts the virtual key
// value before writing to the database.
func (vk *TableVirtualKey) BeforeSave(tx *gorm.DB) error {
entityCount := 0
if vk.TeamID != nil {
entityCount++
}
if vk.CustomerID != nil {
entityCount++
}
if vk.AccessProfileID != nil && *vk.AccessProfileID > 0 {
entityCount++
}
if entityCount > 1 {
return fmt.Errorf("virtual key can only be attached to one of: team, customer, or access profile")
// Enforce mutual exclusion: VK can belong to either Team OR Customer, not both
if vk.TeamID != nil && vk.CustomerID != nil {
return fmt.Errorf("virtual key cannot belong to both team and customer")
}
Comment thread
BearTS marked this conversation as resolved.

// Hash must be computed before encryption (from plaintext value)
Expand Down
82 changes: 17 additions & 65 deletions transports/bifrost-http/handlers/governance.go
Original file line number Diff line number Diff line change
Expand Up @@ -92,10 +92,9 @@ type CreateVirtualKeyRequest struct {
MCPClientName string `json:"mcp_client_name" validate:"required"`
ToolsToExecute schemas.WhiteList `json:"tools_to_execute,omitempty"`
} `json:"mcp_configs,omitempty"` // Empty means no MCP clients allowed (deny-by-default)
TeamID *string `json:"team_id,omitempty"` // Mutually exclusive with CustomerID and AccessProfileID
CustomerID *string `json:"customer_id,omitempty"` // Mutually exclusive with TeamID and AccessProfileID
AccessProfileID *uint `json:"access_profile_id,omitempty"` // Mutually exclusive with TeamID and CustomerID
Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget: each must have a unique reset_duration
TeamID *string `json:"team_id,omitempty"` // Mutually exclusive with CustomerID
CustomerID *string `json:"customer_id,omitempty"` // Mutually exclusive with TeamID
Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget: each must have a unique reset_duration
RateLimit *CreateRateLimitRequest `json:"rate_limit,omitempty"`
IsActive *bool `json:"is_active,omitempty"`
CalendarAligned bool `json:"calendar_aligned,omitempty"` // When true, all budgets reset at clean calendar boundaries
Expand All @@ -122,7 +121,6 @@ type UpdateVirtualKeyRequest struct {
} `json:"mcp_configs,omitempty"`
TeamID *string `json:"team_id,omitempty"`
CustomerID *string `json:"customer_id,omitempty"`
AccessProfileID *uint `json:"access_profile_id,omitempty"`
Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget: replaces all VK-level budgets
RateLimit *UpdateRateLimitRequest `json:"rate_limit,omitempty"`
IsActive *bool `json:"is_active,omitempty"`
Expand Down Expand Up @@ -477,13 +475,12 @@ func (h *GovernanceHandler) getVirtualKeys(ctx *fasthttp.RequestCtx) {
search := string(ctx.QueryArgs().Peek("search"))
customerID := string(ctx.QueryArgs().Peek("customer_id"))
teamID := string(ctx.QueryArgs().Peek("team_id"))
accessProfileIDStr := string(ctx.QueryArgs().Peek("access_profile_id"))
sortBy := string(ctx.QueryArgs().Peek("sort_by"))
order := string(ctx.QueryArgs().Peek("order"))
isExport := string(ctx.QueryArgs().Peek("export")) == "true"
excludeAccessProfileManagedVirtual := string(ctx.QueryArgs().Peek("exclude_access_profile_managed_virtual")) == "true"

if limitStr != "" || offsetStr != "" || search != "" || customerID != "" || teamID != "" || accessProfileIDStr != "" || sortBy != "" || isExport || excludeAccessProfileManagedVirtual {
if limitStr != "" || offsetStr != "" || search != "" || customerID != "" || teamID != "" || sortBy != "" || isExport || excludeAccessProfileManagedVirtual {
Comment on lines 478 to +483

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Honor order-only virtual-key list requests.

If a caller sends only ?order=desc, Line 481 falls through to GetVirtualKeys and silently ignores the requested ordering. Include order != "" in this branch condition so sort-direction changes still use GetVirtualKeysPaginated.

Suggested fix
-	if limitStr != "" || offsetStr != "" || search != "" || customerID != "" || teamID != "" || sortBy != "" || isExport || excludeAccessProfileManagedVirtual {
+	if limitStr != "" || offsetStr != "" || search != "" || customerID != "" || teamID != "" || sortBy != "" || order != "" || isExport || excludeAccessProfileManagedVirtual {
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
sortBy := string(ctx.QueryArgs().Peek("sort_by"))
order := string(ctx.QueryArgs().Peek("order"))
isExport := string(ctx.QueryArgs().Peek("export")) == "true"
excludeAccessProfileManagedVirtual := string(ctx.QueryArgs().Peek("exclude_access_profile_managed_virtual")) == "true"
if limitStr != "" || offsetStr != "" || search != "" || customerID != "" || teamID != "" || accessProfileIDStr != "" || sortBy != "" || isExport || excludeAccessProfileManagedVirtual {
if limitStr != "" || offsetStr != "" || search != "" || customerID != "" || teamID != "" || sortBy != "" || isExport || excludeAccessProfileManagedVirtual {
sortBy := string(ctx.QueryArgs().Peek("sort_by"))
order := string(ctx.QueryArgs().Peek("order"))
isExport := string(ctx.QueryArgs().Peek("export")) == "true"
excludeAccessProfileManagedVirtual := string(ctx.QueryArgs().Peek("exclude_access_profile_managed_virtual")) == "true"
if limitStr != "" || offsetStr != "" || search != "" || customerID != "" || teamID != "" || sortBy != "" || order != "" || isExport || excludeAccessProfileManagedVirtual {
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@transports/bifrost-http/handlers/governance.go` around lines 476 - 481, The
branch that decides between GetVirtualKeysPaginated and GetVirtualKeys ignores
requests that only set order (e.g., ?order=desc); update the condition that
currently checks limitStr, offsetStr, search, customerID, teamID, sortBy,
isExport, and excludeAccessProfileManagedVirtual to also include order != "" so
that ordering-only requests route to GetVirtualKeysPaginated; locate the
variables sortBy and order and the call sites of GetVirtualKeysPaginated /
GetVirtualKeys in governance.go to make this change.

// Paginated/filtered path
params := configstore.VirtualKeyQueryParams{
Search: search,
Expand All @@ -494,18 +491,6 @@ func (h *GovernanceHandler) getVirtualKeys(ctx *fasthttp.RequestCtx) {
Export: isExport,
ExcludeAccessProfileManagedVirtual: excludeAccessProfileManagedVirtual,
}
if accessProfileIDStr != "" {
apID, err := strconv.ParseUint(accessProfileIDStr, 10, 0)
if err != nil {
SendError(ctx, 400, "Invalid access_profile_id parameter: must be a number")
return
}
if (customerID != "" || teamID != "") && apID > 0 {
SendError(ctx, 400, "access_profile_id cannot be combined with team_id or customer_id filters")
return
}
params.AccessProfileID = uint(apID)
}
if limitStr != "" {
n, err := strconv.Atoi(limitStr)
if err != nil {
Expand Down Expand Up @@ -580,19 +565,9 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) {
SendError(ctx, 400, "Virtual key name is required")
return
}
// Validate mutually exclusive TeamID, CustomerID, and AccessProfileID
entityCount := 0
if req.TeamID != nil {
entityCount++
}
if req.CustomerID != nil {
entityCount++
}
if req.AccessProfileID != nil && *req.AccessProfileID > 0 {
entityCount++
}
if entityCount > 1 {
SendError(ctx, 400, "VirtualKey can only be attached to one of: Team, Customer, or AccessProfile")
// Validate mutually exclusive TeamID and CustomerID
if req.TeamID != nil && req.CustomerID != nil {
SendError(ctx, 400, "VirtualKey cannot be attached to both Team and Customer")
return
}
// Validate budgets if provided
Expand Down Expand Up @@ -638,7 +613,6 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) {
Description: req.Description,
TeamID: req.TeamID,
CustomerID: req.CustomerID,
AccessProfileID: req.AccessProfileID,
IsActive: isActive,
CalendarAligned: req.CalendarAligned,
}
Expand Down Expand Up @@ -870,19 +844,9 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
SendError(ctx, 400, "Invalid JSON")
return
}
// Validate mutually exclusive TeamID, CustomerID, and AccessProfileID
entityCount := 0
if req.TeamID != nil {
entityCount++
}
if req.CustomerID != nil {
entityCount++
}
if req.AccessProfileID != nil && *req.AccessProfileID > 0 {
entityCount++
}
if entityCount > 1 {
SendError(ctx, 400, "VirtualKey can only be attached to one of: Team, Customer, or AccessProfile")
// Validate mutually exclusive TeamID and CustomerID
if req.TeamID != nil && req.CustomerID != nil {
SendError(ctx, 400, "VirtualKey cannot be attached to both Team and Customer")
return
}
vk, err := h.configStore.GetVirtualKey(ctx, vkID)
Expand Down Expand Up @@ -934,21 +898,16 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
}
if req.TeamID != nil {
vk.TeamID = req.TeamID
vk.CustomerID = nil
vk.AccessProfileID = nil
} else if req.CustomerID != nil {
vk.CustomerID = nil // Clear CustomerID if setting TeamID
}
if req.CustomerID != nil {
vk.CustomerID = req.CustomerID
vk.TeamID = nil
vk.AccessProfileID = nil
} else if req.AccessProfileID != nil {
vk.AccessProfileID = req.AccessProfileID
vk.TeamID = nil
vk.CustomerID = nil
} else {
// All nil — clear entity assignment
vk.TeamID = nil // Clear TeamID if setting CustomerID
}
// When both TeamID and CustomerID are nil
if req.TeamID == nil && req.CustomerID == nil {
vk.TeamID = nil
vk.CustomerID = nil
vk.AccessProfileID = nil
}
if req.IsActive != nil {
vk.IsActive = req.IsActive
Expand Down Expand Up @@ -1084,13 +1043,6 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) {
if err := h.configStore.UpdateVirtualKey(ctx, vk, tx); err != nil {
return err
}
// UpdateVirtualKey's Select list excludes access_profile_id to protect config-sync paths.
// Persist it separately so API-driven entity assignment is always saved.
if err := tx.Model(&configstoreTables.TableVirtualKey{}).
Where("id = ?", vk.ID).
Updates(map[string]interface{}{"access_profile_id": vk.AccessProfileID}).Error; err != nil {
return fmt.Errorf("failed to update virtual key access_profile_id: %w", err)
}
if req.ProviderConfigs != nil {
// Get existing provider configs for comparison
var existingConfigs []configstoreTables.TableVirtualKeyProviderConfig
Expand Down
19 changes: 2 additions & 17 deletions ui/app/_fallbacks/enterprise/lib/store/apis/accessProfileApi.ts
Original file line number Diff line number Diff line change
@@ -1,22 +1,7 @@
import { GetAccessProfilesParams, GetAccessProfilesResponse, GetUserAccessProfilesResponse } from "@enterprise/lib/types/accessProfile";
import { GetUserAccessProfilesResponse } from "@enterprise/lib/types/accessProfile";

// OSS build has no access-profile backend — return undefined data so consumers
// fall back gracefully (empty lists, no AP-specific UI rendered).
export const useGetAccessProfilesQuery = (
_params?: GetAccessProfilesParams | void,
_opts?: { skip?: boolean },
): {
data: GetAccessProfilesResponse | undefined;
isLoading: boolean;
isError: boolean;
error: null;
} => ({
data: undefined,
isLoading: false,
isError: false,
error: null,
});

// (e.g. useVirtualKeyUsage) fall back to VK-owned budget/rate-limit values.
export const useGetUserAccessProfilesQuery = (
_userId: string,
_opts?: { skip?: boolean; pollingInterval?: number },
Expand Down
16 changes: 0 additions & 16 deletions ui/app/_fallbacks/enterprise/lib/types/accessProfile.ts
Original file line number Diff line number Diff line change
Expand Up @@ -38,20 +38,4 @@ export interface UserAccessProfile {

export interface GetUserAccessProfilesResponse {
access_profiles: UserAccessProfile[];
}

export interface GetAccessProfilesParams {
limit?: number;
offset?: number;
search?: string;
tags?: string;
is_active?: boolean;
}

export interface GetAccessProfilesResponse {
access_profiles: { id: number; name: string; [key: string]: unknown }[];
count: number;
total_count: number;
limit: number;
offset: number;
}
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,7 @@ interface VirtualKeyDetailSheetProps {
}

export default function VirtualKeyDetailSheet({ virtualKey, onClose }: VirtualKeyDetailSheetProps) {
const { assignedUsers, managingProfile, hasApRateLimit, displayBudgets, displayRateLimit } = useVirtualKeyUsage(virtualKey);
const isManagedByProfile = !!virtualKey.access_profile_id;
const { assignedUsers, managingProfile, isManagedByProfile, hasApRateLimit, displayBudgets, displayRateLimit } = useVirtualKeyUsage(virtualKey);

const getEntityInfo = () => {
if (virtualKey.team) {
Expand Down
20 changes: 2 additions & 18 deletions ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@ import { Customer, Team, VirtualKey } from "@/lib/types/governance";
import { cn } from "@/lib/utils";
import { formatCurrency } from "@/lib/utils/governance";
import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib";
import { useGetAccessProfilesQuery } from "@enterprise/lib/store/apis/accessProfileApi";
import {
ArrowDown,
ArrowUp,
Expand Down Expand Up @@ -80,8 +79,6 @@ function virtualKeysToCSV(vks: VirtualKey[], accessProfileNames: Record<number,
? `Team: ${vk.team.name}`
: vk.customer
? `Customer: ${vk.customer.name}`
: vk.access_profile_id
? `Access Profile: ${accessProfileNames[vk.access_profile_id] ?? vk.access_profile_id}`
: "";
const budgetLimit = vk.budgets?.length ? vk.budgets.map((b) => formatCurrency(b.max_limit)).join("; ") : "";
const budgetSpent = vk.budgets?.length ? vk.budgets.map((b) => formatCurrency(b.current_usage)).join("; ") : "";
Expand Down Expand Up @@ -295,15 +292,6 @@ export default function VirtualKeysTable({
const [showBulkRotateDialog, setShowBulkRotateDialog] = useState(false);
const [fetchVirtualKeys, { isFetching: isExporting }] = useLazyGetVirtualKeysQuery();

const { data: accessProfilesData } = useGetAccessProfilesQuery({ limit: 100 });
const accessProfileNames = useMemo(() => {
const map: Record<number, string> = {};
for (const ap of accessProfilesData?.access_profiles ?? []) {
map[ap.id] = ap.name;
}
return map;
}, [accessProfilesData]);

// Derive objects from props so they stay in sync with RTK cache updates
const editingVirtualKey = useMemo(
() => (editingVirtualKeyId ? (virtualKeys.find((vk) => vk.id === editingVirtualKeyId) ?? null) : null),
Expand Down Expand Up @@ -477,7 +465,7 @@ export default function VirtualKeysTable({

const handleExportCSV = async () => {
if (exportScope === "current_page") {
downloadCSV(virtualKeysToCSV(virtualKeys, accessProfileNames));
downloadCSV(virtualKeysToCSV(virtualKeys));
toast.success(`Exported ${virtualKeys.length} virtual keys`);
setShowExportDialog(false);
return;
Expand All @@ -499,7 +487,7 @@ export default function VirtualKeysTable({
export: true,
}).unwrap();

downloadCSV(virtualKeysToCSV(result.virtual_keys, accessProfileNames));
downloadCSV(virtualKeysToCSV(result.virtual_keys));
toast.success(`Exported ${result.virtual_keys.length} virtual keys`);
setShowExportDialog(false);
} catch (error) {
Expand Down Expand Up @@ -797,10 +785,6 @@ export default function VirtualKeysTable({
<Badge variant="outline" className="block max-w-full truncate text-left">
Customer: {vk.customer.name}
</Badge>
) : vk.access_profile_id ? (
<Badge variant="outline" className="block max-w-full truncate text-left">
AP: {accessProfileNames[vk.access_profile_id] ?? vk.access_profile_id}
</Badge>
) : (
<span className="text-muted-foreground max-w-full truncate text-left text-sm">-</span>
)}
Expand Down
3 changes: 0 additions & 3 deletions ui/components/ui/combobox.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -295,7 +295,6 @@ interface ComboboxSelectBaseProps {
hideClear?: boolean;
className?: string;
emptyMessage?: string;
"data-testid"?: string;
}

interface ComboboxSelectSingleProps extends ComboboxSelectBaseProps {
Expand All @@ -321,7 +320,6 @@ function ComboboxSelect(props: ComboboxSelectProps) {
className,
emptyMessage = "No results found.",
noPortal,
"data-testid": dataTestId,
} = props;

const [open, setOpen] = React.useState(false);
Expand Down Expand Up @@ -444,7 +442,6 @@ function ComboboxSelect(props: ComboboxSelectProps) {
role="combobox"
aria-expanded={open}
disabled={disabled}
data-testid={dataTestId}
className={cn(
"h-8 w-full justify-between !bg-transparent font-normal active:scale-none",
!selectedLabel && "text-muted-foreground",
Expand Down
1 change: 0 additions & 1 deletion ui/lib/store/apis/baseApi.ts
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,6 @@ export const baseApi = createApi({
"Versions",
"Sessions",
"AccessProfiles",
"AccessProfileVirtualKeys",
"BusinessUnits",
"PromptDeployments",
"AuthType",
Expand Down
1 change: 0 additions & 1 deletion ui/lib/store/apis/governanceApi.ts
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,6 @@ export const governanceApi = baseApi.injectEndpoints({
...(params?.search && { search: params.search }),
...(params?.customer_id && { customer_id: params.customer_id }),
...(params?.team_id && { team_id: params.team_id }),
...(params?.access_profile_id !== undefined && { access_profile_id: params.access_profile_id }),
...(params?.exclude_access_profile_managed_virtual === true && {
exclude_access_profile_managed_virtual: "true",
}),
Expand Down
Loading
Loading