From 70041ea52911272f5b2102a59120eef913e231b5 Mon Sep 17 00:00:00 2001 From: Vaibhav mittal Date: Fri, 3 Jul 2026 15:29:22 +0530 Subject: [PATCH] fix(governance): support expiring virtual keys --- framework/configstore/migrations.go | 24 ++++++ framework/configstore/rdb.go | 2 +- framework/configstore/tables/virtualkey.go | 10 +++ plugins/governance/main.go | 24 +++++- plugins/governance/resolver.go | 7 ++ .../bifrost-http/handlers/governance.go | 30 +++++++ .../views/virtualKeyDetailsSheet.tsx | 21 ++++- .../virtual-keys/views/virtualKeySheet.tsx | 84 +++++++++++++++++++ .../virtual-keys/views/virtualKeysTable.tsx | 13 ++- ui/lib/types/governance.ts | 4 + 10 files changed, 211 insertions(+), 8 deletions(-) diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 101cab4e53d..561b6bdbf33 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -435,6 +435,7 @@ var configstoreMigrationSteps = []migrationStep{ {IDs: []string{"add_bedrock_mantle_key_columns"}, run: migrationAddBedrockMantleKeyColumns}, {IDs: []string{"add_model_pricing_is_deprecated_column"}, run: migrationAddModelPricingIsDeprecatedColumn}, {IDs: []string{"add_mcp_client_tool_execution_timeout_column"}, run: migrationAddMCPClientToolExecutionTimeoutColumn}, + {IDs: []string{"add_virtual_key_expires_at_column"}, run: migrationAddVirtualKeyExpiresAtColumn}, } // quoteSQLiteIdentifier quotes a SQLite identifier, escaping any double quotes. @@ -10292,3 +10293,26 @@ func migrationAddMCPClientToolExecutionTimeoutColumn(ctx context.Context, db *go } return nil } + +// migrationAddVirtualKeyExpiresAtColumn adds nullable expires_at to governance_virtual_keys. +// No index: expiry is checked in-memory from the already-loaded VK, never queried by column. +func migrationAddVirtualKeyExpiresAtColumn(ctx context.Context, db *gorm.DB, logger schemas.Logger) error { + migrationName := "add_virtual_key_expires_at_column" + logger.Info("[configstore] starting migration %s", migrationName) + defer logger.Info("[configstore] finished migration %s", migrationName) + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: migrationName, + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + return addColumnIfNotExists(tx, logger, &tables.TableVirtualKey{}, "expires_at") + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + return dropColumnIfExists(tx, logger, &tables.TableVirtualKey{}, "expires_at") + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running %s migration: %w", migrationName, err) + } + return nil +} diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index a3c47668759..a3946895700 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -3376,7 +3376,7 @@ func (s *RDBConfigStore) UpdateVirtualKey(ctx context.Context, virtualKey *table } else { virtualKey.ID = existing.ID if err := txDB.WithContext(ctx). - Select("name", "description", "value", "is_active", "team_id", "customer_id", "rate_limit_id", "calendar_aligned", "config_hash", "updated_at", "encryption_status", "value_hash"). + Select("name", "description", "value", "is_active", "expires_at", "team_id", "customer_id", "rate_limit_id", "calendar_aligned", "config_hash", "updated_at", "encryption_status", "value_hash"). Updates(virtualKey).Error; err != nil { return s.parseGormError(err) } diff --git a/framework/configstore/tables/virtualkey.go b/framework/configstore/tables/virtualkey.go index 33e3bf2466d..217b5a73c60 100644 --- a/framework/configstore/tables/virtualkey.go +++ b/framework/configstore/tables/virtualkey.go @@ -211,6 +211,7 @@ type TableVirtualKey struct { Description string `gorm:"type:text" json:"description,omitempty"` Value schemas.SecretVar `gorm:"uniqueIndex:idx_virtual_key_value;type:text;not null" json:"value"` IsActive *bool `gorm:"default:true" json:"is_active,omitempty"` // Nil means true (DB default); false means inactive + ExpiresAt *time.Time `gorm:"type:timestamp;null" json:"expires_at,omitempty"` // Optional expiry; nil means never expires 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"` @@ -276,6 +277,15 @@ func (vk TableVirtualKey) MarshalJSON() ([]byte, error) { }) } +// IsExpiredAt reports whether the virtual key has passed its expiry. +// now == expires_at is treated as expired; nil ExpiresAt means never expires. +func (vk *TableVirtualKey) IsExpiredAt(now time.Time) bool { + if vk == nil || vk.ExpiresAt == nil { + return false + } + return !now.UTC().Before(vk.ExpiresAt.UTC()) +} + // 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. diff --git a/plugins/governance/main.go b/plugins/governance/main.go index 6e383953dc1..b1e4277dc29 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -1204,7 +1204,7 @@ func (p *GovernancePlugin) PreRequestHook(ctx *schemas.BifrostContext, req *sche if virtualKeyValue != "" { var ok bool virtualKey, ok = p.store.GetVirtualKey(ctx, virtualKeyValue) - if !ok || virtualKey == nil || !virtualKey.IsActiveValue() { + if !ok || virtualKey == nil || virtualKey.IsExpiredAt(time.Now().UTC()) { return nil } } @@ -1440,7 +1440,7 @@ func (p *GovernancePlugin) PreMCPHook(ctx *schemas.BifrostContext, req *schemas. // This runs independently of EvaluateGovernanceRequest to enforce execution-time allow-list. if virtualKeyValue != "" { vk, ok := p.store.GetVirtualKey(ctx, virtualKeyValue) - if !ok || vk == nil || !vk.IsActiveValue() { + if !ok || vk == nil { // VK became invalid after initial check - fail closed for security ctx.SetValue(governanceRejectedContextKey, true) return req, &schemas.MCPPluginShortCircuit{Error: &schemas.BifrostError{ @@ -1451,6 +1451,26 @@ func (p *GovernancePlugin) PreMCPHook(ctx *schemas.BifrostContext, req *schemas. }, }}, nil } + if !vk.IsActiveValue() { + ctx.SetValue(governanceRejectedContextKey, true) + return req, &schemas.MCPPluginShortCircuit{Error: &schemas.BifrostError{ + Type: bifrost.Ptr(string(DecisionVirtualKeyBlocked)), + StatusCode: bifrost.Ptr(403), + Error: &schemas.ErrorField{ + Message: "Virtual key is inactive", + }, + }}, nil + } + if vk.IsExpiredAt(time.Now().UTC()) { + ctx.SetValue(governanceRejectedContextKey, true) + return req, &schemas.MCPPluginShortCircuit{Error: &schemas.BifrostError{ + Type: bifrost.Ptr(string(DecisionVirtualKeyBlocked)), + StatusCode: bifrost.Ptr(403), + Error: &schemas.ErrorField{ + Message: "Virtual key has expired", + }, + }}, nil + } if !p.isMCPToolAllowedByVK(vk, toolName) { ctx.SetValue(governanceRejectedContextKey, true) return req, &schemas.MCPPluginShortCircuit{Error: &schemas.BifrostError{ diff --git a/plugins/governance/resolver.go b/plugins/governance/resolver.go index 59cbc0aede7..b8101b1cb4d 100644 --- a/plugins/governance/resolver.go +++ b/plugins/governance/resolver.go @@ -4,6 +4,7 @@ package governance import ( "context" "fmt" + "time" "github.com/maximhq/bifrost/core/schemas" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" @@ -271,6 +272,12 @@ func (r *BudgetResolver) EvaluateVirtualKeyRequest(ctx *schemas.BifrostContext, Reason: "Virtual key is inactive", } } + if vk.IsExpiredAt(time.Now().UTC()) { + return &EvaluationResult{ + Decision: DecisionVirtualKeyBlocked, + Reason: "Virtual key has expired", + } + } // 2. Check provider filtering if requestType != schemas.MCPToolExecutionRequest && requestType != schemas.ListModelsRequest && !r.isProviderAllowed(vk, provider) { return &EvaluationResult{ diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index 37756023a24..2c873ff0e94 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -165,6 +165,7 @@ type CreateVirtualKeyRequest struct { 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 + ExpiresAt *time.Time `json:"expires_at,omitempty"` // Optional expiry; nil means never expires } // UpdateVirtualKeyRequest represents the request body for updating a virtual key @@ -193,6 +194,8 @@ type UpdateVirtualKeyRequest struct { IsActive *bool `json:"is_active,omitempty"` CalendarAligned *bool `json:"calendar_aligned,omitempty"` // When true, all budgets reset at clean calendar boundaries ResetBudgetUsage *bool `json:"reset_budget_usage,omitempty"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` // Set a new expiry; nil means "leave unchanged" + ClearExpiresAt bool `json:"clear_expires_at,omitempty"` // true to remove an existing expiry } var errVirtualKeyDualAssociation = errors.New("VirtualKey cannot be attached to both Team and Customer") @@ -1271,6 +1274,14 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) { seenDurations[b.ResetDuration] = true } } + // Validate expires_at: must be in the future if provided + if req.ExpiresAt != nil { + now := time.Now().UTC() + if !req.ExpiresAt.After(now) { + SendError(ctx, 400, "expires_at must be a future timestamp") + return + } + } // Set defaults: nil means "use DB default (true)" isActive := req.IsActive if isActive == nil { @@ -1297,6 +1308,7 @@ func (h *GovernanceHandler) createVirtualKey(ctx *fasthttp.RequestCtx) { CustomerID: req.CustomerID, IsActive: isActive, CalendarAligned: req.CalendarAligned, + ExpiresAt: req.ExpiresAt, } if err := h.configStore.CreateVirtualKey(ctx, &vk, tx); err != nil { return err @@ -1495,6 +1507,19 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { SendError(ctx, 400, "VirtualKey cannot be attached to both Team and Customer") return } + // Validate mutually exclusive ExpiresAt and ClearExpiresAt + if req.ExpiresAt != nil && req.ClearExpiresAt { + SendError(ctx, 400, "cannot set both expires_at and clear_expires_at") + return + } + // Validate expires_at: must be in the future if provided + if req.ExpiresAt != nil { + now := time.Now().UTC() + if !req.ExpiresAt.After(now) { + SendError(ctx, 400, "expires_at must be a future timestamp") + return + } + } vk, err := h.configStore.GetVirtualKey(ctx, vkID) if err != nil { if errors.Is(err, configstore.ErrNotFound) { @@ -1551,6 +1576,11 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { if req.IsActive != nil { vk.IsActive = req.IsActive } + if req.ClearExpiresAt { + vk.ExpiresAt = nil + } else if req.ExpiresAt != nil { + vk.ExpiresAt = req.ExpiresAt + } if req.CalendarAligned != nil { vk.CalendarAligned = *req.CalendarAligned } diff --git a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx index cc0e5974861..0a4166dd0e2 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx @@ -136,12 +136,27 @@ export default function VirtualKeyDetailSheet({
Status
- - {virtualKey.is_active ? (isExhausted ? "Exhausted" : "Active") : "Inactive"} - + {(() => { + const isExpired = !!virtualKey.expires_at && Date.now() >= new Date(virtualKey.expires_at).getTime(); + const variant = !virtualKey.is_active ? "secondary" : isExpired || isExhausted ? "destructive" : "default"; + const label = !virtualKey.is_active ? "Inactive" : isExpired ? "Expired" : isExhausted ? "Exhausted" : "Active"; + return {label}; + })()}
+ {virtualKey.expires_at && ( +
+ Expires +
+ {formatDistanceToNow(new Date(virtualKey.expires_at), { + addSuffix: true, + })} + ({new Date(virtualKey.expires_at).toLocaleString()}) +
+
+ )} +
Created
diff --git a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx index 6d57f0ab196..3a614ca3b4b 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx @@ -13,6 +13,7 @@ import { } from "@/components/ui/alertDialog"; import { AsyncMultiSelect } from "@/components/ui/asyncMultiselect"; import { Button } from "@/components/ui/button"; +import { DateTimePicker } from "@/components/ui/datePickerWithRange"; import { ComboboxSelect } from "@/components/ui/combobox"; import { ConfigSyncAlert } from "@/components/ui/configSyncAlert"; import { Form, FormControl, FormField, FormItem, FormLabel, FormMessage } from "@/components/ui/form"; @@ -49,6 +50,7 @@ import { CreateVirtualKeyRequest, Customer, Team, UpdateVirtualKeyRequest, Virtu import { RbacOperation, RbacResource, useRbac } from "@enterprise/lib"; import { zodResolver } from "@hookform/resolvers/zod"; import { useNavigate } from "@tanstack/react-router"; +import { formatDistanceToNow } from "date-fns"; import { Info, Lock, RotateCcw, Trash2, Users, X } from "lucide-react"; import { useEffect, useState } from "react"; import { useForm } from "react-hook-form"; @@ -113,6 +115,7 @@ const formSchema = z teamId: z.string().optional(), customerId: z.string().optional(), isActive: z.boolean(), + expiresAt: z.string().nullable().optional(), // ISO 8601 datetime-local string, or null to clear // Budget budgetCalendarAligned: z.boolean(), budgets: z @@ -164,6 +167,61 @@ type VirtualKeyType = { provider: string; }; +const pad2 = (n: number) => n.toString().padStart(2, "0"); + +const toDatetimeLocal = (d: Date) => + `${d.getFullYear()}-${pad2(d.getMonth() + 1)}-${pad2(d.getDate())}T${pad2(d.getHours())}:${pad2(d.getMinutes())}`; + +const presetFromNow = (offsetMs: number) => toDatetimeLocal(new Date(Date.now() + offsetMs)); + +const EXPIRY_PRESETS = [ + { label: "30 min", ms: 30 * 60_000 }, + { label: "1 hour", ms: 60 * 60_000 }, + { label: "24 hours", ms: 24 * 60 * 60_000 }, + { label: "7 days", ms: 7 * 24 * 60 * 60_000 }, +] as const; + +interface ExpiryFieldProps { + value: string | null | undefined; + onChange: (v: string | null) => void; +} + +function ExpiryPickerField({ value, onChange }: ExpiryFieldProps) { + const summary = value ? formatDistanceToNow(new Date(value), { addSuffix: true }) : null; + + return ( + +
+ Expiry + {value && ( + + )} +
+

Leave empty for a key that never expires.

+ {summary &&

{summary}

} +
+ + {EXPIRY_PRESETS.map(({ label, ms }) => ( + + ))} + onChange(toDatetimeLocal(dt))} + /> +
+ + + ); +} + export default function VirtualKeySheet({ virtualKey, teams, customers, defaultTeamId, onSave, onCancel }: VirtualKeySheetProps) { const [isOpen, setIsOpen] = useState(true); const navigate = useNavigate(); @@ -241,6 +299,12 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT teamId: virtualKey?.team_id || (!isEditing ? defaultTeamId || "" : ""), customerId: virtualKey?.customer_id || "", isActive: virtualKey?.is_active ?? true, + expiresAt: virtualKey?.expires_at + ? (() => { + const d = new Date(virtualKey.expires_at); + return new Date(d.getTime() - d.getTimezoneOffset() * 60000).toISOString().slice(0, 16); + })() + : null, budgets: virtualKey?.budgets && virtualKey.budgets.length > 0 ? virtualKey.budgets.map((b) => ({ @@ -646,6 +710,18 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT : []; if (isEditing && virtualKey) { // Update existing virtual key + // Only include expiry fields when the user actually changed the expiry field. + // Pre-filled defaultValues are not dirty, so an unchanged expired key won't + // resend its old expired timestamp and cause the backend to reject the edit. + const expiryChanged = !!form.formState.dirtyFields.expiresAt; + const expiryPayload = expiryChanged + ? data.expiresAt + ? { expires_at: new Date(data.expiresAt).toISOString() } + : virtualKey?.expires_at + ? { clear_expires_at: true } + : {} + : {}; + const updateData: UpdateVirtualKeyRequest = { name: data.name, description: data.description, @@ -670,6 +746,7 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT is_active: data.isActive, calendar_aligned: data.budgetCalendarAligned, reset_budget_usage: resetBudgetUsage, + ...expiryPayload, }; // Add budgets if enabled @@ -716,6 +793,8 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT is_active: data.isActive, // VK-level setting that governs both budget and rate-limit calendar alignment. calendar_aligned: data.budgetCalendarAligned, + // Optional expiry: send as UTC ISO string, or omit for no expiry + ...(data.expiresAt ? { expires_at: new Date(data.expiresAt).toISOString() } : {}), }; // Add budgets if enabled @@ -870,6 +949,11 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT )} /> + } + />
{/* Provider Configurations */}
diff --git a/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx b/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx index 7290a8baacf..61f07eedd9f 100644 --- a/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx +++ b/ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx @@ -80,7 +80,8 @@ function virtualKeysToCSV(vks: VirtualKey[], accessProfileNames: Record= vk.rate_limit.request_max_limit); - const status = vk.is_active ? (isExhausted ? "Exhausted" : "Active") : "Inactive"; + const isExpired = !!vk.expires_at && Date.now() >= new Date(vk.expires_at).getTime(); + const status = !vk.is_active ? "Inactive" : isExpired ? "Expired" : isExhausted ? "Exhausted" : "Active"; const assignedTo = vk.team ? `Team: ${vk.team.name}` : vk.customer ? `Customer: ${vk.customer.name}` : ""; 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("; ") : ""; @@ -845,6 +846,8 @@ export default function VirtualKeysTable({ ) : ( virtualKeys.map((vk) => { const isRevealed = revealedKeys.has(vk.id); + const isExpired = !!vk.expires_at && Date.now() >= new Date(vk.expires_at).getTime(); + const showExpiredBadge = vk.is_active && isExpired; return ( e.stopPropagation()}> - + {showExpiredBadge ? ( + + Expired + + ) : ( + + )}