Skip to content

[codex] Support API key spend sorting#853

Open
PierreLeGuen wants to merge 3 commits into
mainfrom
codex/sort-api-keys-by-spend-api
Open

[codex] Support API key spend sorting#853
PierreLeGuen wants to merge 3 commits into
mainfrom
codex/sort-api-keys-by-spend-api

Conversation

@PierreLeGuen

@PierreLeGuen PierreLeGuen commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Adds order_by and order_direction query params to the workspace API-key list endpoint.
  • Threads API-key sort options through the workspace service and repository traits.
  • Orders by either the existing created_at DESC default or total API-key spend, now including both inference usage and service usage.
  • Adds an index for service-usage lookups by workspace/API key.
  • Adds e2e coverage for default created-at ordering plus descending/ascending usage ordering with inference and service spend.
  • Addresses review feedback by updating the services test stub to the new repository trait signature.

Related PRs

Validation

  • rustfmt --edition 2021 --check crates/api/tests/e2e_all/api_keys.rs crates/database/src/repositories/api_key.rs crates/services/src/auth/mod.rs
  • cargo test -p services --lib --no-run
  • cargo test -p api --test e2e_all api_keys::test_list_workspace_api_keys_orders_by_usage --no-run
  • cargo test --lib --bins --no-run
  • cargo clippy --all-targets --all-features -- -D warnings
  • git diff --check

Notes

  • I did not count a full runtime e2e pass locally. The shared local e2e database has a pre-existing V63 migration checksum mismatch, and a fresh DB gets through migrations but the server cannot boot here because dstack/TEE signing files are unavailable outside the CVM environment.

@PierreLeGuen PierreLeGuen temporarily deployed to Cloud API test env June 25, 2026 16:11 — with GitHub Actions Inactive
@PierreLeGuen PierreLeGuen temporarily deployed to Cloud API test env June 25, 2026 16:11 — with GitHub Actions Inactive
@claude

claude Bot commented Jun 25, 2026

Copy link
Copy Markdown

PR Review: Support API key spend sorting

Focused review on critical issues. The core implementation is solid — the dynamic ORDER BY is built from closed enum matches (no user-controlled strings reach format!, so no SQL injection), the new params are optional with the same defaults as before (backward compatible), ordering is deterministic via tie-breakers, and all three production callers were updated.

⚠️ Critical — test build will not compile

crates/services/src/auth/mod.rs:732 — the StubApiKeyRepo test mock still implements the old 3-parameter signature of list_by_workspace_paginated, but the trait in ports.rs now requires 5 parameters (order_by, order_direction). This stub lives inside the cfg(test) mod tests of the services crate, so it fails to compile under cargo test --lib / cargo test -p services.

The PR validation (cargo check -p api, cargo test -p api --test e2e_all ... --no-run) does not compile the services crate test code, which is why this was missed.

Fix — add the two missing params to the stub method (matching the trait):

async fn list_by_workspace_paginated(
    &self,
    _: WorkspaceId,
    _: i64,
    _: i64,
    _: Option<ApiKeyOrderBy>,
    _: Option<ApiKeyOrderDirection>,
) -> Result<Vec<ApiKey>, RepositoryError> {
    unimplemented!()
}

(add the ApiKeyOrderBy / ApiKeyOrderDirection imports to the test module as needed).

Minor (non-blocking)

  • The new ApiKeyOrderBy / ApiKeyOrderDirection enums derive ToSchema and are referenced in the utoipa::path params, but are not registered in components(schemas(...)). Verify they still render in the OpenAPI spec / Scalar UI.

Please confirm with a full cargo test --lib --bins (or at least cargo test -p services --lib --no-run) before merge.

⚠️ (issues found)

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces sorting capabilities when listing workspace API keys, allowing users to order keys by creation date or usage in ascending or descending order. These parameters are integrated into the API query parameters, service layers, and database repository, which now dynamically constructs the SQL query with the specified sorting and tie-breaker logic. An end-to-end test has also been added to verify sorting by usage. I have no further feedback to provide as there are no review comments.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 491441c070

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment on lines +199 to +200
order_by: Option<ApiKeyOrderBy>,
order_direction: Option<ApiKeyOrderDirection>,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Update the test ApiKeyRepository implementation

With the trait now requiring these two sort parameters, the #[cfg(test)] StubApiKeyRepo in crates/services/src/auth/mod.rs:731-736 still implements list_by_workspace_paginated with only workspace_id, limit, and offset (I checked repo-wide impls with rg "impl ApiKeyRepository for"). That makes the services crate's unit-test target fail to compile when cargo test -p services is run, so the auth tests cannot run until the stub signature is updated.

Useful? React with 👍 / 👎.


