Feat/record usage api

[henrypark133]

• Refactored image usage recording to share helper logic and updated usage service traits/mocks to return UsageLogEntry and support API-driven recording.
• Added a new authenticated /v1/usage endpoint with request/response types, validation, pricing lookup, and usage recording that returns the created log entry.

---

Note

Medium Risk
Touches billing/usage accounting and adds a new write endpoint plus a DB uniqueness constraint; mistakes could cause under/over-charging or rejected inserts, though changes are scoped and covered by new E2E tests.

Overview
Adds a new authenticated POST /v1/usage endpoint (and OpenAPI schema) to let clients record chat-completion and image-generation usage directly, with server-side model pricing lookup, request validation, and a tagged-union response.

Makes usage recording idempotent per organization by hashing the caller-supplied id into inference_id, enforcing it with a new partial unique index, and updating the usage repository/service flow to return the existing record (and skip balance updates + cost metrics) on duplicates; also refactors image usage recording in completion routes to reuse shared helper/timeout+async-retry logic and updates mocks/tests accordingly.

^{Written by Cursor Bugbot for commit c971212. This will update automatically on new commits. Configure here.}

[gemini-code-assist]

Summary of Changes

Hello @henrypark133, 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 platform's usage tracking capabilities by introducing a dedicated API endpoint for recording usage events. It streamlines internal usage recording processes through refactoring and ensures that both internal and external usage logging is consistent, robust, and provides rich data through updated service interfaces and detailed response structures. The changes facilitate better integration with external systems for usage reporting and billing.

Highlights

• New Usage Recording API Endpoint: Introduced a new authenticated /v1/usage API endpoint that allows external systems to record usage events, including chat completions and image generations. This endpoint handles request validation, pricing lookup, and usage logging.
• Refactored Image Usage Recording: Existing image generation and editing usage recording logic has been refactored into shared helper functions (build_image_usage_request and record_usage_with_sync_fallback) to promote code reuse and improve robustness with synchronous attempts and asynchronous retries.
• Enhanced Usage Service Traits and Mocks: The UsageServiceTrait and its mock implementations (MockUsageService, CapturingUsageService) have been updated to return a UsageLogEntry object upon successful usage recording, providing more detailed information about the recorded event. A new record_usage_from_api method was added to the trait to specifically handle API-driven requests.
• Comprehensive API Request and Response Structures: Defined RecordUsageApiRequest as a tagged union to differentiate between chat completion and image generation usage requests, each with specific fields. Similarly, RecordUsageResponse provides detailed, type-specific responses including calculated costs and token/image counts.
• End-to-End Testing for New API: New end-to-end tests have been added to validate the functionality of the /v1/usage API, covering successful recording for different usage types, handling of external IDs, and various validation scenarios like model not found or zero tokens.

Changelog

• crates/api/src/lib.rs
  • Added build_usage_recording_routes function to encapsulate the new /v1/usage API endpoint setup.
  • Integrated the newly defined usage recording routes into the main application router.
• crates/api/src/openapi.rs
  • Included the record_usage route in the OpenAPI specification for API documentation.
  • Added RecordUsageResponse to the OpenAPI schema definitions to describe the API's output.
• crates/api/src/routes/completions.rs
  • Introduced build_image_usage_request to standardize the creation of usage requests for image operations.
  • Added record_usage_with_sync_fallback for resilient usage recording with synchronous attempts and asynchronous retries.
  • Refactored image_generations and image_edits handlers to use the new helper functions for usage recording.
  • Updated the expected return type of usage_service.record_usage from Result<(), UsageError> to Result<UsageLogEntry, UsageError>.
• crates/api/src/routes/usage.rs
  • Imported AuthenticatedApiKey for API key-based authentication on the new endpoint.
  • Defined RecordUsageResponse enum to structure the API response for different usage types (chat completion, image generation).
  • Implemented the record_usage async function for the new /v1/usage POST endpoint, including authentication, request parsing, and delegation to the usage service.
  • Added OpenAPI documentation for the new /v1/usage endpoint, detailing request body, responses, and security.
• crates/api/tests/e2e_usage_recording.rs
  • Added a new test file dedicated to end-to-end testing of the /v1/usage API.
  • Included tests for recording chat completion usage, image generation usage, usage with an external ID, and validation for scenarios like model not found and zero tokens.
• crates/services/src/test_utils.rs
  • Updated MockUsageService and CapturingUsageService to align with the new record_usage signature (returning UsageLogEntry).
  • Implemented the new record_usage_from_api trait method in mock services.
  • Added InferenceType and RecordUsageApiRequest to imports for updated service traits.
• crates/services/src/usage/mod.rs
  • Modified the record_usage function to return a UsageLogEntry instead of ().
  • Implemented record_usage_from_api to handle API requests, including model name resolution, input validation (e.g., non-negative tokens, positive image count), and delegation to the internal record_usage method.
  • Added logic to derive provider_request_id and inference_id from an optional external id provided in API requests.
• crates/services/src/usage/ports.rs
  • Updated the UsageServiceTrait to change the return type of record_usage to Result<UsageLogEntry, UsageError>.
  • Added a new trait method record_usage_from_api to abstract API-driven usage recording.
  • Defined RecordUsageApiRequest enum as a tagged union for API requests, supporting ChatCompletion and ImageGeneration with their specific fields.
  • Added ValidationError to the UsageError enum to handle specific validation failures.

Activity

• The pull request introduces a new API endpoint, /v1/usage, to allow external recording of usage data.
• Existing image usage recording logic has been refactored into reusable helper functions, improving code organization and maintainability.
• The core usage service interface (UsageServiceTrait) has been updated to return detailed UsageLogEntry objects, providing richer data post-recording.
• A new record_usage_from_api method was added to the usage service to specifically handle and validate requests originating from the new API endpoint.
• Comprehensive end-to-end tests have been added to ensure the new API endpoint functions correctly across various scenarios, including different usage types and validation checks.

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 new API endpoint for recording usage and refactors existing usage recording logic for image generation. The changes are well-structured, and the addition of end-to-end tests for the new endpoint is excellent. The core logic for the new /v1/usage endpoint, including request validation and delegation to the service layer, is solid. The refactoring in completions.rs significantly improves code clarity and reduces duplication.

I have a couple of suggestions for improvement:

• Refactoring duplicated code in the record_usage_with_sync_fallback helper function to improve maintainability.
• Making the handling of image_count in the new record_usage API handler more robust by logging an error if it's unexpectedly missing.

Overall, this is a great contribution that enhances the API's capabilities and improves the codebase.

[claude]

Code Review - PR #421

I've reviewed this PR focusing on critical issues. The overall architecture and implementation are solid, but I found one critical issue that needs to be addressed before merging:

---

⚠️ CRITICAL: Duplicate Usage Recording Risk

Location: crates/api/src/routes/usage.rs:588-640 (the new /v1/usage endpoint)

Issue: The endpoint has no duplicate detection for usage records. Unlike the inference pipeline which generates unique inference_id values for each request, this API allows clients to submit the same usage data multiple times, potentially leading to:

1. Double billing: A client could accidentally retry failed requests and record usage twice
2. Intentional abuse: Malicious actors could drain credits without actual usage
3. Data inconsistency: The same provider_request_id can create multiple usage log entries

Evidence from schema:

• The organization_usage_log table has no UNIQUE constraint on inference_id or provider_request_id (see V0037__usage_log_new_columns.sql and V0004__add_organization_usage_tracking.sql)
• Only non-unique indexes exist: idx_org_usage_inference_id, idx_org_usage_provider_request_id

Recommendation:
Add idempotency protection by checking for existing records when an id is provided:

// In record_usage_from_api, after deriving inference_id:
if let Some(inf_id) = inference_id {
    // Check if this inference_id already exists
    if self.usage_repository.get_usage_by_inference_id(inf_id).await?.is_some() {
        return Err(UsageError::ValidationError(
            "Usage already recorded for this id".into()
        ));
    }
}

Alternatively, add a UNIQUE constraint on inference_id in the database and handle the conflict gracefully by returning the existing record instead of an error (true idempotency).

---

Minor Observations (Non-blocking):

1. Error logging lacks context (crates/api/src/routes/usage.rs:631):

   • tracing::error!("Failed to record usage"); should log the actual error for debugging
   • Suggestion: tracing::error!(error = %e, "Failed to record usage");

2. Validation inconsistency (crates/services/src/usage/mod.rs:225):

   • Token validation checks < 0 but tokens are i32, so this is a good defensive check
   • However, consider using u32 in the API contract to enforce non-negativity at the type level

3. Missing transaction boundary:

   • The record_usage operation updates both organization_usage_log and organization_balance
   • Verify that the repository implementation uses a transaction to ensure atomicity

---

Security & Production Safety ✅

• Authentication: Properly uses API key auth with workspace context
• Rate limiting: Applied via api_key_rate_limit_middleware
• Usage checks: Pre-flight credit validation via usage_check_middleware
• Logging compliance: No customer data logged (follows CLAUDE.md privacy rules)
• Middleware ordering: Correct (auth → rate limit → usage check)

---

Architecture & Code Quality ✅

• Refactoring: The extraction of build_image_usage_request and record_usage_with_sync_fallback is excellent - reduces duplication and improves maintainability
• Tagged unions: RecordUsageApiRequest and RecordUsageResponse are well-designed for extensibility
• Test coverage: Comprehensive e2e tests covering happy paths and edge cases

---

Verdict

⚠️ Issues Found - Critical fix required

The duplicate usage recording issue is a production safety concern that could lead to billing discrepancies. Please add idempotency protection before merging.

Once the duplicate detection is added, this will be ready to merge. The refactoring work is excellent and the API design is solid.

---

cc @henrypark133

[henrypark133]

Issue: The endpoint has no duplicate detection for usage records. Unlike the inference pipeline which generates unique inference_id values for each request, this API allows clients to submit the same usage data multiple times, potentially leading to:

Made external id required so every new usage record has some idempotency protection

[nickpismenkov]

LGTM, thanks!

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.}