fix: fixes calendar_aligned migration and UI for VK

[roroghost17]

Summary

Calendar alignment (calendar_aligned) is now a VK-only setting stored on governance_virtual_keys. Previously, governance_budgets and governance_rate_limits each carried their own calendar_aligned column, which led to inconsistencies after schema migrations. This PR removes those legacy columns, derives the alignment value from the owning virtual key at reset time, and extends the UI toggle to cover rate limits in addition to budgets.

Changes

• DB migration (migrationDropLegacyCalendarAlignedColumns): Drops calendar_aligned from governance_budgets and governance_rate_limits using DROP COLUMN IF EXISTS so it is safe to run on any DB state.
• Rate-limit struct (TableRateLimit): Removed the CalendarAligned field since the value is now sourced from the owning VK.
• Reset logic (ResetExpiredRateLimitsInMemory): Builds a reverse map from rate-limit ID → owning VK's CalendarAligned value at reset time, covering both VK-attached and provider-config-attached rate limits. Rate limits not reachable from any VK remain non-aligned.
• VK update query (UpdateVirtualKey): Replaced budget_id with calendar_aligned in the Select list so the VK-level alignment flag is persisted on update.
• UI (virtualKeySheet.tsx): The calendar alignment toggle and its warning dialog are now shown when any budget or rate limit has a calendar-alignable duration. The toggle and dialog were relocated to a shared section below both the budget and rate-limit fields. Warning copy updated to mention both budget and rate-limit usage resets.

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

# Core/Transports
go test ./framework/configstore/... ./plugins/governance/...

# UI
cd ui
pnpm i
pnpm build

1. Create a virtual key with a rate limit using a daily or longer reset duration.
2. Verify the "Align to calendar cycle" toggle appears in the rate-limit section.
3. Enable calendar alignment, save, and confirm the rate-limit counters reset and snap to the period boundary.
4. Confirm the toggle also appears (and works) when only budgets with alignable durations are configured.
5. Run the migration against a DB that previously ran migrate_calendar_aligned and confirm the calendar_aligned columns are dropped from governance_budgets and governance_rate_limits without error.

Breaking changes

• Yes
• No

The calendar_aligned column is dropped from governance_budgets and governance_rate_limits. Any code or query that references those columns directly will break. The migration handles existing databases safely via DROP COLUMN IF EXISTS.

Related issues

Security considerations

None. No auth, secrets, or PII are involved.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 4c200358-a948-49ac-815f-396a0599b998

📥 Commits

Reviewing files that changed from the base of the PR and between 6b6d59a and 81193ef.

📒 Files selected for processing (5)

• framework/configstore/migrations.go
• framework/configstore/rdb.go
• framework/configstore/tables/ratelimit.go
• plugins/governance/store.go
• ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx

💤 Files with no reviewable changes (1)

• framework/configstore/tables/ratelimit.go

🚧 Files skipped from review as they are similar to previous changes (4)

• framework/configstore/rdb.go
• framework/configstore/migrations.go
• ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx
• plugins/governance/store.go

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes

  • Safely remove legacy calendar-aligned storage via migration when present.
  • Rate-limit resets now derive calendar alignment from the virtual-key configuration, not per-rate-limit records.

• New Features

  • Calendar alignment is now a virtual-key-wide setting affecting budgets and rate limits.

• UI Updates

  • Alignment toggle moved to the rate-limit section, shown only when relevant, with updated confirmation text.

Walkthrough

Calendar alignment was moved from per-rate-limit storage into a virtual-key–level property: legacy DB columns are dropped, the rate-limit model field was removed, virtual-key updates carry calendar_aligned, runtime logic derives alignment from owning VKs, and the UI toggle was moved to the VK rate-limiting section. (50 words)

Changes

Calendar Alignment Refactor

