perf(usage): overlap the pre-inference credit checks

[PierreLeGuen]

Reduces sequential waiting in the pre-inference credit gate while preserving rejection precedence and staking synchronization order.

• The per-key spend gate runs first and rejects before any organization work.
• The staking preflight and organization balance read overlap. Organizations with a staking source, or a failed preflight, re-read their balance after preflight; active limits are read afterward.
• UsageCheckResult::evaluate shares the existing credit decision between middleware and service callers.

This saves one sequential database-call wait for non-staking organizations. It does not reduce their query count. Staking organizations and failed preflights perform one extra discarded balance read, and the overlap can use two pool connections. Pool waits and latency under concurrency should be checked during rollout.

Validation: current-head middleware tests passed (7), service tests matching usage passed (42), and GitHub CI including E2E is green. No SQL changes.

Related: batched signature writes in #1045 and prepared-statement reuse in #1046.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-10T23:56:05.800606Z	3e4df6f	PR opened

ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[claude]

Review: perf(usage): overlap the pre-inference credit checks

Good change — the ordering constraint that matters (limit read strictly after the staking preflight) is real and correctly preserved. I verified against staking_farm.rs: sync_for_source only writes organization_limits_history (via update_staking_farm_limit) and the source sync state — it never touches the balance/usage rows. So overlapping get_balance with the preflight is safe, and get_balance/get_limit are the same two reads check_can_use was already doing. Extracting UsageCheckResult::evaluate is a clean way to keep the two paths from drifting.

Two things I'd want fixed before merge, plus notes.

⚠️ 1. The ordering regression test cannot fail (crates/api/src/middleware/usage.rs:513)

assert_limits_read_after_staking is the guard for the one invariant this PR is careful about, but it can't detect a violation:

fn sync_organization_if_stale(&self, organization_id: Uuid) -> Pin<Box<...>> {
    self.calls.lock().unwrap().push(organization_id);
    self.events.lock().unwrap().push("staking");   // <-- at construction, not at await
    let should_fail = self.should_fail;
    Box::pin(async move { ... })                    // <-- resolves with no yield point
}

"staking" is recorded when the future is built, and the returned future is immediately ready. tokio::join! polls in argument order, so "staking" lands before anything else no matter how the code is sequenced. Someone moving get_limit(...) into the join! — exactly the regression the doc comment warns about — would still see staking < limits and a green test, while reintroducing the stale-limit read.

Make the preflight observably span the window:

fn sync_organization_if_stale(&self, organization_id: Uuid) -> Pin<Box<...>> {
    self.calls.lock().unwrap().push(organization_id);
    let events = self.events.clone();
    let should_fail = self.should_fail;
    Box::pin(async move {
        events.lock().unwrap().push("staking-start");
        tokio::task::yield_now().await;   // let any concurrent limit read run
        events.lock().unwrap().push("staking-done");
        if should_fail { anyhow::bail!("staking sync failed"); }
        Ok(None)
    })
}

then assert position("staking-done") < position("limits"). With the yield_now, the balance assertion also becomes meaningful (it proves overlap rather than just presence).

⚠️ 2. Per-key enforcement can be silently skipped (crates/api/src/middleware/usage.rs:126, :153)

run_usage_checks receives the limit and the spend read as two independent Options:

