feat: add budget reset endpoint and functionality

[Purvi09]

Summary

Adds a governance endpoint to manually reset a budget's usage on demand, without waiting for its reset window to elapse. Operators can now zero out a budget's current_usage immediately (e.g. after a billing correction or a mistaken spend) via POST /api/governance/budgets/{budget_id}/reset. The reset is applied to the in-memory governance store and persisted to the database, reusing the same side effects as an expiry-driven reset.

Ref #4892

Changes

• plugins/governance/store.go: added ResetBudget(ctx, budgetID) to the GovernanceStore interface and its LocalGovernanceStore implementation. It advances LastReset to now (restarting rolling windows; calendar-aligned budgets still reset at their next boundary), zeroes usage, zeroes the LastDBUsagesBudgets baseline so the next dump doesn't re-add stale usage, refreshes embedded VK budget references, fires the reset hook, and persists via ResetExpiredBudgets. Added the ErrBudgetNotFound sentinel error, returned for unknown IDs. A concurrent reset that loses the CAS returns the latest snapshot instead of erroring (idempotent).
• transports/bifrost-http/handlers/governance.go: added ResetBudget to the GovernanceManager interface, registered the POST /api/governance/budgets/{budget_id}/reset route, and added the resetBudget handler — 400 on invalid ID, 404 on ErrBudgetNotFound, 500 on other errors, 200 with { "budget": ... } on success.
• transports/bifrost-http/server/server.go: added ResetBudget to ServerCallbacks and implemented it on BifrostHTTPServer, delegating to the governance plugin's store.
• docs/openapi/openapi.json: documented the new endpoint (resetBudget operation, ManagementBearerAuth, 200/400/404/500 responses).
• Tests: added plugins/governance/budgetreset_test.go (manual reset side effects, not-found, calendar-aligned, idempotent) and handler tests in governance_test.go (success, not-found). Updated pricing_override_test.go to satisfy the extended interface.

Design notes / trade-offs

• Manual reset uses the same ResetBudgetAt + ResetExpiredBudgets path as expiry-driven resets, so persistence and hook behavior stay consistent.
• For calendar-aligned budgets the manual reset only zeroes current usage; the next scheduled reset still fires at the calendar boundary (covered by a dedicated test).
• The CAS-loses case is treated as success (returns the current snapshot) rather than an error, making the endpoint safe to retry / call concurrently.

Type of change

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

Affected areas

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

How to test

# Unit tests
go test ./plugins/governance/... -run ResetBudget
go test ./transports/bifrost-http/handlers/... -run ResetBudget

# Full suite
go test ./...

Manual check against a running server:

curl -X POST \
  -H "Authorization: Bearer $MANAGEMENT_TOKEN" \
  http://localhost:8080/api/governance/budgets/<budget_id>/reset

Expected: 200 with { "budget": { ... "current_usage": 0, "last_reset": "<now>" ... } }. An unknown budget_id returns 404; an empty/invalid ID returns 400.

Breaking changes

• Yes
• No

Related issues

Security considerations

The endpoint is protected by ManagementBearerAuth, consistent with other governance management routes. Resetting a budget clears spend tracking, so access should remain restricted to management tokens.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed (OpenAPI spec)
• I verified builds succeed (Go and UI)

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 8eeda138-d80a-4270-801e-a52b51f46f7c

📥 Commits

Reviewing files that changed from the base of the PR and between eca53c3 and e41a414.

📒 Files selected for processing (2)

• plugins/governance/budgetreset_test.go
• plugins/governance/store.go

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

• plugins/governance/budgetreset_test.go

---

📝 Walkthrough

Summary by CodeRabbit

• New Features
  • Added a new endpoint to reset a governance budget and return the updated budget details.
  • Resetting now reports clear outcomes for success, missing budgets, and invalid requests.
• Bug Fixes
  • Improved reset handling so budget usage is cleared consistently and linked records stay in sync.
  • Calendar-based budgets now keep their expected reset timing after a manual reset.
• Tests
  • Added coverage for successful resets, missing budgets, error handling, and date-aligned behavior.

Walkthrough

Adds a manual budget reset API and store method, wires it through the HTTP server and handler, documents the new endpoint in OpenAPI, and adds coverage for success and error cases.

Changes

Manual Budget Reset

Layer / File(s)	Summary
Governance store reset implementation and tests plugins/governance/store.go, plugins/governance/budgetreset_test.go	Adds ErrBudgetNotFound, updates the GovernanceStore.ResetBudget contract, implements LocalGovernanceStore.ResetBudget, and adds tests for manual reset, not-found, calendar alignment, and idempotency.
HTTP handler, routing, and server callback wiring transports/bifrost-http/handlers/governance.go, transports/bifrost-http/handlers/governance_test.go, transports/bifrost-http/handlers/pricing_override_test.go, transports/bifrost-http/server/server.go	Adds ResetBudget to GovernanceManager and ServerCallbacks, registers the reset route, implements the handler with 400/404/500 mapping, wires BifrostHTTPServer.ResetBudget to the governance plugin, and updates handler tests plus mocks.
OpenAPI documentation docs/openapi/openapi.json	Adds the OpenAPI path for POST /api/governance/budgets/{budget_id}/reset with parameters, security, and response schemas.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• maximhq/bifrost#3434: Both affect governance budget reset flow and calendar-aligned reset behavior in plugins/governance/store.go.
• maximhq/bifrost#3476: Both touch budget reset timing and calendar alignment semantics used by the reset path.
• maximhq/bifrost#4820: Both modify budget reset side effects around reset hooks in the governance store.

Suggested reviewers: akshaydeo, roroghost17

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately describes the main change: adding a budget reset endpoint and supporting functionality.
Description check	✅ Passed	The PR description closely follows the template and covers summary, changes, testing, security, breaking changes, and checklist items.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

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"

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🤖 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 `@plugins/governance/budgetreset_test.go`:
- Around line 109-127: The idempotent reset test in
TestGovernanceStore_ResetBudget_Idempotent only checks the returned budget ID
and misses the main reset contract in the CAS-decline path. Update this test to
also assert the returned snapshot from ResetBudget has CurrentUsage reset to
zero (or the expected post-reset value) when ResetBudgetAt declines due to a
future LastReset, using buildBudgetWithUsage and ResetBudget as the key symbols
to locate the case.

In `@plugins/governance/store.go`:
- Around line 1867-1897: ResetBudget currently treats a failed ResetBudgetAt CAS
as success by returning the latest snapshot even when CurrentUsage may still be
non-zero. Update LocalGovernanceStore.ResetBudget so the fallback branch does
not silently succeed on a stale LastReset; either retry with a reset target that
is guaranteed to advance beyond the current timestamp or return an error when
the reset cannot be applied. Also update the idempotent test around ResetBudget
to assert that the returned budget from a manual reset has CurrentUsage equal to
zero.

🪄 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: 25716481-f3bd-472f-91a4-ccf88d342952

📥 Commits

Reviewing files that changed from the base of the PR and between f244691 and eca53c3.

📒 Files selected for processing (7)

• docs/openapi/openapi.json
• plugins/governance/budgetreset_test.go
• plugins/governance/store.go
• transports/bifrost-http/handlers/governance.go
• transports/bifrost-http/handlers/governance_test.go
• transports/bifrost-http/handlers/pricing_override_test.go
• transports/bifrost-http/server/server.go

[greptile-apps]

Confidence Score: 5/5

Safe to merge; the new endpoint correctly reuses the existing reset machinery and is guarded by management auth.

The CAS loop, baseline zeroing, VK reference refresh, hook firing, and DB persistence are all wired correctly and follow established patterns in the store. The two findings are about test coverage gaps and a subtle UX question (does reset on a zero-usage budget advance LastReset?), neither of which affects runtime correctness.

plugins/governance/store.go (early-return behaviour on zero usage) and plugins/governance/budgetreset_test.go (calendar-aligned and already-zero test gaps)

Important Files Changed