Layer / File(s)	Summary
Rate-limit model removal + migration framework/configstore/tables/ratelimit.go, framework/configstore/migrations.go	Removed CalendarAligned field from TableRateLimit. Added drop_legacy_calendar_aligned_columns migration that conditionally drops calendar_aligned columns from budgets and rate_limits; triggerMigrations invokes it and propagates errors.
Virtual key persistence whitelist framework/configstore/rdb.go	UpdateVirtualKey selects calendar_aligned (added) and removes budget_id from the update whitelist so VK rows can carry the alignment flag during updates.
Runtime resolution of alignment plugins/governance/store.go	ResetExpiredRateLimitsInMemory builds a reverse map from rate-limit ID → owning VK CalendarAligned (covering VK-owned and provider-config rate-limits) and uses that derived calendarAligned value when computing token/request reset targets instead of a per-rate-limit field.
UI: VK-level toggle and gating logic ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx	Watches for token/request reset-duration fields added; showCalendarAlignToggle now computed from budgets and enabled token/request rate-limits; moved calendar-alignment toggle and confirmation dialog from budget area into the rate-limiting section and updated messaging to reference both budgets and rate-limit counters.

sequenceDiagram
    participant UI as Client (UI)
    participant API as Server (API)
    participant DB as Database
    participant GS as GovernanceStore (in-memory)

    rect rgba(0,128,0,0.5)
    UI->>API: Submit VK update (includes calendar_aligned)
    API->>DB: UpdateVirtualKey (persist calendar_aligned)
    DB-->>API: OK
    API-->>UI: Response
    end

    rect rgba(0,0,255,0.5)
    Note over GS,DB: Runtime reset resolution (periodic / on-start)
    GS->>DB: Read cached VKs and rate-limits
    DB-->>GS: VKs, rate-limits
    GS->>GS: Build rateLimitID -> VK.calendar_aligned map
    GS->>GS: resolvePeriodStart(rateLimit, calendarAlignedFromMap)
    end

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested reviewers

• akshaydeo

Poem

🐰 I hop from rows of old to new,
One flag for keys, not each small view.
Columns fall and toggles move,
In-memory maps now find their groove,
A gentle thump — alignment true.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately reflects the main change: removing legacy calendar_aligned columns from the database and updating the UI to consolidate calendar alignment as a VK-only setting.
Description check	✅ Passed	The description comprehensively covers all required sections: summary, detailed changes across all affected files, type of change selection, affected areas, testing instructions, breaking changes acknowledgment, and security considerations.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 05-13-fix_fixes_calendar_aligned_migration_and_ui_for_vk

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 @coderabbitai help to get the list of available commands and usage tips.}

[CLAassistant]

All committers have signed the CLA.

[roroghost17]

• feat: implements calendar align feature at team level #3476
• fix: fixes calendar_aligned migration and UI for VK #3452 👈 (View in Graphite)
• remove budget level calendar alignemnt #3434
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

[greptile-apps]

Confidence Score: 4/5

Safe to merge once the non-deterministic calendar-alignment behavior for shared rate limits is resolved.

The migration and DB-layer changes are clean and correct. The main open question is in ResetExpiredRateLimitsInMemory: if two virtual keys reference the same rate-limit ID but differ in CalendarAligned, the last VK visited by sync.Map.Range wins — making alignment arbitrary on every reset cycle. This was flagged in a prior thread and remains unaddressed. No new blocking issues were found in this pass.

plugins/governance/store.go — the reverse-map building logic in ResetExpiredRateLimitsInMemory for shared rate limits.

Important Files Changed

