Skip to content

feat: add expiration support to virtual keys #4882

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Pratham-Mishra04 wants to merge 1 commit into from
Closed

feat: add expiration support to virtual keys #4882

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
24 changes: 24 additions & 0 deletions framework/configstore/migrations.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
}
2 changes: 1 addition & 1 deletion framework/configstore/rdb.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down
10 changes: 10 additions & 0 deletions framework/configstore/tables/virtualkey.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"`

Expand Down Expand Up @@ -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.
Expand Down
24 changes: 22 additions & 2 deletions plugins/governance/main.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Comment on lines +1207 to 1209

@greptile-apps greptile-apps Bot Jul 3, 2026

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.

P1 Inactive VK check dropped from PreRequestHook short-circuit

The condition previously returned early for inactive keys (!virtualKey.IsActiveValue()); this PR replaces that with the expiry check but doesn't preserve the inactive check. As a result, inactive (but not expired) virtual keys now fall through the guard and enter the routing pipeline — stampGovernanceCtxFromVK, routing-rule evaluation, loadBalanceProvider, and MCP tool-allowlist computation all execute before PreLLMHook eventually blocks the key via EvaluateVirtualKeyRequest. Load-balancer counters and MCP context values can be mutated for requests that will never be served. Both checks should be in the condition:

if !ok || virtualKey == nil || !virtualKey.IsActiveValue() || virtualKey.IsExpiredAt(time.Now().UTC()) {
    return nil
}

}
Expand Down Expand Up @@ -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{
Expand All @@ -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{
Expand Down
7 changes: 7 additions & 0 deletions plugins/governance/resolver.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ package governance
import (
"context"
"fmt"
"time"

"github.com/maximhq/bifrost/core/schemas"
configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables"
Expand Down Expand Up @@ -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{
Expand Down
30 changes: 30 additions & 0 deletions transports/bifrost-http/handlers/governance.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@coderabbitai coderabbitai Bot Jul 3, 2026

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.

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the relevant handler and schema sections with line numbers.
git ls-files transports/bifrost-http/handlers/governance.go transports/config.schema.json

echo "---- governance.go ----"
sed -n '140,230p' transports/bifrost-http/handlers/governance.go

echo "---- config.schema.json search ----"
rg -n '"expires_at"|"clear_expires_at"|virtual_keys|additionalProperties' transports/config.schema.json

Repository: maximhq/bifrost

Length of output: 16596

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Show the exact virtual_key object schema section if present.
python3 - <<'PY'
from pathlib import Path
p = Path('transports/config.schema.json')
text = p.read_text()
for needle in ['"virtual_keys"', '"expires_at"', '"clear_expires_at"']:
    idx = text.find(needle)
    print(f"\n== {needle} ==")
    if idx == -1:
        print("not found")
        continue
    start = max(0, text.rfind('\n', 0, idx-500))
    end = text.find('\n', idx+1000)
    print(text[start:end if end!=-1 else len(text)])
PY

Repository: maximhq/bifrost

Length of output: 1776

Add expires_at and clear_expires_at to transports/config.schema.json. governance.virtual_keys[] is still additionalProperties: false and doesn’t define either field, so config-backed virtual keys using this new expiry support will fail schema validation.

🤖 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` at line 168, The governance
virtual key model now includes ExpiresAt, but the config schema still rejects it
because governance.virtual_keys[] is locked down with additionalProperties
false. Update transports/config.schema.json to explicitly add expires_at and
clear_expires_at to the virtual key object definition, using the same naming and
semantics as the governance types so config-backed keys validate correctly.

Source: Path instructions

}

// UpdateVirtualKeyRequest represents the request body for updating a virtual key
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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
}
Expand Down
21 changes: 18 additions & 3 deletions ui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,12 +136,27 @@ export default function VirtualKeyDetailSheet({
<div className="grid grid-cols-3 items-center gap-4">
<span className="text-muted-foreground text-sm">Status</span>
<div className="col-span-2">
<Badge variant={virtualKey.is_active ? (isExhausted ? "destructive" : "default") : "secondary"}>
{virtualKey.is_active ? (isExhausted ? "Exhausted" : "Active") : "Inactive"}
</Badge>
{(() => {
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 <Badge variant={variant}>{label}</Badge>;
})()}
</div>
</div>

{virtualKey.expires_at && (
<div className="grid grid-cols-3 items-center gap-4">
<span className="text-muted-foreground text-sm">Expires</span>
<div className="col-span-2 text-sm">
{formatDistanceToNow(new Date(virtualKey.expires_at), {
addSuffix: true,
})}
<span className="text-muted-foreground ml-1 text-xs">({new Date(virtualKey.expires_at).toLocaleString()})</span>
</div>
</div>
)}

<div className="grid grid-cols-3 items-center gap-4">
<span className="text-muted-foreground text-sm">Created</span>
<div className="col-span-2 text-sm">
Expand Down
84 changes: 84 additions & 0 deletions ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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";
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 (
<FormItem>
<div className="flex items-center justify-between">
<FormLabel>Expiry</FormLabel>
{value && (
<Button type="button" variant="ghost" size="sm" onClick={() => onChange(null)}>
Clear
</Button>
)}
</div>
<p className="text-muted-foreground text-xs">Leave empty for a key that never expires.</p>
{summary && <p className="text-sm font-medium">{summary}</p>}
<div className="flex flex-wrap gap-1.5">
<Button type="button" variant={!value ? "secondary" : "outline"} size="sm" onClick={() => onChange(null)}>
Never
</Button>
{EXPIRY_PRESETS.map(({ label, ms }) => (
<Button key={label} type="button" variant="outline" size="sm" onClick={() => onChange(presetFromNow(ms))}>
{label}
</Button>
))}
<DateTimePicker
buttonClassName="h-8 text-sm px-3"
dateTime={value ? new Date(value) : undefined}
disabledBefore={new Date()}
onDateTimeUpdate={(dt) => onChange(toDatetimeLocal(dt))}
/>
</div>
<FormMessage />
</FormItem>
);
}

export default function VirtualKeySheet({ virtualKey, teams, customers, defaultTeamId, onSave, onCancel }: VirtualKeySheetProps) {
const [isOpen, setIsOpen] = useState(true);
const navigate = useNavigate();
Expand Down Expand Up @@ -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) => ({
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -870,6 +949,11 @@ export default function VirtualKeySheet({ virtualKey, teams, customers, defaultT
</FormItem>
)}
/>
<FormField
control={form.control}
name="expiresAt"
render={({ field }) => <ExpiryPickerField value={field.value} onChange={field.onChange} />}
/>
</div>
{/* Provider Configurations */}
<div className="space-y-2">
Expand Down
Loading
Loading