Filename	Overview
plugins/governance/store.go	Adds ResetBudget with a CAS loop that correctly handles future LastReset via nanosecond advance. Early return on CurrentUsage == 0 silently skips LastReset advancement, which may surprise callers expecting a window restart.
plugins/governance/budgetreset_test.go	New test file covering manual reset, not-found, idempotent, and calendar-aligned cases. The calendar-aligned test doesn't set budget.IsCalendarAligned = true, so it doesn't unambiguously exercise the calendar-aligned expiry path. The already-zero early-return path lacks a dedicated test.
transports/bifrost-http/handlers/governance.go	Adds resetBudget handler and route registration. Error mapping (400/404/500) is correct and consistent with other governance handlers.
transports/bifrost-http/handlers/governance_test.go	Adds four handler-level tests: success, not-found, store error (500), and malformed encoding (400). Coverage is thorough for the handler layer.
transports/bifrost-http/server/server.go	Adds ResetBudget to ServerCallbacks interface and BifrostHTTPServer, delegating to the governance plugin's store. Consistent with other governance delegation patterns.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Client
    participant GovernanceHandler
    participant BifrostHTTPServer
    participant LocalGovernanceStore
    participant BudgetsMap as sync.Map (budgets)
    participant DB

    Client->>GovernanceHandler: "POST /api/governance/budgets/{id}/reset"
    GovernanceHandler->>BifrostHTTPServer: ResetBudget(ctx, budgetID)
    BifrostHTTPServer->>LocalGovernanceStore: ResetBudget(ctx, budgetID)
    LocalGovernanceStore->>BudgetsMap: LoadBudget (not-found check)
    alt Budget not found
        LocalGovernanceStore-->>GovernanceHandler: ErrBudgetNotFound
        GovernanceHandler-->>Client: 404
    else Budget found
        loop "CAS retry until ok or CurrentUsage==0"
            LocalGovernanceStore->>BudgetsMap: LoadBudget (snapshot)
            alt "CurrentUsage == 0"
                LocalGovernanceStore-->>GovernanceHandler: current snapshot (early return)
            else "CurrentUsage > 0"
                LocalGovernanceStore->>BudgetsMap: ResetBudgetAt (CAS zero usage, advance LastReset)
            end
        end
        LocalGovernanceStore->>LocalGovernanceStore: Zero LastDBUsagesBudgets baseline
        LocalGovernanceStore->>LocalGovernanceStore: updateBudgetReferences (VK refresh)
        LocalGovernanceStore->>LocalGovernanceStore: Fire onBudgetsReset hook
        LocalGovernanceStore->>DB: ResetExpiredBudgets (persist)
        LocalGovernanceStore-->>BifrostHTTPServer: resetBudget (zeroed snapshot)
        BifrostHTTPServer-->>GovernanceHandler: resetBudget
        GovernanceHandler-->>Client: 200 budget current_usage 0
    end

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Client
    participant GovernanceHandler
    participant BifrostHTTPServer
    participant LocalGovernanceStore
    participant BudgetsMap as sync.Map (budgets)
    participant DB

    Client->>GovernanceHandler: "POST /api/governance/budgets/{id}/reset"
    GovernanceHandler->>BifrostHTTPServer: ResetBudget(ctx, budgetID)
    BifrostHTTPServer->>LocalGovernanceStore: ResetBudget(ctx, budgetID)
    LocalGovernanceStore->>BudgetsMap: LoadBudget (not-found check)
    alt Budget not found
        LocalGovernanceStore-->>GovernanceHandler: ErrBudgetNotFound
        GovernanceHandler-->>Client: 404
    else Budget found
        loop "CAS retry until ok or CurrentUsage==0"
            LocalGovernanceStore->>BudgetsMap: LoadBudget (snapshot)
            alt "CurrentUsage == 0"
                LocalGovernanceStore-->>GovernanceHandler: current snapshot (early return)
            else "CurrentUsage > 0"
                LocalGovernanceStore->>BudgetsMap: ResetBudgetAt (CAS zero usage, advance LastReset)
            end
        end
        LocalGovernanceStore->>LocalGovernanceStore: Zero LastDBUsagesBudgets baseline
        LocalGovernanceStore->>LocalGovernanceStore: updateBudgetReferences (VK refresh)
        LocalGovernanceStore->>LocalGovernanceStore: Fire onBudgetsReset hook
        LocalGovernanceStore->>DB: ResetExpiredBudgets (persist)
        LocalGovernanceStore-->>BifrostHTTPServer: resetBudget (zeroed snapshot)
        BifrostHTTPServer-->>GovernanceHandler: resetBudget
        GovernanceHandler-->>Client: 200 budget current_usage 0
    end

Loading

_{Reviews (2): Last reviewed commit: "fix(governance): force manual budget res..." | Re-trigger Greptile}