diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 72e958f0653..3191ea5e684 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -746,6 +746,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationDropLegacyCalendarAlignedColumns(ctx, db); err != nil { return err } + if err := migrationAddVKAccessProfileIDColumn(ctx, db); err != nil { + return err + } return nil } @@ -7661,3 +7664,38 @@ func migrationAddTeamCalendarAlignedColumn(ctx context.Context, db *gorm.DB) err } return nil } + +// migrationAddVKAccessProfileIDColumn adds access_profile_id to governance_virtual_keys +// so that existing VKs can be attached directly to an access profile template (enterprise feature). +func migrationAddVKAccessProfileIDColumn(ctx context.Context, db *gorm.DB) error { + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: "add_vk_access_profile_id_column", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mig := tx.Migrator() + if !mig.HasColumn(&tables.TableVirtualKey{}, "access_profile_id") { + if err := mig.AddColumn(&tables.TableVirtualKey{}, "AccessProfileID"); err != nil { + return fmt.Errorf("failed to add access_profile_id column to governance_virtual_keys: %w", err) + } + } + if !mig.HasIndex(&tables.TableVirtualKey{}, "idx_governance_virtual_keys_access_profile_id") { + if err := mig.CreateIndex(&tables.TableVirtualKey{}, "AccessProfileID"); err != nil { + return fmt.Errorf("failed to create index on governance_virtual_keys.access_profile_id: %w", err) + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mig := tx.Migrator() + if mig.HasColumn(&tables.TableVirtualKey{}, "access_profile_id") { + return mig.DropColumn(&tables.TableVirtualKey{}, "access_profile_id") + } + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running add_vk_access_profile_id_column migration: %s", err.Error()) + } + return nil +} diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index 22a84d8fd72..8b702367ebf 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -2281,6 +2281,8 @@ 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) + "%" diff --git a/framework/configstore/store.go b/framework/configstore/store.go index d6ad96db9f1..71f474140c4 100644 --- a/framework/configstore/store.go +++ b/framework/configstore/store.go @@ -20,6 +20,7 @@ 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) diff --git a/framework/configstore/tables/virtualkey.go b/framework/configstore/tables/virtualkey.go index d2015d9cf85..d018dd99c5c 100644 --- a/framework/configstore/tables/virtualkey.go +++ b/framework/configstore/tables/virtualkey.go @@ -205,10 +205,11 @@ 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: 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"` + // 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"` CalendarAligned bool `gorm:"default:false" json:"calendar_aligned"` @@ -243,13 +244,22 @@ func (vk *TableVirtualKey) IsActiveValue() bool { return *vk.IsActive } -// 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. +// 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. func (vk *TableVirtualKey) BeforeSave(tx *gorm.DB) error { - // 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") + 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") } // Hash must be computed before encryption (from plaintext value) diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 835133e41d0..51eb0050f8f 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -91,9 +91,10 @@ 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 - CustomerID *string `json:"customer_id,omitempty"` // Mutually exclusive with TeamID - Budgets []CreateBudgetRequest `json:"budgets,omitempty"` // Multi-budget: each must have a unique reset_duration + 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 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 @@ -119,6 +120,7 @@ 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"` @@ -372,12 +374,13 @@ 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 != "" || sortBy != "" || isExport || excludeAccessProfileManagedVirtual { + if limitStr != "" || offsetStr != "" || search != "" || customerID != "" || teamID != "" || accessProfileIDStr != "" || sortBy != "" || isExport || excludeAccessProfileManagedVirtual { // Paginated/filtered path params := configstore.VirtualKeyQueryParams{ Search: search, @@ -388,6 +391,18 @@ 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 { @@ -462,9 +477,19 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "Virtual key name is required") return } - // 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") + // 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") return } // Validate budgets if provided @@ -510,6 +535,7 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) { Description: req.Description, TeamID: req.TeamID, CustomerID: req.CustomerID, + AccessProfileID: req.AccessProfileID, IsActive: isActive, CalendarAligned: req.CalendarAligned, } @@ -737,9 +763,19 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "Invalid JSON") return } - // 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") + // 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") return } vk, err := h.configStore.GetVirtualKey(ctx, vkID) @@ -791,16 +827,21 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { } if req.TeamID != nil { vk.TeamID = req.TeamID - vk.CustomerID = nil // Clear CustomerID if setting TeamID - } - if req.CustomerID != nil { + vk.CustomerID = nil + vk.AccessProfileID = nil + } else if req.CustomerID != nil { vk.CustomerID = req.CustomerID - 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.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 + vk.CustomerID = nil + vk.AccessProfileID = nil } if req.IsActive != nil { vk.IsActive = req.IsActive @@ -935,6 +976,13 @@ 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 diff --git a/ui/app/_fallbacks/enterprise/lib/store/apis/accessProfileApi.ts b/ui/app/_fallbacks/enterprise/lib/store/apis/accessProfileApi.ts index c5038baeee0..7d162e5a33a 100644 --- a/ui/app/_fallbacks/enterprise/lib/store/apis/accessProfileApi.ts +++ b/ui/app/_fallbacks/enterprise/lib/store/apis/accessProfileApi.ts @@ -1,7 +1,22 @@ -import { GetUserAccessProfilesResponse } from "@enterprise/lib/types/accessProfile"; +import { GetAccessProfilesParams, GetAccessProfilesResponse, GetUserAccessProfilesResponse } from "@enterprise/lib/types/accessProfile"; // OSS build has no access-profile backend — return undefined data so consumers -// (e.g. useVirtualKeyUsage) fall back to VK-owned budget/rate-limit values. +// 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, +}); + export const useGetUserAccessProfilesQuery = ( _userId: string, _opts?: { skip?: boolean; pollingInterval?: number }, diff --git a/ui/app/_fallbacks/enterprise/lib/types/accessProfile.ts b/ui/app/_fallbacks/enterprise/lib/types/accessProfile.ts index 414743dafb1..22d245a4fda 100644 --- a/ui/app/_fallbacks/enterprise/lib/types/accessProfile.ts +++ b/ui/app/_fallbacks/enterprise/lib/types/accessProfile.ts @@ -38,4 +38,20 @@ 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; } \ No newline at end of file diff --git a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx index 22a984b7b54..c9d6d3566f7 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx @@ -45,11 +45,12 @@ import { } from "@/lib/store"; import { KnownProvider } from "@/lib/types/config"; import { CreateVirtualKeyRequest, Customer, Team, UpdateVirtualKeyRequest, VirtualKey } from "@/lib/types/governance"; +import { useGetAccessProfilesQuery } from "@enterprise/lib/store/apis/accessProfileApi"; import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib"; import { zodResolver } from "@hookform/resolvers/zod"; import { useNavigate } from "@tanstack/react-router"; import { Info, Lock, RotateCcw, Trash2, Users, X } from "lucide-react"; -import { useEffect, useState } from "react"; +import { useEffect, useMemo, useState } from "react"; import { useForm } from "react-hook-form"; import { components, MultiValueProps, OptionProps } from "react-select"; import { toast } from "sonner"; @@ -62,6 +63,9 @@ interface VirtualKeySheetProps { // When set, the new VK is created under this team. The entity assignment is pre-set // and cannot be changed (but all other fields remain editable). defaultTeamId?: string; + // When set, the new VK is created under this access profile. The entity assignment is + // pre-set and cannot be changed. + defaultAccessProfileId?: number; onSave: () => void; onCancel: () => void; } @@ -106,9 +110,10 @@ const formSchema = z description: z.string().optional(), providerConfigs: z.array(providerConfigSchema).optional(), mcpConfigs: z.array(mcpConfigSchema).optional(), - entityType: z.enum(["team", "customer", "none"]), + entityType: z.enum(["team", "customer", "access_profile", "none"]), teamId: z.string().optional(), customerId: z.string().optional(), + accessProfileId: z.string().optional(), isActive: z.boolean(), // Budget budgetCalendarAligned: z.boolean(), @@ -137,11 +142,15 @@ const formSchema = z if (data.entityType === "customer") { return data.customerId && data.customerId.trim() !== ""; } + // If entityType is "access_profile", accessProfileId must be provided and not empty + if (data.entityType === "access_profile") { + return data.accessProfileId && data.accessProfileId.trim() !== ""; + } return true; }, { - message: "Please select a valid team or customer when assignment type is chosen", - path: ["entityType"], // This will show the error on the entityType field + message: "Please select a valid team, customer, or access profile when assignment type is chosen", + path: ["entityType"], }, ); @@ -154,7 +163,7 @@ type VirtualKeyType = { provider: string; }; -export default function VirtualKeySheet({ virtualKey, teams, customers, defaultTeamId, onSave, onCancel }: VirtualKeySheetProps) { +export default function VirtualKeySheet({ virtualKey, teams, customers, defaultTeamId, defaultAccessProfileId, onSave, onCancel }: VirtualKeySheetProps) { const [isOpen, setIsOpen] = useState(true); const navigate = useNavigate(); const isEditing = !!virtualKey; @@ -172,6 +181,7 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT const attachedTeamId = isEditing ? virtualKey?.team_id || "" : defaultTeamId || ""; const attachedTeam = attachedTeamId ? teams.find((t) => t.id === attachedTeamId) : undefined; const isTeamLocked = !isEditing && !!defaultTeamId; + const isAPLocked = !isEditing && !!defaultAccessProfileId; const handleClose = () => { setIsOpen(false); @@ -186,6 +196,8 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT const [createVirtualKey, { isLoading: isCreating }] = useCreateVirtualKeyMutation(); const [updateVirtualKey, { isLoading: isUpdating }] = useUpdateVirtualKeyMutation(); const { data: mcpClientsResponse, error: mcpClientsError } = useGetMCPClientsQuery(); + const { data: accessProfilesData } = useGetAccessProfilesQuery({ limit: 100 }); + const accessProfiles = accessProfilesData?.access_profiles ?? []; const mcpClientsData = mcpClientsResponse?.clients || []; const isLoading = isCreating || isUpdating; @@ -224,9 +236,19 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT mcp_client_name: config.mcp_client?.name || "", tools_to_execute: config.tools_to_execute || [], })) || [], - entityType: virtualKey?.team_id ? "team" : virtualKey?.customer_id ? "customer" : !isEditing && defaultTeamId ? "team" : "none", + entityType: virtualKey?.team_id ? "team" + : virtualKey?.customer_id ? "customer" + : virtualKey?.access_profile_id ? "access_profile" + : !isEditing && defaultTeamId ? "team" + : !isEditing && defaultAccessProfileId ? "access_profile" + : "none", teamId: virtualKey?.team_id || (!isEditing ? defaultTeamId || "" : ""), customerId: virtualKey?.customer_id || "", + accessProfileId: virtualKey?.access_profile_id + ? String(virtualKey.access_profile_id) + : !isEditing && defaultAccessProfileId + ? String(defaultAccessProfileId) + : "", isActive: virtualKey?.is_active ?? true, budgets: virtualKey?.budgets && virtualKey.budgets.length > 0 @@ -261,16 +283,22 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT } }, [mcpClientsError]); - // Clear team/customer IDs when entityType changes to "none" + // Clear entity ID fields when entityType changes useEffect(() => { const entityType = form.watch("entityType"); if (entityType === "none") { form.setValue("teamId", "", { shouldDirty: true }); form.setValue("customerId", "", { shouldDirty: true }); + form.setValue("accessProfileId", "", { shouldDirty: true }); } else if (entityType === "team") { form.setValue("customerId", "", { shouldDirty: true }); + form.setValue("accessProfileId", "", { shouldDirty: true }); } else if (entityType === "customer") { form.setValue("teamId", "", { shouldDirty: true }); + form.setValue("accessProfileId", "", { shouldDirty: true }); + } else if (entityType === "access_profile") { + form.setValue("teamId", "", { shouldDirty: true }); + form.setValue("customerId", "", { shouldDirty: true }); } }, [form.watch("entityType"), form]); @@ -368,6 +396,27 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT const [showCalendarAlignWarning, setShowCalendarAlignWarning] = useState(false); const [showReassignTeamWarning, setShowReassignTeamWarning] = useState(false); const [pendingTeamId, setPendingTeamId] = useState(null); + const [showReassignAPWarning, setShowReassignAPWarning] = useState(false); + const [pendingAccessProfileId, setPendingAccessProfileId] = useState(null); + const [showReassignTypeWarning, setShowReassignTypeWarning] = useState(false); + const [pendingEntityType, setPendingEntityType] = useState<"team" | "customer" | "access_profile" | "none" | null>(null); + + const currentAssignmentLabel = useMemo(() => { + if (!isEditing) return null; + if (virtualKey?.team_id) { + const team = teams?.find((t) => t.id === virtualKey.team_id); + return team ? `Team: ${team.name}` : "a team"; + } + if (virtualKey?.customer_id) { + const customer = customers?.find((c) => c.id === virtualKey.customer_id); + return customer ? `Customer: ${customer.name}` : "a customer"; + } + if (virtualKey?.access_profile_id) { + const ap = accessProfiles?.find((p) => p.id === virtualKey.access_profile_id); + return ap ? `Access Profile: ${ap.name}` : "an access profile"; + } + return null; + }, [isEditing, virtualKey, teams, customers, accessProfiles]); const handleCalendarAlignedChange = (checked: boolean) => { if (checked && isEditing) { @@ -451,6 +500,7 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT mcp_configs: data.mcpConfigs, team_id: data.entityType === "team" && data.teamId && data.teamId.trim() !== "" ? data.teamId : undefined, customer_id: data.entityType === "customer" && data.customerId && data.customerId.trim() !== "" ? data.customerId : undefined, + access_profile_id: data.entityType === "access_profile" && data.accessProfileId && data.accessProfileId.trim() !== "" ? Number(data.accessProfileId) : undefined, is_active: data.isActive, calendar_aligned: data.budgetCalendarAligned, }; @@ -493,6 +543,7 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT mcp_configs: data.mcpConfigs, team_id: data.entityType === "team" && data.teamId && data.teamId.trim() !== "" ? data.teamId : undefined, customer_id: data.entityType === "customer" && data.customerId && data.customerId.trim() !== "" ? data.customerId : undefined, + access_profile_id: data.entityType === "access_profile" && data.accessProfileId && data.accessProfileId.trim() !== "" ? Number(data.accessProfileId) : undefined, is_active: data.isActive, // VK-level setting that governs both budget and rate-limit calendar alignment. calendar_aligned: data.budgetCalendarAligned, @@ -1287,6 +1338,99 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT + + { + setShowReassignAPWarning(open); + if (!open) { + setPendingAccessProfileId(null); + } + }} + > + + + Reassign to a different access profile? + + This key is currently assigned to another access profile. Reassigning it will move budget tracking to this access profile — future requests through this key will count against this profile's budget, not the previous one. + + + + setPendingAccessProfileId(null)}> + Cancel + + { + if (pendingAccessProfileId !== null) { + form.setValue("accessProfileId", pendingAccessProfileId, { shouldDirty: true }); + } + setPendingAccessProfileId(null); + setShowReassignAPWarning(false); + }} + > + Reassign + + + + + + {/* Cross-type assignment warning */} + { + if (!open) setPendingEntityType(null); + setShowReassignTypeWarning(open); + }} + > + + + Change assignment? + + This key is currently assigned to {currentAssignmentLabel}. Changing the assignment type will move budget + tracking — future requests will count against the new entity, not the previous one. + + + + setPendingEntityType(null)}> + Cancel + + { + if (pendingEntityType) { + form.setValue("entityType", pendingEntityType, { shouldDirty: true }); + if (pendingEntityType === "team" && teams?.length > 0) { + form.setValue("teamId", teams[0].id, { shouldDirty: true, shouldValidate: true }); + form.setValue("customerId", "", { shouldDirty: true, shouldValidate: true }); + form.setValue("accessProfileId", "", { shouldDirty: true, shouldValidate: true }); + } else if (pendingEntityType === "customer" && customers?.length > 0) { + form.setValue("customerId", customers[0].id, { shouldDirty: true, shouldValidate: true }); + form.setValue("teamId", "", { shouldDirty: true, shouldValidate: true }); + form.setValue("accessProfileId", "", { shouldDirty: true, shouldValidate: true }); + } else if (pendingEntityType === "access_profile") { + const apId = accessProfiles?.length > 0 + ? String(accessProfiles[0].id) + : virtualKey?.access_profile_id ? String(virtualKey.access_profile_id) : ""; + form.setValue("accessProfileId", apId, { shouldDirty: true, shouldValidate: true }); + form.setValue("teamId", "", { shouldDirty: true, shouldValidate: true }); + form.setValue("customerId", "", { shouldDirty: true, shouldValidate: true }); + } else { + form.setValue("teamId", "", { shouldDirty: true, shouldValidate: true }); + form.setValue("customerId", "", { shouldDirty: true, shouldValidate: true }); + form.setValue("accessProfileId", "", { shouldDirty: true, shouldValidate: true }); + } + await form.trigger(["teamId", "customerId", "accessProfileId", "entityType"]); + } + setPendingEntityType(null); + setShowReassignTypeWarning(false); + }} + > + Change Assignment + + + + {/* Rate Limiting Configuration */}
@@ -1398,7 +1542,7 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT - {(teams?.length > 0 || customers?.length > 0) && ( + {(teams?.length > 0 || customers?.length > 0 || accessProfiles?.length > 0 || !!virtualKey?.access_profile_id || !!defaultAccessProfileId) && ( <> @@ -1418,26 +1562,47 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT { value: "none", label: "No Assignment" }, ...(teams?.length > 0 ? [{ value: "team", label: "Assign to Team" }] : []), ...(customers?.length > 0 ? [{ value: "customer", label: "Assign to Customer" }] : []), + ...(accessProfiles?.length > 0 || virtualKey?.access_profile_id || defaultAccessProfileId ? [{ value: "access_profile", label: "Assign to Access Profile" }] : []), ]} value={field.value} onValueChange={async (value) => { const val = value ?? "none"; + const originalType = virtualKey?.team_id ? "team" + : virtualKey?.customer_id ? "customer" + : virtualKey?.access_profile_id ? "access_profile" + : "none"; + if (isEditing && currentAssignmentLabel && val !== "none" && val !== originalType) { + setPendingEntityType(val as "team" | "customer" | "access_profile"); + setShowReassignTypeWarning(true); + return; + } field.onChange(val); if (val === "team" && teams?.length > 0) { form.setValue("teamId", teams[0].id, { shouldDirty: true, shouldValidate: true }); form.setValue("customerId", "", { shouldDirty: true, shouldValidate: true }); - await form.trigger(["teamId", "customerId", "entityType"]); + form.setValue("accessProfileId", "", { shouldDirty: true, shouldValidate: true }); + await form.trigger(["teamId", "customerId", "accessProfileId", "entityType"]); } else if (val === "customer" && customers?.length > 0) { form.setValue("customerId", customers[0].id, { shouldDirty: true, shouldValidate: true }); form.setValue("teamId", "", { shouldDirty: true, shouldValidate: true }); - await form.trigger(["teamId", "customerId", "entityType"]); + form.setValue("accessProfileId", "", { shouldDirty: true, shouldValidate: true }); + await form.trigger(["teamId", "customerId", "accessProfileId", "entityType"]); + } else if (val === "access_profile") { + const apId = accessProfiles?.length > 0 + ? String(accessProfiles[0].id) + : virtualKey?.access_profile_id ? String(virtualKey.access_profile_id) : ""; + form.setValue("accessProfileId", apId, { shouldDirty: true, shouldValidate: true }); + form.setValue("teamId", "", { shouldDirty: true, shouldValidate: true }); + form.setValue("customerId", "", { shouldDirty: true, shouldValidate: true }); + await form.trigger(["teamId", "customerId", "accessProfileId", "entityType"]); } else { form.setValue("teamId", "", { shouldDirty: true, shouldValidate: true }); form.setValue("customerId", "", { shouldDirty: true, shouldValidate: true }); - await form.trigger(["teamId", "customerId", "entityType"]); + form.setValue("accessProfileId", "", { shouldDirty: true, shouldValidate: true }); + await form.trigger(["teamId", "customerId", "accessProfileId", "entityType"]); } }} - disabled={isTeamLocked} + disabled={isTeamLocked || isAPLocked} disableSearch hideClear className="h-9" @@ -1502,6 +1667,37 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT )} /> )} + + {form.watch("entityType") === "access_profile" && (accessProfiles?.length > 0 || virtualKey?.access_profile_id || defaultAccessProfileId) && ( + ( + + Select Access Profile + ({ value: String(ap.id), label: ap.name }))} + value={field.value || null} + onValueChange={(val) => { + const newVal = val ?? ""; + if (isEditing && virtualKey?.access_profile_id && newVal && newVal !== String(virtualKey.access_profile_id)) { + setPendingAccessProfileId(newVal); + setShowReassignAPWarning(true); + } else { + field.onChange(newVal); + } + }} + placeholder="Select an access profile" + disabled={isAPLocked} + emptyMessage="No access profiles found." + className="h-9" + /> + + + )} + /> + )}
diff --git a/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx b/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx index 807468f5672..db70d96d4ec 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx @@ -53,6 +53,7 @@ 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, @@ -83,7 +84,7 @@ const formatResetDuration = (duration: string) => type ExportScope = "current_page" | "all"; -function virtualKeysToCSV(vks: VirtualKey[]): string { +function virtualKeysToCSV(vks: VirtualKey[], accessProfileNames: Record = {}): string { const headers = [ "Name", "Status", @@ -112,7 +113,9 @@ function virtualKeysToCSV(vks: VirtualKey[]): string { ? `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("; ") : ""; @@ -366,6 +369,15 @@ export default function VirtualKeysTable({ const [fetchVirtualKeys, { isFetching: isExporting }] = useLazyGetVirtualKeysQuery(); + const { data: accessProfilesData } = useGetAccessProfilesQuery({ limit: 100 }); + const accessProfileNames = useMemo(() => { + const map: Record = {}; + 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( () => @@ -480,7 +492,7 @@ export default function VirtualKeysTable({ const handleExportCSV = async () => { if (exportScope === "current_page") { - downloadCSV(virtualKeysToCSV(virtualKeys)); + downloadCSV(virtualKeysToCSV(virtualKeys, accessProfileNames)); toast.success(`Exported ${virtualKeys.length} virtual keys`); setShowExportDialog(false); return; @@ -504,7 +516,7 @@ export default function VirtualKeysTable({ export: true, }).unwrap(); - downloadCSV(virtualKeysToCSV(result.virtual_keys)); + downloadCSV(virtualKeysToCSV(result.virtual_keys, accessProfileNames)); toast.success(`Exported ${result.virtual_keys.length} virtual keys`); setShowExportDialog(false); } catch (error) { @@ -829,6 +841,13 @@ export default function VirtualKeysTable({ > Customer: {vk.customer.name} + ) : vk.access_profile_id ? ( + + AP: {accessProfileNames[vk.access_profile_id] ?? vk.access_profile_id} + ) : ( - diff --git a/ui/components/ui/combobox.tsx b/ui/components/ui/combobox.tsx index 3cda97d1679..35dfd6a17e1 100644 --- a/ui/components/ui/combobox.tsx +++ b/ui/components/ui/combobox.tsx @@ -295,6 +295,7 @@ interface ComboboxSelectBaseProps { hideClear?: boolean; className?: string; emptyMessage?: string; + "data-testid"?: string; } interface ComboboxSelectSingleProps extends ComboboxSelectBaseProps { @@ -320,6 +321,7 @@ function ComboboxSelect(props: ComboboxSelectProps) { className, emptyMessage = "No results found.", noPortal, + "data-testid": dataTestId, } = props; const [open, setOpen] = React.useState(false); @@ -442,6 +444,7 @@ 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", diff --git a/ui/lib/store/apis/baseApi.ts b/ui/lib/store/apis/baseApi.ts index 6214e434b0a..2113698735a 100644 --- a/ui/lib/store/apis/baseApi.ts +++ b/ui/lib/store/apis/baseApi.ts @@ -163,6 +163,7 @@ export const baseApi = createApi({ "Versions", "Sessions", "AccessProfiles", + "AccessProfileVirtualKeys", "BusinessUnits", "PromptDeployments", "AuthType", diff --git a/ui/lib/store/apis/governanceApi.ts b/ui/lib/store/apis/governanceApi.ts index c22932ec6df..f446204e960 100644 --- a/ui/lib/store/apis/governanceApi.ts +++ b/ui/lib/store/apis/governanceApi.ts @@ -61,6 +61,7 @@ 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", }), diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index 695b0eccf06..b47dd8d355a 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -72,6 +72,7 @@ export interface VirtualKey { mcp_configs?: VirtualKeyMCPConfig[]; team_id?: string; customer_id?: string; + access_profile_id?: number; rate_limit_id?: string; is_active: boolean; calendar_aligned?: boolean; @@ -157,6 +158,7 @@ export interface CreateVirtualKeyRequest { mcp_configs?: VirtualKeyMCPConfigRequest[]; team_id?: string; customer_id?: string; + access_profile_id?: number; budgets?: CreateBudgetRequest[]; rate_limit?: CreateRateLimitRequest; is_active?: boolean; @@ -170,6 +172,7 @@ export interface UpdateVirtualKeyRequest { mcp_configs?: VirtualKeyMCPConfigRequest[]; team_id?: string; customer_id?: string; + access_profile_id?: number; budgets?: CreateBudgetRequest[]; rate_limit?: UpdateRateLimitRequest; is_active?: boolean; @@ -241,6 +244,7 @@ export interface GetVirtualKeysParams { search?: string; customer_id?: string; team_id?: string; + access_profile_id?: number; exclude_access_profile_managed_virtual?: boolean; sort_by?: "name" | "budget_spent" | "created_at" | "status"; order?: "asc" | "desc";