let order_by_column = match order_by {
ApiKeyOrderBy::CreatedAt => "ak.created_at",
ApiKeyOrderBy::Usage => "usage",

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Include service charges in usage ordering

When an API key has spend recorded through service usage (for example web_search, which writes api_key_id/total_cost rows to organization_service_usage_log and updates org balance), this new order_by=usage path still ranks by the usage alias that only sums organization_usage_log. In that scenario a service-heavy key is sorted as zero spend, so the spend ordering returned to the dashboard can be wrong even though the key has billable usage.

Useful? React with 👍 / 👎.

@Evrard-Nil Evrard-Nil left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review: Changes Requested

Overall the design is clean — the sort params are properly typed enums (no raw user input in the query), auth is unchanged (existing check_workspace_permission gate), and the order_by_column / order_dir values are hardcoded Rust match arms so there is no SQL injection risk. The tie-breakers are correct. One build-breaking bug found.


Bug (blocks merge): StubApiKeyRepo has stale method signature

crates/services/src/auth/mod.rs around line 733 — StubApiKeyRepo::list_by_workspace_paginated still has the old 4-param signature (self + WorkspaceId + 2 × i64). The trait now requires 6 params (+ order_by + order_direction).

cargo check passes because Rust only compiles test code during cargo test. cargo test -p services emits:

error[E0050]: method `list_by_workspace_paginated` has 4 parameters but the declaration in trait
  `workspace::ports::ApiKeyRepository::list_by_workspace_paginated` has 6

Fix: add _: Option<ApiKeyOrderBy> and _: Option<ApiKeyOrderDirection> to the stub signature.


Minor: no test for the created_at default sort path

The new test covers order_by=usage in both directions, and the unused_key (zero spend) correctly appears first on the ascending asc path. There's no test that omits order_by entirely and verifies the legacy created_at DESC default is preserved. Not a blocker, but worth adding to protect against future regressions.


Performance note (informational)

The ORDER BY usage path sorts a post-GROUP-BY computed column — an index on organization_usage_log(api_key_id) exists (idx_org_usage_api_key, V0004) so the join aggregation is index-assisted. Per-workspace API key counts are small in practice, so the sort will hit a hash-agg plan that fits in work_mem without a filesort. No new migration needed.


Cross-component consistency

Frontend PR nearai/nearai-cloud-ui#249 sends orderBy=usage and orderDirection (camelCase), which get forwarded as order_by=usage and order_direction (snake_case) to this API — the route's serde(rename_all = "snake_case") on ListApiKeysParams matches what the frontend will deliver. No mismatch.

@Evrard-Nil Evrard-Nil requested a review from lloydmak99 June 25, 2026 16:25
@PierreLeGuen PierreLeGuen deployed to Cloud API test env June 25, 2026 16:53 — with GitHub Actions Active

Copy link
Copy Markdown
Contributor Author

Review feedback has been addressed in b88b5488:

  • Updated the StubApiKeyRepo::list_by_workspace_paginated test implementation to include the new sort parameters, and cargo test -p services --lib --no-run passes.
  • Updated usage sorting to include both organization_usage_log and organization_service_usage_log, using pre-aggregated joins to avoid double-counting.
  • Added a service usage workspace/API-key index.
  • Extended e2e coverage to prove default created_at DESC ordering plus usage asc/desc ordering with inference and service spend.

CI is green: security audit, lint, unit tests, integration tests, e2e tests, and Test Suite all passed. Ready for re-review.

@PierreLeGuen PierreLeGuen requested a review from Evrard-Nil June 25, 2026 17:00

@Evrard-Nil Evrard-Nil left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The blocker from the previous review — StubApiKeyRepo::list_by_workspace_paginated missing the two new Option<ApiKeyOrderBy> / Option<ApiKeyOrderDirection> params — is fixed in b88b548. cargo test -p services now passes (430 + 6 tests, 0 failures). cargo check is clean.

The rest of the implementation is unchanged and correct:

  • order_by_column / order_dir / tie_breaker are all hardcoded Rust match-arm strings; no user input reaches the format! call — no SQL injection risk.
  • Service spend is now included in the usage sort key via a second LEFT JOIN on organization_service_usage_log, which was the gap that prompted commit b88b548.
  • The new migration V0067__add_service_usage_api_key_index.sql adds a covering index on (workspace_id, api_key_id) for the subquery join.
  • E2E test test_list_workspace_api_keys_orders_by_usage compiles and covers both ASC and DESC paths with mixed inference + service spend keys; default sort (no params) also asserts the created_at DESC behaviour.
  • serde(rename_all = "snake_case") on ListApiKeysParams matches the snake_case query params the frontend (nearai-cloud-ui#249) will send.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants