feat(proxy): add rerank, responses, passthrough, audio, images, messages + budget enforcement by moonming · Pull Request #19 · api7/aisix

moonming · 2026-04-19T12:59:25Z

Summary

New proxy endpoints: POST /v1/rerank (Cohere-style), POST /v1/responses (OpenAI Responses API), POST /v1/messages (Anthropic native pass-through), POST /v1/images/generations, POST /v1/audio/{transcriptions,translations,speech}, ANY /passthrough/:provider/*rest (raw provider pass-through)
Budget enforcement: pre-dispatch 429 when monthly spend ≥ cap; post-dispatch cost accumulation via BudgetTracker; Model.cost field for per-token cost calculation
Admin plane: GET /admin/v1/health, GET /admin/v1/spend, Budget/Credential/Team CRUD, playground handler, API-key rotation; HealthTracker and BudgetTracker shared between proxy and admin via Arc
Multipart forwarding: audio endpoints collect axum Multipart fields, rewrite the model field, and rebuild reqwest::multipart::Form before forwarding to upstream
Shared HTTP client: OnceLock<reqwest::Client> in http_client.rs used across messages, audio, rerank, responses, and passthrough handlers
Model schema: Model.cost: Option<ModelCost> added; ApiKey.max_budget_usd: Option<f64> added for inline budget caps

Test plan

All 387 workspace tests pass with cargo test --workspace
budget_exceeded_returns_429 — confirms pre-dispatch budget check returns 429
budget_accumulates_cost_on_success — confirms post-dispatch cost is tracked
passthrough, rerank, responses, messages — 401/403/404/happy-path tests for each new endpoint
Audio: 401, 404, happy-path returning audio/mpeg bytes
cargo check --workspace clean with no warnings
Route /passthrough/:provider/*rest wildcard works (catches all HTTP methods via axum::routing::any)

🤖 Generated with Claude Code

…ges, budget enforcement Add the remaining proxy API surface and enforce budget caps. New proxy endpoints: - POST /v1/rerank — Cohere-style rerank pass-through - POST /v1/responses — OpenAI Responses API pass-through (OpenAI-only) - POST /v1/messages — Anthropic native Messages API pass-through - POST /v1/images/generations — image generation via Bridge.generate_image() - POST /v1/audio/{transcriptions,translations,speech} — audio endpoints - ANY /passthrough/:provider/*rest — raw provider pass-through for any HTTP method Budget enforcement in chat handler: - Pre-dispatch: ProxyError::BudgetExceeded (429) when spend >= monthly cap - Post-dispatch: accumulate cost in BudgetTracker via cost_for(total_tokens) New model data: - Model.cost: {input_per_1k, output_per_1k} for budget tracking - ApiKey.max_budget_usd: inline per-key budget cap (alongside Budget entity) Shared HTTP client (OnceLock<reqwest::Client>) for messages/audio/rerank/responses/passthrough. 102 proxy tests pass; 0 failures across 387 workspace tests.

Copilot

Pull request overview

Adds a broad set of new proxy endpoints (OpenAI-compatible + provider-native passthrough), introduces in-process budget + health tracking shared with the admin plane, and expands the admin API surface for managing new entities and operational reporting.

Changes:

Add new proxy routes: /v1/models, /v1/responses, /v1/rerank, /v1/messages, /v1/{completions,embeddings,images,generations,audio/*}, and /passthrough/:provider/*rest
Add budget + health trackers (shared via Arc) and expose admin endpoints /admin/v1/{health,spend} plus CRUD for Budgets/Credentials/Teams and API key rotation
Extend gateway/provider layers to support embeddings, completions, and image generation via the Bridge trait; update schemas/snapshot tables for new entities

Reviewed changes

Copilot reviewed 49 out of 50 changed files in this pull request and generated 16 comments.

Show a summary per file

File	Description
crates/aisix-server/src/main.rs	Wires proxy/admin state together and shares trackers + proxy router with admin.
crates/aisix-ratelimit/src/limiter.rs	Adds `RateLimitStatus` and `Limiter::peek()` for rate-limit header injection.
crates/aisix-ratelimit/src/lib.rs	Re-exports `RateLimitStatus`.
crates/aisix-proxy/src/state.rs	Adds `BudgetTracker` + `HealthTracker` into `ProxyState`.
crates/aisix-proxy/src/responses.rs	Implements `POST /v1/responses` OpenAI pass-through.
crates/aisix-proxy/src/rerank.rs	Implements `POST /v1/rerank` pass-through.
crates/aisix-proxy/src/render.rs	Adds helper to inject `x-ratelimit-*` response headers.
crates/aisix-proxy/src/passthrough.rs	Adds raw provider pass-through route `/passthrough/:provider/*rest`.
crates/aisix-proxy/src/models.rs	Implements OpenAI-compatible `GET /v1/models`.
crates/aisix-proxy/src/messages.rs	Implements `POST /v1/messages` Anthropic-native pass-through.
crates/aisix-proxy/src/lib.rs	Mounts new proxy routes; adds tests for ratelimit headers + budget enforcement (chat).
crates/aisix-proxy/src/images.rs	Implements `POST /v1/images/generations` via bridge.
crates/aisix-proxy/src/http_client.rs	Introduces shared process-wide `reqwest::Client`.
crates/aisix-proxy/src/health.rs	Adds `HealthTracker` implementation + tests.
crates/aisix-proxy/src/error.rs	Adds `BudgetExceeded` proxy error mapping.
crates/aisix-proxy/src/embeddings.rs	Implements `POST /v1/embeddings` via bridge.
crates/aisix-proxy/src/completions.rs	Implements `POST /v1/completions` via bridge.
crates/aisix-proxy/src/chat.rs	Injects rate-limit headers; adds budget pre-check + spend accumulation; records health success/failure.
crates/aisix-proxy/src/budget.rs	Adds in-process monthly spend tracker with rollover behavior.
crates/aisix-proxy/src/audio.rs	Implements OpenAI audio endpoints (multipart forwarding + JSON speech).
crates/aisix-proxy/Cargo.toml	Adds dependencies needed for new proxy functionality (reqwest/chrono/anthropic devdep).
crates/aisix-provider-openai/src/wire.rs	Adds OpenAI embeddings wire types + conversions.
crates/aisix-provider-openai/src/bridge.rs	Implements `embed`, plus passthrough `complete` and `generate_image` for OpenAI bridge.
crates/aisix-gateway/src/lib.rs	Re-exports embeddings types from `chat` module.
crates/aisix-gateway/src/chat.rs	Adds normalized embeddings request/response types.
crates/aisix-gateway/src/bridge.rs	Extends `Bridge` trait with `embed`, `complete`, `generate_image` defaults.
crates/aisix-etcd/src/loader.rs	Loads new etcd kinds: credentials + budgets.
crates/aisix-core/src/models/team.rs	Adds `Team` entity definition + tests.
crates/aisix-core/src/models/snapshot.rs	Adds new tables (Credential/Budget/Team) to snapshot and updates tests.
crates/aisix-core/src/models/schema.rs	Adds JSON schemas + validators for Credential/Budget/Team.
crates/aisix-core/src/models/model.rs	Adds optional `Model.cost` (`ModelCost`) schema.
crates/aisix-core/src/models/mod.rs	Exposes new model modules/types/validators.
crates/aisix-core/src/models/credential.rs	Adds `Credential` entity definition + tests.
crates/aisix-core/src/models/budget.rs	Adds `Budget` entity definition + token cost helper + tests.
crates/aisix-core/src/models/apikey.rs	Adds `max_budget_usd`; adds wildcard model access + accessible model listing.
crates/aisix-core/src/lib.rs	Re-exports new entities/validators.
crates/aisix-admin/src/teams_handlers.rs	Adds Teams CRUD handlers.
crates/aisix-admin/src/store.rs	Extends config store trait + in-memory store for Budgets/Credentials/Teams.
crates/aisix-admin/src/state.rs	Adds optional shared budget/health trackers and optional proxy router wiring.
crates/aisix-admin/src/spend_handler.rs	Adds `GET /admin/v1/spend` reporting from `BudgetTracker`.
crates/aisix-admin/src/playground_handler.rs	Adds `/playground/chat/completions` in-process forwarding to proxy router.
crates/aisix-admin/src/lib.rs	Mounts new admin routes: rotate, budgets/credentials/teams CRUD, health/spend, playground.
crates/aisix-admin/src/health_handler.rs	Adds `GET /admin/v1/health` reporting using `HealthTracker`.
crates/aisix-admin/src/etcd_store.rs	Adds etcd CRUD plumbing for budgets/credentials/teams.
crates/aisix-admin/src/credentials_handlers.rs	Adds Credentials CRUD handlers.
crates/aisix-admin/src/budgets_handlers.rs	Adds Budgets CRUD handlers.
crates/aisix-admin/src/apikeys_handlers.rs	Adds API key rotation endpoint.
crates/aisix-admin/Cargo.toml	Adds chrono and dev dependencies used by new handler tests.
Cargo.toml	Enables axum multipart and reqwest multipart features.
Cargo.lock	Updates lockfile for new dependencies/features.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-19T13:05:40Z

+/// Per-token cost for budget tracking. Both values are in USD per 1,000 tokens.
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[serde(deny_unknown_fields)]
+pub struct ModelCost {
+    /// Input (prompt) token cost in USD per 1,000 tokens.
+    pub input_per_1k: f64,
+    /// Output (completion) token cost in USD per 1,000 tokens.
+    pub output_per_1k: f64,
+}
+
+impl ModelCost {
+    /// Calculate USD cost for the given token counts.
+    pub fn calculate(&self, input_tokens: u64, output_tokens: u64) -> f64 {
+        let input_cost = self.input_per_1k * (input_tokens as f64) / 1000.0;
+        let output_cost = self.output_per_1k * (output_tokens as f64) / 1000.0;
+        input_cost + output_cost
+    }
+}


The PR description mentions using Model.cost / per-token costs for budget calculation, but ModelCost is only defined here and not used anywhere in the proxy’s budget enforcement (current logic uses Budget.usd_per_1k_tokens). This makes the new schema field effectively dead. Either integrate Model.cost into cost accumulation or remove/defers the field so config doesn’t drift from behavior.

Copilot · 2026-04-19T13:05:40Z

+    // Budget pre-check. Refuse if the previous request already pushed
+    // monthly spend past the cap. Mid-request overshoot is bounded by
+    // one request worth of tokens — acceptable for V1; a future
+    // pre-debit-by-prompt-tokens-only mode can tighten it.
+    let budget_for_key = snapshot
+        .budgets
+        .entries()
+        .into_iter()
+        .find(|b| b.value.api_key_id == auth.entry.id);
+    if let Some(b) = budget_for_key.as_ref() {
+        if state
+            .budgets
+            .would_exceed(&auth.entry.id, b.value.monthly_usd_cap)
+        {
+            return Err(ProxyError::BudgetExceeded(auth.entry.id.clone()));
+        }
+    }


Budget enforcement is currently only applied in /v1/chat/completions (pre-check + post-spend accumulation). The new endpoints added in this PR (/v1/completions, /v1/embeddings, /v1/images/generations, /v1/messages, /v1/rerank, /v1/responses, audio, passthrough) don’t perform the same budget check, so callers can bypass the monthly cap by using those routes. Consider moving the budget check into shared middleware/auth extraction or duplicating it across all billable endpoints.

Copilot · 2026-04-19T13:05:41Z

+    match dispatch(&state, &auth, body, &request_id).await {
+        Ok((resp, provider)) => {
+            let elapsed = started.elapsed();
+            let status = 200u16;
+            emit_access_log(
+                &model_name,
+                &provider,
+                &api_key_id,
+                status,
+                elapsed,
+                &request_id,
+            );
+            state.metrics.record_request(&provider, &model_name, status, RequestOutcome::Success, elapsed);
+            resp


embeddings logs/metrics hard-code status=200 on the Ok path, but dispatch can return a non-200 response (e.g. 501 Not Implemented when the bridge returns the default “does not support embeddings” config error). This will skew access logs and metrics. Use resp.status() from the returned Response when recording.

Copilot · 2026-04-19T13:05:41Z

+    /// True if `(current spend + projected_cost) > cap`. The check
+    /// excludes the projected request itself — used for pre-commit
+    /// short-circuit when the *previous* month's tail already
+    /// over-shot the cap.
+    pub fn would_exceed(&self, api_key_id: &str, cap_usd: f64) -> bool {
+        self.spend(api_key_id) >= cap_usd
+    }


BudgetTracker::would_exceed docs mention (current spend + projected_cost) > cap, but the method signature doesn’t accept a projected cost and the implementation checks only spend(api_key_id) >= cap_usd. Either update the doc comment to match the actual semantics or extend the API to accept/compute a projected cost.

Copilot · 2026-04-19T13:05:41Z

+        // Seconds remaining in the current minute-window. Zero if the
+        // window just started or has already rolled.
+        let minute_reset = MINUTE_SECS - (now % MINUTE_SECS);
+


The minute_reset comment says the reset seconds are “Zero if the window just started”, but the calculation MINUTE_SECS - (now % MINUTE_SECS) yields MINUTE_SECS at the start of a window (e.g. when now % 60 == 0). Please fix the comment (or adjust the calculation if you truly want 0 at boundaries).

Copilot · 2026-04-19T13:05:43Z

+    // Inject upstream Authorization; strip the incoming proxy auth.
+    if api_key.is_empty() {
+        // Some providers use special headers (anthropic uses x-api-key).
+        if provider_lower == "anthropic" {
+            builder = builder.header("x-api-key", &api_key);
+        }
+    } else {
+        builder = builder.header(header::AUTHORIZATION, format!("Bearer {api_key}"));
+        if provider_lower == "anthropic" {
+            builder = builder.header("x-api-key", &api_key);
+            builder = builder.header("anthropic-version", "2023-06-01");
+        }
+    }


If provider_config.api_key is empty, passthrough still forwards the request (and for Anthropic even sets x-api-key to an empty value). Other proxy handlers treat an empty upstream key as a configuration error. Consider returning a BridgeError::Config/400 here when api_key is empty to avoid leaking misconfiguration as upstream 401s and to prevent sending malformed auth headers.

Copilot · 2026-04-19T13:05:43Z

+//! 3. Enforces that the model uses the `anthropic/` provider — non-Anthropic
+//!    models are rejected with 422 ("model is not an Anthropic provider").
+//! 4. Rewrites the `model` field to the upstream model name (strips the
+//!    `anthropic/` prefix).


The module-level docs say non-Anthropic models are rejected with 422, but the implementation returns ProxyError::InvalidRequest(...) which maps to 400. Please align the documentation with the actual behavior (or change the status mapping if 422 is intended).

Copilot · 2026-04-19T13:05:43Z

+    /// Maximum USD spend per calendar month. When the accumulated spend
+    /// for this key reaches or exceeds this cap the proxy returns 429.
+    /// Absent = no budget enforcement.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub max_budget_usd: Option<f64>,


The PR description mentions ApiKey.max_budget_usd for inline budget caps, but in this change set the field is not referenced anywhere (search shows only its definition/tests). As-is, setting max_budget_usd won’t actually enforce anything. Either wire this into the proxy budget check logic or remove it from the schema/description to avoid a misleading configuration knob.

Copilot · 2026-04-19T13:05:43Z

+    match dispatch(&state, &auth, body, &request_id).await {
+        Ok((resp, provider)) => {
+            let elapsed = started.elapsed();
+            emit_access_log(&model_name, &provider, &api_key_id, 200, elapsed, &request_id);
+            state.metrics.record_request(
+                &provider,
+                &model_name,
+                200,
+                RequestOutcome::Success,
+                elapsed,
+            );


completions logs/metrics hard-code status=200 in the success path, but dispatch can return a non-200 response (e.g. 501 Not Implemented when the provider doesn’t support /completions). This will emit incorrect access logs and Prometheus metrics for those cases. Consider reading resp.status() from the returned Response and using that for emit_access_log + record_request instead of a constant.

Copilot · 2026-04-19T13:05:43Z

+    // Find a model for this provider to grab api_key + api_base.
+    let provider_lower = provider.to_lowercase();
+    let all_models = snapshot.models.entries();
+    let model_entry = all_models
+        .into_iter()
+        .find(|e| {
+            e.value
+                .model
+                .to_lowercase()
+                .starts_with(&format!("{provider_lower}/"))
+        })
+        .ok_or_else(|| {
+            ProxyError::ModelNotFound(format!("no model found for provider `{provider}`"))
+        })?;


/passthrough/:provider/*rest selects the first Model whose model starts with {provider}/ by iterating snapshot.models.entries(). ResourceTable::entries() iterates a DashMap, so the iteration order is not deterministic; if multiple Models exist for the same provider, the upstream api_key/api_base chosen can vary between runs. Consider rejecting ambiguous configs (return 400 if >1 match), or require a specific model identifier in the route so selection is deterministic.

Fixes CI on feat/credentials-budgets: - cargo fmt --check failures across admin, core, gateway, proxy, ratelimit crates - clippy::needless_borrows_for_generic_args in audio.rs (reqwest post) - clippy::too_many_arguments on audio::emit_access_log (allow attribute) - clippy::map_flatten + clippy::redundant_closure in rerank.rs All 387 workspace tests still pass.

Copilot AI review requested due to automatic review settings April 19, 2026 12:59

Copilot started reviewing on behalf of moonming April 19, 2026 12:59 View session

Copilot AI reviewed Apr 19, 2026

View reviewed changes

moonming merged commit e25d5a2 into main Apr 19, 2026
6 checks passed

moonming deleted the feat/credentials-budgets branch April 19, 2026 23:52

juzhiyuan mentioned this pull request May 29, 2026

security: enforce model ACLs for routing targets and passthrough providers #449

Closed

13 tasks

jarvis9443 mentioned this pull request Jun 2, 2026

fix(passthrough): enforce caller model ACL on provider passthrough (#449) #489

Merged

Navigation Menu

Uh oh!

feat(proxy): add rerank, responses, passthrough, audio, images, messages + budget enforcement - #19

feat(proxy): add rerank, responses, passthrough, audio, images, messages + budget enforcement#19
moonming merged 2 commits into
mainfrom
feat/credentials-budgets

moonming commented Apr 19, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Copilot AI Apr 19, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

moonming commented Apr 19, 2026

Summary

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants