feat: add expiration support to virtual keys - #4882
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds virtual key expiration support end-to-end: a new database migration and ChangesVirtual Key Expiry
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant GovernancePlugin
participant ConfigStore
Client->>GovernancePlugin: PreRequestHook / PreMCPHook
GovernancePlugin->>ConfigStore: fetch virtual key
ConfigStore-->>GovernancePlugin: virtual key record
GovernancePlugin->>GovernancePlugin: IsExpiredAt(now)
GovernancePlugin-->>Client: DecisionVirtualKeyBlocked (403, "Virtual key has expired")
sequenceDiagram
participant User
participant VirtualKeySheet
participant ExpiryPickerField
participant GovernanceAPI
User->>ExpiryPickerField: pick date or clear
ExpiryPickerField->>VirtualKeySheet: update expiresAt field
User->>VirtualKeySheet: submit form
VirtualKeySheet->>VirtualKeySheet: check dirtyFields.expiresAt
VirtualKeySheet->>GovernanceAPI: send expires_at or clear_expires_at
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Comment |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if !ok || virtualKey == nil || virtualKey.IsExpiredAt(time.Now().UTC()) { | ||
| return nil | ||
| } |
There was a problem hiding this comment.
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
}| virtualKeys.map((vk) => { | ||
| const isRevealed = revealedKeys.has(vk.id); |
There was a problem hiding this comment.
Both virtualKeysTable.tsx and virtualKeyDetailsSheet.tsx compute isExpired with Date.now() during the render pass. A key that expires minutes after the table loads will continue showing "Active" until the component re-renders (e.g. navigation or manual refresh). Consider either a short polling interval or a useMemo/useEffect pair that schedules a state flip at new Date(vk.expires_at).getTime() - Date.now() milliseconds so the badge transitions automatically.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/governance/main.go (1)
1203-1210: 🔒 Security & Privacy | 🟡 Minor | ⚡ Quick winRestore the inactive VK guard in
PreRequestHook. Inactive keys now reach routing, allowlist stamping, and MCP injection beforePreLLMHookrejects them. Keep this hook aligned withPreMCPHookby short-circuiting on!IsActiveValue()before any VK-scoped setup.🤖 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 `@plugins/governance/main.go` around lines 1203 - 1210, The PreRequestHook virtual key lookup currently only filters missing or expired keys, so inactive keys can still proceed into routing, allowlist stamping, and MCP setup. Update the PreRequestHook logic around p.store.GetVirtualKey and the virtualKey.IsExpiredAt check to also short-circuit when the key is not active by using IsActiveValue(), keeping it aligned with PreMCPHook and ensuring VK-scoped setup only happens for active keys.Source: Path instructions
🧹 Nitpick comments (4)
ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx (2)
189-223: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
data-testidto the new Expiry picker's interactive elements.
ExpiryPickerFieldintroduces several new interactive controls (Clear button, Never button, preset buttons,DateTimePicker) with nodata-testidattributes. E2E tests will need stable selectors for these.Based on learnings, `data-testid` conventions of `--` should be applied consistently; adjust names to match existing `vk-*` patterns in this file.🏷️ Suggested testids
- {value && ( - <Button type="button" variant="ghost" size="sm" onClick={() => onChange(null)}> + {value && ( + <Button type="button" variant="ghost" size="sm" data-testid="vk-expiry-clear-btn" 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)}> + <Button type="button" variant={!value ? "secondary" : "outline"} size="sm" data-testid="vk-expiry-never-btn" onClick={() => onChange(null)}> Never </Button> {EXPIRY_PRESETS.map(({ label, ms }) => ( - <Button key={label} type="button" variant="outline" size="sm" onClick={() => onChange(presetFromNow(ms))}> + <Button key={label} type="button" variant="outline" size="sm" data-testid={`vk-expiry-preset-${label}`} onClick={() => onChange(presetFromNow(ms))}> {label} </Button> ))} <DateTimePicker + data-testid="vk-expiry-datetime-picker" buttonClassName="h-8 text-sm px-3"🤖 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 `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 189 - 223, Add stable data-testid attributes to the new interactive controls in ExpiryPickerField so E2E tests can target them reliably. Update the Clear button, Never button, each preset button in EXPIRY_PRESETS, and the DateTimePicker with vk-* testid names that match the existing virtual key naming pattern in this file. Use consistent entity-element-qualifier-style identifiers so the selectors remain predictable even if the UI layout changes.Source: Coding guidelines
213-218: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winNo client-side guard against picking a past time on the current day.
disabledBefore={new Date()}only disables prior days in the calendar; the pairedTimePickerstill allows selecting an already-passed time on today's date. The backend rejects this with a 400 (expires_at must be a future timestamp), so submission fails, surfaced only via a generic toast. Adding a same-day future-time check (zod.refineor inline validation) would give the user immediate, actionable feedback instead of a round-trip failure.🤖 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 `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx` around lines 213 - 218, The DateTimePicker flow in virtualKeySheet.tsx only blocks past days, so users can still pick an earlier time today and hit the backend 400. Add client-side validation in the virtual key expiration form path around DateTimePicker/onDateTimeUpdate (or the zod schema/refine used by the sheet) to reject any timestamp that is not strictly in the future, including same-day times. Surface an inline validation message before submit so the user gets immediate feedback instead of a generic toast.ui/lib/types/governance.ts (1)
183-184: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win
expires_atonUpdateVirtualKeyRequestshouldn't allownull.Unlike
team_id/customer_id(typed asschemas.OptionalJSON[string]on the backend to distinguish "not provided" from "explicit null/clear"), the backendExpiresAtfield is a plain*time.Timewithomitempty:ExpiresAt *time.Time json:"expires_at,omitempty" // Set a new expiry; nil means "leave unchanged". During JSON unmarshaling, sendingexpires_at: nullis indistinguishable from omitting the field entirely — both decode to a nil pointer, meaning "leave unchanged", not "clear". Clearing is only achievable viaclear_expires_at: true. Allowingstring | nullhere invites a future caller to (incorrectly) assumenullclears the expiry.Current call sites in
virtualKeySheet.tsxcorrectly avoid sendingnull(they useclear_expires_at), so this isn't exploited today — but the type should reflect the real contract.🛠️ Suggested type tightening
- expires_at?: string | null; // ISO 8601 UTC timestamp; omit to leave unchanged + expires_at?: string; // ISO 8601 UTC timestamp to set a new expiry; use clear_expires_at to remove clear_expires_at?: boolean; // true to remove an existing expiry (mutually exclusive with expires_at)🤖 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 `@ui/lib/types/governance.ts` around lines 183 - 184, The `UpdateVirtualKeyRequest` type allows `expires_at` to be `null`, but the backend `ExpiresAt` contract only treats omission as “leave unchanged” and uses `clear_expires_at` to clear the value. Tighten the `expires_at` field in `governance.ts` to accept only a string timestamp (or be omitted), and keep `clear_expires_at` as the only way to remove an expiry. Check related callers like `virtualKeySheet.tsx` to ensure they continue using `clear_expires_at` instead of relying on `null`.ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx (1)
904-912: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a
data-testidto the new Expired badge; replacing the active-switch testid may break E2E selectors.When
showExpiredBadgeis true, the row no longer rendersVKActiveSwitch(data-testid="vk-active-switch-${vk.name}"), and the newBadgehas no testid of its own. Any E2E test asserting on the active-switch selector for all rows, or needing to detect the "Expired" state, has no stable hook here.Based on path instructions, "UI `data-testid` attributes are load-bearing for E2E tests" and removal/replacement of an existing testid requires updating `tests/e2e/` references.🏷️ Suggested fix
{showExpiredBadge ? ( - <Badge variant="destructive" className="text-xs"> + <Badge variant="destructive" className="text-xs" data-testid={`vk-expired-badge-${vk.name}`}> Expired </Badge> ) : ( <VKActiveSwitch vk={vk} hasUpdateAccess={hasUpdateAccess} onToggle={handleToggleActive} /> )}🤖 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 `@ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx` around lines 904 - 912, The new Expired badge in virtualKeysTable should get its own stable data-testid because it replaces VKActiveSwitch and removes the existing vk-active-switch-${vk.name} hook for expired rows. Update the Badge branch in the showExpiredBadge conditional to add a unique testid for the expired state, and then adjust any affected tests in tests/e2e/ that currently rely on the active-switch selector so they assert the correct state using the new Badge identifier. Use VKActiveSwitch and the Badge render path as the key locations.Source: Path instructions
🤖 Prompt for all review comments with 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.
Inline comments:
In `@transports/bifrost-http/handlers/governance.go`:
- 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.
---
Outside diff comments:
In `@plugins/governance/main.go`:
- Around line 1203-1210: The PreRequestHook virtual key lookup currently only
filters missing or expired keys, so inactive keys can still proceed into
routing, allowlist stamping, and MCP setup. Update the PreRequestHook logic
around p.store.GetVirtualKey and the virtualKey.IsExpiredAt check to also
short-circuit when the key is not active by using IsActiveValue(), keeping it
aligned with PreMCPHook and ensuring VK-scoped setup only happens for active
keys.
---
Nitpick comments:
In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 189-223: Add stable data-testid attributes to the new interactive
controls in ExpiryPickerField so E2E tests can target them reliably. Update the
Clear button, Never button, each preset button in EXPIRY_PRESETS, and the
DateTimePicker with vk-* testid names that match the existing virtual key naming
pattern in this file. Use consistent entity-element-qualifier-style identifiers
so the selectors remain predictable even if the UI layout changes.
- Around line 213-218: The DateTimePicker flow in virtualKeySheet.tsx only
blocks past days, so users can still pick an earlier time today and hit the
backend 400. Add client-side validation in the virtual key expiration form path
around DateTimePicker/onDateTimeUpdate (or the zod schema/refine used by the
sheet) to reject any timestamp that is not strictly in the future, including
same-day times. Surface an inline validation message before submit so the user
gets immediate feedback instead of a generic toast.
In `@ui/app/workspace/virtual-keys/views/virtualKeysTable.tsx`:
- Around line 904-912: The new Expired badge in virtualKeysTable should get its
own stable data-testid because it replaces VKActiveSwitch and removes the
existing vk-active-switch-${vk.name} hook for expired rows. Update the Badge
branch in the showExpiredBadge conditional to add a unique testid for the
expired state, and then adjust any affected tests in tests/e2e/ that currently
rely on the active-switch selector so they assert the correct state using the
new Badge identifier. Use VKActiveSwitch and the Badge render path as the key
locations.
In `@ui/lib/types/governance.ts`:
- Around line 183-184: The `UpdateVirtualKeyRequest` type allows `expires_at` to
be `null`, but the backend `ExpiresAt` contract only treats omission as “leave
unchanged” and uses `clear_expires_at` to clear the value. Tighten the
`expires_at` field in `governance.ts` to accept only a string timestamp (or be
omitted), and keep `clear_expires_at` as the only way to remove an expiry. Check
related callers like `virtualKeySheet.tsx` to ensure they continue using
`clear_expires_at` instead of relying on `null`.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: de2d686b-666b-446b-a2e9-dc81bc1ee68e
📒 Files selected for processing (10)
framework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/tables/virtualkey.goplugins/governance/main.goplugins/governance/resolver.gotransports/bifrost-http/handlers/governance.goui/app/workspace/virtual-keys/views/virtualKeyDetailsSheet.tsxui/app/workspace/virtual-keys/views/virtualKeySheet.tsxui/app/workspace/virtual-keys/views/virtualKeysTable.tsxui/lib/types/governance.ts
| 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 |
There was a problem hiding this comment.
🗄️ 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.jsonRepository: 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)])
PYRepository: 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
2 participants
Summary
Adds optional expiry support for virtual keys. Operators can now set an
expires_attimestamp on a virtual key; once that time passes, the key is automatically treated as blocked without needing to manually deactivate it.Changes
expires_atnullable timestamp column togovernance_virtual_keysvia a new database migration, with rollback support.IsExpiredAt(now time.Time) boolmethod onTableVirtualKeyfor in-memory expiry checks (no DB index needed since expiry is evaluated from the already-loaded key).UpdateVirtualKeyto includeexpires_atin the set of persisted fields.EvaluateVirtualKeyRequestin the budget resolver to block expired keys with aDecisionVirtualKeyBlockedresult.PreRequestHookto reject expired keys at the pre-request stage.PreMCPHookto separately checkIsActiveValue()andIsExpiredAt(), returning distinct 403 error messages for inactive vs. expired keys.ExpiresAttoCreateVirtualKeyRequestandUpdateVirtualKeyRequestHTTP handler types, with validation that the timestamp must be in the future.ClearExpiresAtboolean toUpdateVirtualKeyRequestto explicitly remove an existing expiry, mutually exclusive with setting a newexpires_at.ExpiryPickerFieldcomponent with quick-select presets (30 min, 1 hour, 24 hours, 7 days) and a custom date-time picker.Type of change
Affected areas
How to test
expires_atset to a time 1 minute in the future. Confirm requests succeed before expiry."Virtual key has expired".expires_atin the past — confirm the API returns a 400 validation error.clear_expires_at: trueand confirm the expiry is removed.expires_atandclear_expires_at: truein the same update request — confirm a 400 error is returned.Screenshots/Recordings
Add before/after screenshots of the virtual key table and detail sheet showing the "Expired" badge and expiry field.
Breaking changes
Related issues
Security considerations
Expired keys are rejected at both the pre-request and pre-MCP hook stages, ensuring fail-closed behavior. Expiry is evaluated in-memory from the loaded key object, so there is no risk of a time-of-check/time-of-use gap introduced by additional DB queries. The
clear_expires_atandexpires_atfields are mutually exclusive to prevent ambiguous update semantics.Checklist
docs/contributing/README.mdand followed the guidelines