Filename	Overview
framework/configstore/migrations.go	Adds migrationDropLegacyCalendarAlignedColumns: uses GORM HasColumn guards before each drop, wraps all errors properly, and is idempotent — safe on any DB state.
framework/configstore/rdb.go	Replaces dead budget_id entry (no matching struct field on TableVirtualKey) with calendar_aligned in the UpdateVirtualKey SELECT list, making VK-level alignment persisted on update for the first time.
framework/configstore/tables/ratelimit.go	Removes the CalendarAligned field; alignment is now sourced from the owning VK at reset time. Simple, clean removal with no other dependents in this struct.
plugins/governance/store.go	Builds reverse map (rate-limit ID → owning VK's CalendarAligned) at reset time; correctly handles both VK-direct and provider-config rate limits, but last-writer-wins assignment is non-deterministic when two VKs share a rate limit (flagged in previous review thread).
ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx	Extends the calendar-alignment toggle to rate limits by adding hasAnyAlignableRateLimit; dialog relocated to a shared section below both fields. Toggle visibility logic is sound except when watchedTokenMaxLimit is 0 (flagged in previous review thread).

_{Reviews (7): Last reviewed commit: "fix: fixes calendar_aligned migration an..." | Re-trigger Greptile}

[coderabbitai]

Actionable comments posted: 3

🤖 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 `@framework/configstore/migrations.go`:
- Around line 7546-7552: The migration currently swallows DDL errors by logging
warnings when executing tx.Exec for "ALTER TABLE governance_budgets DROP COLUMN
IF EXISTS calendar_aligned" and "ALTER TABLE governance_rate_limits DROP COLUMN
IF EXISTS calendar_aligned"; change this so that any non-nil error from tx.Exec
is returned (or aggregated and returned) so the migration fails instead of
succeeding silently—i.e., in the migration function that calls tx.Exec(...) for
those ALTER TABLE statements, remove the log.Printf warning-only branches and
propagate the error (return err or wrap both errors) so the migration
transaction aborts and the failure is recorded.

In `@plugins/governance/store.go`:
- Around line 1461-1473: The loop that builds rateLimitCalendarAligned from
gs.virtualKeys currently overwrites entries non-deterministically because
sync.Map.Range order is undefined; change the logic in the iteration (the block
that inspects TableVirtualKey.RateLimitID and each
ProviderConfigs[i].RateLimitID) to check for an existing entry before writing
and handle conflicts deterministically: if the ID is not present, set
rateLimitCalendarAligned[id] = vk.CalendarAligned; if it is present and the
existing value differs, record or surface the conflict (e.g., log an error
and/or return a construction error) instead of overwriting so the behavior is
deterministic and invariant violations are detected; apply the same check for
both vk.RateLimitID and each provider-config RateLimitID and add tests verifying
conflicting CalendarAligned values are handled.

In `@ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx`:
- Around line 303-307: The calendar_aligned toggle is shown for alignable rate
limits (hasAnyAlignableRateLimit / showCalendarAlignToggle) but onSubmit only
persists calendar_aligned in the budget branch, so toggles on rate-limit-only
VKs get lost; update the onSubmit handler to always include calendar_aligned in
the payload when the toggle is visible (i.e., when showCalendarAlignToggle or
hasAnyAlignableRateLimit is true), not just when budgets exist, and ensure you
read the value from the same form state used to render the toggle
(calendar_aligned form field) so changes to watchedTokenMaxLimit /
watchedRequestMaxLimit and their reset durations (watchedTokenResetDuration,
watchedRequestResetDuration) are respected.

🪄 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: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 0e5a49a8-5fd6-422a-a525-f3294cd10a45

📥 Commits

Reviewing files that changed from the base of the PR and between 63802ac and a47983d.

📒 Files selected for processing (5)

• framework/configstore/migrations.go
• framework/configstore/rdb.go
• framework/configstore/tables/ratelimit.go
• plugins/governance/store.go
• ui/app/workspace/virtual-keys/views/virtualKeySheet.tsx

💤 Files with no reviewable changes (1)

• framework/configstore/tables/ratelimit.go

[akshaydeo]

Merge activity

• May 15, 3:20 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• May 15, 3:21 PM UTC: @akshaydeo merged this pull request with Graphite.

The base branch was changed.