Background Quota Refresh & Automated Token Management

[dacsang97]

Overview

This pull request introduces a robust background worker for periodic quota refreshing and enhances the authentication handling within the quota management system. These changes eliminate the need for manual API calls to update quota data and resolve authentication failures caused by expired OAuth tokens. It also exposes several management API endpoints for monitoring and controlling quotas.

Key Changes

1. Quota Management API (internal/api/handlers/management/quota_fetchers.go)

Exposed new management endpoints to interact with the quota system:

• GET /v0/management/quotas: Retrieves quota data for all connected accounts across all providers.
• GET /v0/management/quotas/:provider: Filters and returns quota data for a specific provider.
• GET /v0/management/quotas/:provider/:account: Fetches detailed quota information for a specific account identifier (e.g., email).
• POST /v0/management/quotas/refresh: Invalidate cache and force an immediate refresh of quota data. Supports an optional JSON body {"providers": ["antigravity"]} to limit the scope.
• GET /v0/management/subscription-info: Provides subscription tier and project details for supported providers (Antigravity/Gemini-CLI).

2. Background Quota Refresh Worker

• Implemented a new Worker system in internal/quota/worker.go that periodically fetches quota data for all configured providers in the background.
• Integrated the worker with the server lifecycle, ensuring it starts automatically when the server is ready and shuts down gracefully.
• Configurable via the quota-refresh-interval (seconds) setting in config.yaml.

3. Automated OAuth Token Refresh

• Antigravity & Codex Fetchers: Updated antigravity.go and codex.go to proactively check for token expiration.
• Self-Healing: The system automatically refreshes expired or expiring access tokens using stored refresh_token credentials before attempting to fetch quota data, resolving previous 401 UNAUTHENTICATED errors.

4. Cache & Timestamp Consistency

• Synchronized TTL: The QuotaCache TTL is now automatically adjusted to match the background refresh interval (plus a 10% safety buffer).
• Refined last_updated: Refactored the response logic to ensure the last_updated timestamp reflects the time the data was actually fetched from the upstream provider (the cached time) rather than the timestamp of the API request.

Verification

• Verified 401 errors are resolved via automated token refresh.
• Confirmed API endpoints return correct cached data with stable last_updated timestamps.
• Verified background refresh worker triggers at configured intervals.
• All packages build and pass existing tests.

[gemini-code-assist]

Summary of Changes

Hello @dacsang97, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the quota management system by introducing automation for both data refreshing and authentication. It establishes a background worker to periodically update quota information across all providers, ensuring data is always current without manual intervention. Concurrently, it implements an intelligent OAuth token refresh mechanism for Antigravity and Codex providers, preventing service disruptions caused by expired access tokens. These changes aim to improve the reliability and efficiency of quota reporting, providing users with up-to-date and accurate usage statistics.

Highlights

• Background Quota Refresh Worker: A new worker system has been implemented to periodically fetch and cache quota data for all configured providers, eliminating the need for manual updates. This worker integrates with the server's lifecycle, starting and stopping gracefully.
• Automated OAuth Token Refresh: The Antigravity and Codex quota fetchers now proactively manage OAuth tokens. They check for token expiration (with a 5-minute buffer) before making API calls and automatically refresh tokens using stored "refresh_token"s, resolving authentication failures due to expired tokens.
• Dynamic Cache TTL and Accurate Timestamps: The QuotaCache now dynamically adjusts its Time-To-Live (TTL) to synchronize with the background worker's refresh interval, ensuring cached data remains valid. Additionally, the "last_updated" timestamp in API responses now accurately reflects when the data was retrieved from the provider, rather than the current request time.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a significant and well-structured feature for background quota refreshing and automated token management. The implementation of the background worker, caching layer, and provider-specific fetchers is robust. The automatic token refresh for expired OAuth credentials is a great enhancement that will improve reliability.

My review includes a few key points:

• A critical security concern regarding hardcoded OAuth client secrets.
• A high-severity issue with a hardcoded client ID.
• A few medium-severity issues related to API correctness and potential bugs in the caching and request handling logic.

Overall, this is a solid contribution. Addressing the feedback will further improve the security and correctness of the new quota management system.