api_key_spend_limit: Option<i64>,
api_key_spend: F,   // F: Future<Output = Option<Result<i64, ()>>>
...
if let (Some(api_key_limit), Some(spend)) = (api_key_spend_limit, api_key_spend) {

If those ever disagree — limit Some, future yields None — the key-level spend gate is skipped entirely and an over-budget key is billed to the org with no error. check_usage_for_api_key keeps them consistent today, but nothing in the signature enforces it, and this is a fail-open path in a billing gate. Couple them in the type:

api_key_spend: Option<(i64 /* limit */, F)>,

or have the future yield Option<(i64, Result<i64, ()>)> so the limit travels with the read. Either way the mismatch becomes unrepresentable.

Notes (non-blocking)

• Extra work on the 402 path. A key over its own spend limit, or a failed get_api_key_spend, used to short-circuit before the staking preflight. Now the preflight (which can make an external NEAR RPC and write a limits-history row) and the balance read always run first. try_start_sync + the staleness window bound this, but a client retry-looping on 402 now generates external work per attempt. Intentional per the description — worth a line in the PR body.
• UsageServiceTrait::check_can_use now has zero production callers (repo-wide grep: trait def, UsageServiceImpl, and three test mocks). The description says the decision is "shared by the middleware and check_can_use" — true of evaluate, but check_can_use itself is dead in prod now. Fine to keep as API surface; just don't rely on it as the thing exercising evaluate.
• UsageCheckResult::evaluate has no direct unit tests. It is now the single source of truth for the credit gate and the four-arm match is pure and trivially testable — four asserts in ports.rs would be cheap insurance, especially for the (Some(balance), None) => NoLimitSet and (None, Some(0)) => NoCredits arms.
• Pool pressure (3 concurrent connections/request) is correctly called out and can't deadlock here — the join! members acquire and release independently with no nesting, so the worst case is queuing, not a cycle.
• No privacy concerns: the diff logs only IDs and amounts, and dropping api_key_id from the "within spend limit" debug line is a small improvement.

I reviewed this statically — cargo clippy/cargo test weren't available in this environment, so I did not independently reproduce the test results in the description.

⚠️

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3e4df6f4d3

ℹ️ 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".

[ironloopai]

Review · Status

🟩 Completed

IronLoop completed the review and posted it to GitHub.

Result

Open submitted review →

Run details

• Run: fd6acf6c-86c7-40c8-a659-8d7df7a5952c
• Base: main at 7c852f2
• Head: perf/chat-completions-db-calls at 3e4df6f
• Created: 2026-09-10 23:56 UTC
• Updated: 2026-09-10 23:59 UTC

_{Automatic trigger · attempt 1 of 3 · completed in 2m 32s}

[ironloopai]

Review · Summary

🟢 No actionable findings

No new actionable findings identified.

Validation

• ✅ Static behavior review — Traced the concurrent spend, staking-preflight, balance, and post-preflight limit paths through their service and repository implementations.
• ✅ Diff hygiene — The proposed changes contain no whitespace errors.

Review details

• Run: fd6acf6c-86c7-40c8-a659-8d7df7a5952c
• Attempts: 1

[PierreLeGuen]

Addressed the review in 62546da:

1. Ordering test could not fail — the staking mock now records staking-start, yields, then staking-done, and the assertion is staking-done < limits. Moving the limit read into the join! makes it fail. The balance-overlap test asserts the full event order (staking-start, balance, staking-done, limits).
2. Per-key gate could be skipped — run_usage_checks now takes Option<(limit, spend_read)>, so the limit travels with the read and the mismatch is unrepresentable.
3. 402 path — the per-key gate is awaited first and rejects before any organization work (Codex thread). Also re-reads the balance after the preflight for organizations that have a staking source (other Codex thread), so the pre-PR read order holds where a slow sync could matter.
4. Added four unit tests for UsageCheckResult::evaluate, one per arm.

Re-ran the credit e2e set locally (120 passed) plus the middleware and usage unit tests, clippy and fmt.

[github-actions]

🔍 OpenCodeReview found 4 issue(s) in this PR.

• ✅ 4 posted as inline comment(s)
• 📝 0 posted as summary

⚠️ 1 warning(s) occurred during review.

[lloydmak99]

Solid, well-scoped perf change: the pre-inference credit checks are re-ordered so the staking-farm preflight and the org balance read overlap via tokio::join!, with the credit decision extracted into UsageCheckResult::evaluate. The extraction is faithful and the ordering/error-precedence invariants are preserved.

• The per-key spend gate is still awaited first (ports.rs:139-175), so 402 api_key_limit_exceeded and 500 spend-read failures short-circuit before any RPC/DB work, and it's now type-coupled to its spend read (Option<(i64, F)>) so it can't be silently skipped.
• get_limit runs after the join!, preserving the limit-after-preflight invariant; staking orgs re-read balance after the preflight (ports.rs:210-215), matching prior behavior.
• Non-blocking: staking orgs do one extra, discarded balance read from the join! before re-reading (ports.rs:210-215) — wasted work, not a correctness bug, and called out in the PR body. Fine to leave as a follow-up.

Checks: git diff --check and cargo fmt --check passed; targeted unit tests couldn't build locally (no cc linker in sandbox), but GitHub CI lint/unit/integration/audit/deny passed on head 15b45427 (E2E pending at review time).