diff --git a/Cargo.lock b/Cargo.lock index 44c7fe44..68308a68 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -250,7 +250,9 @@ dependencies = [ "aisix-gateway", "aisix-guardrails", "aisix-obs", + "aisix-provider-openai", "aisix-ratelimit", + "async-stream", "async-trait", "axum", "bytes", @@ -267,6 +269,7 @@ dependencies = [ "tower-http", "tracing", "uuid", + "wiremock", ] [[package]] @@ -294,6 +297,10 @@ dependencies = [ "aisix-etcd", "aisix-gateway", "aisix-obs", + "aisix-provider-anthropic", + "aisix-provider-deepseek", + "aisix-provider-gemini", + "aisix-provider-openai", "aisix-proxy", "anyhow", "axum", diff --git a/crates/aisix-proxy/Cargo.toml b/crates/aisix-proxy/Cargo.toml index 861c8219..8105ee27 100644 --- a/crates/aisix-proxy/Cargo.toml +++ b/crates/aisix-proxy/Cargo.toml @@ -28,9 +28,12 @@ serde_json.workspace = true futures.workspace = true futures-util.workspace = true async-trait.workspace = true +async-stream = "0.3" thiserror.workspace = true tracing.workspace = true uuid.workspace = true [dev-dependencies] tokio = { workspace = true, features = ["macros", "rt"] } +aisix-provider-openai = { path = "../aisix-provider-openai" } +wiremock.workspace = true diff --git a/crates/aisix-proxy/src/auth.rs b/crates/aisix-proxy/src/auth.rs new file mode 100644 index 00000000..44bbe79e --- /dev/null +++ b/crates/aisix-proxy/src/auth.rs @@ -0,0 +1,139 @@ +//! Bearer-token authentication for the proxy surface. +//! +//! The extractor [`AuthenticatedKey`] parses `Authorization: Bearer ` +//! (or `x-api-key: ` as a convenience alternative), looks the key +//! up in the current `AisixSnapshot`, and yields the matching `ApiKey` +//! entity. Handlers take `AuthenticatedKey` as an argument — if parsing +//! or lookup fails the request is short-circuited with a 401 envelope +//! before the handler runs. + +use aisix_core::resource::ResourceEntry; +use aisix_core::ApiKey; +use axum::extract::{FromRef, FromRequestParts}; +use axum::http::request::Parts; +use std::sync::Arc; + +use crate::error::ProxyError; +use crate::state::ProxyState; + +#[derive(Debug, Clone)] +pub struct AuthenticatedKey { + pub entry: Arc>, +} + +impl AuthenticatedKey { + pub fn key(&self) -> &ApiKey { + &self.entry.value + } +} + +#[axum::async_trait] +impl FromRequestParts for AuthenticatedKey +where + S: Send + Sync, + ProxyState: FromRef, +{ + type Rejection = ProxyError; + + async fn from_request_parts(parts: &mut Parts, state: &S) -> Result { + let token = extract_bearer(parts)?; + let proxy_state = ProxyState::from_ref(state); + let snapshot = proxy_state.snapshot.load(); + let entry = snapshot + .apikeys + .get_by_name(&token) + .ok_or(ProxyError::InvalidApiKey)?; + Ok(AuthenticatedKey { entry }) + } +} + +fn extract_bearer(parts: &Parts) -> Result { + if let Some(auth) = parts.headers.get(axum::http::header::AUTHORIZATION) { + let s = auth.to_str().map_err(|_| ProxyError::MissingAuth)?; + if let Some(rest) = s.strip_prefix("Bearer ") { + let rest = rest.trim(); + if rest.is_empty() { + return Err(ProxyError::MissingAuth); + } + return Ok(rest.to_string()); + } + return Err(ProxyError::MissingAuth); + } + if let Some(raw) = parts.headers.get("x-api-key") { + let s = raw.to_str().map_err(|_| ProxyError::MissingAuth)?; + let s = s.trim(); + if s.is_empty() { + return Err(ProxyError::MissingAuth); + } + return Ok(s.to_string()); + } + Err(ProxyError::MissingAuth) +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::http::{HeaderMap, HeaderValue, Request}; + + fn parts_with(headers: HeaderMap) -> Parts { + let mut req = Request::builder().uri("/").body(()).unwrap(); + *req.headers_mut() = headers; + req.into_parts().0 + } + + #[test] + fn extract_bearer_happy_path() { + let mut h = HeaderMap::new(); + h.insert( + axum::http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer sk-abc"), + ); + let parts = parts_with(h); + assert_eq!(extract_bearer(&parts).unwrap(), "sk-abc"); + } + + #[test] + fn extract_bearer_accepts_x_api_key_as_alternative() { + let mut h = HeaderMap::new(); + h.insert("x-api-key", HeaderValue::from_static("sk-abc")); + let parts = parts_with(h); + assert_eq!(extract_bearer(&parts).unwrap(), "sk-abc"); + } + + #[test] + fn extract_bearer_rejects_missing_header() { + let parts = parts_with(HeaderMap::new()); + assert!(matches!( + extract_bearer(&parts), + Err(ProxyError::MissingAuth) + )); + } + + #[test] + fn extract_bearer_rejects_wrong_scheme() { + let mut h = HeaderMap::new(); + h.insert( + axum::http::header::AUTHORIZATION, + HeaderValue::from_static("Basic dXNlcjpwdw=="), + ); + let parts = parts_with(h); + assert!(matches!( + extract_bearer(&parts), + Err(ProxyError::MissingAuth) + )); + } + + #[test] + fn extract_bearer_rejects_empty_bearer() { + let mut h = HeaderMap::new(); + h.insert( + axum::http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer "), + ); + let parts = parts_with(h); + assert!(matches!( + extract_bearer(&parts), + Err(ProxyError::MissingAuth) + )); + } +} diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs new file mode 100644 index 00000000..4df92b82 --- /dev/null +++ b/crates/aisix-proxy/src/chat.rs @@ -0,0 +1,118 @@ +//! `POST /v1/chat/completions` handler. +//! +//! Flow: +//! 1. [`AuthenticatedKey`] extractor runs first — rejects unauthenticated +//! requests with a 401 envelope. +//! 2. Parse [`ChatFormat`] from the JSON body. +//! 3. Resolve `req.model` against the snapshot's Model table → 404 if +//! absent. +//! 4. Check the ApiKey's `allowed_models` whitelist → 403 if disallowed. +//! 5. Look up the matching `Bridge` on the Hub by `Model::provider()` → +//! 503 if no bridge registered. +//! 6. Build a [`BridgeContext`] and dispatch: +//! - `stream == true` → `chat_stream` + Sse response +//! - otherwise → `chat` + JSON response rendered as OpenAI +//! 7. Any `BridgeError` surfaces through [`ProxyError::Bridge`] which +//! supplies the right HTTP status and OpenAI-style error type. + +use aisix_gateway::{BridgeContext, ChatFormat}; +use axum::extract::State; +use axum::response::sse::{Event, KeepAlive, Sse}; +use axum::response::{IntoResponse, Response}; +use axum::Json; +use futures::{Stream, StreamExt}; +use std::convert::Infallible; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use uuid::Uuid; + +use crate::auth::AuthenticatedKey; +use crate::error::ProxyError; +use crate::render::{render_chunk, render_response}; +use crate::state::ProxyState; + +pub async fn chat_completions( + State(state): State, + auth: AuthenticatedKey, + Json(req): Json, +) -> Result { + if req.messages.is_empty() { + return Err(ProxyError::InvalidRequest( + "messages array must not be empty".into(), + )); + } + + let snapshot = state.snapshot.load(); + let model_entry = snapshot + .models + .get_by_name(&req.model) + .ok_or_else(|| ProxyError::ModelNotFound(req.model.clone()))?; + + if !auth.key().can_access(&req.model) { + return Err(ProxyError::ModelForbidden(req.model.clone())); + } + + let provider = model_entry + .value + .provider() + .ok_or_else(|| ProxyError::InvalidRequest("model has no provider prefix".into()))?; + let bridge = state + .hub + .get(provider) + .ok_or(ProxyError::ProviderUnavailable)?; + + let request_id = format!("req-{}", Uuid::new_v4()); + let model_arc = std::sync::Arc::new(model_entry.value.clone()); + let ctx = BridgeContext::new(&request_id, model_arc); + + let now = created_ts(); + + if req.is_streaming() { + let upstream = bridge.chat_stream(&req, &ctx).await?; + let model_name = req.model.clone(); + let sse_stream = build_sse_stream(upstream, model_name, now); + let response = + Sse::new(sse_stream).keep_alive(KeepAlive::new().interval(Duration::from_secs(15))); + return Ok(response.into_response()); + } + + let upstream = bridge.chat(&req, &ctx).await?; + let rendered = render_response(now, upstream); + Ok(Json(rendered).into_response()) +} + +fn created_ts() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs() as i64) + .unwrap_or(0) +} + +fn build_sse_stream( + upstream: aisix_gateway::ChatChunkStream, + _model: String, + created: i64, +) -> impl Stream> { + async_stream::stream! { + futures::pin_mut!(upstream); + while let Some(item) = upstream.next().await { + let ev = match item { + Ok(chunk) => { + let rendered = render_chunk(created, chunk); + match serde_json::to_string(&rendered) { + Ok(json) => Event::default().data(json), + Err(err) => Event::default() + .event("error") + .data(err.to_string()), + } + } + Err(err) => Event::default() + .event("error") + .data(err.to_string()), + }; + yield Ok::<_, Infallible>(ev); + } + // Emit the OpenAI-style [DONE] sentinel so clients that terminate + // on it behave correctly. + yield Ok::<_, Infallible>(Event::default().data("[DONE]")); + } +} diff --git a/crates/aisix-proxy/src/error.rs b/crates/aisix-proxy/src/error.rs new file mode 100644 index 00000000..36cd887c --- /dev/null +++ b/crates/aisix-proxy/src/error.rs @@ -0,0 +1,163 @@ +//! OpenAI-compatible error envelope used by every proxy endpoint. +//! +//! OpenAI's clients expect this exact shape (spec §3): +//! +//! ```json +//! { +//! "error": { +//! "message": "…", +//! "type": "invalid_request_error", +//! "param": null, +//! "code": null +//! } +//! } +//! ``` +//! +//! `ProxyError` is the internal error taxonomy; it implements +//! `IntoResponse` so handlers can `?`-propagate without touching +//! JSON shape boilerplate. + +use aisix_gateway::BridgeError; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::Json; +use serde::Serialize; + +#[derive(Debug, Serialize, Clone)] +pub struct ErrorEnvelope { + pub error: ErrorBody, +} + +#[derive(Debug, Serialize, Clone)] +pub struct ErrorBody { + pub message: String, + #[serde(rename = "type")] + pub kind: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + pub param: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub code: Option, +} + +impl ErrorEnvelope { + pub fn new(message: impl Into, kind: &'static str) -> Self { + Self { + error: ErrorBody { + message: message.into(), + kind, + param: None, + code: None, + }, + } + } + + pub fn with_code(mut self, code: impl Into) -> Self { + self.error.code = Some(code.into()); + self + } +} + +#[derive(Debug, thiserror::Error)] +pub enum ProxyError { + #[error("missing or malformed Authorization header")] + MissingAuth, + #[error("invalid API key")] + InvalidApiKey, + #[error("model {0:?} not found")] + ModelNotFound(String), + #[error("API key is not allowed to use model {0:?}")] + ModelForbidden(String), + #[error("request payload is invalid: {0}")] + InvalidRequest(String), + #[error("no bridge registered for provider")] + ProviderUnavailable, + #[error(transparent)] + Bridge(#[from] BridgeError), +} + +impl ProxyError { + pub fn status(&self) -> StatusCode { + match self { + ProxyError::MissingAuth | ProxyError::InvalidApiKey => StatusCode::UNAUTHORIZED, + ProxyError::ModelForbidden(_) => StatusCode::FORBIDDEN, + ProxyError::ModelNotFound(_) => StatusCode::NOT_FOUND, + ProxyError::InvalidRequest(_) => StatusCode::BAD_REQUEST, + ProxyError::ProviderUnavailable => StatusCode::SERVICE_UNAVAILABLE, + ProxyError::Bridge(b) => { + StatusCode::from_u16(b.http_status()).unwrap_or(StatusCode::BAD_GATEWAY) + } + } + } + + pub fn kind(&self) -> &'static str { + match self { + ProxyError::MissingAuth | ProxyError::InvalidApiKey => "invalid_api_key", + ProxyError::ModelForbidden(_) => "permission_denied", + ProxyError::ModelNotFound(_) => "model_not_found", + ProxyError::InvalidRequest(_) => "invalid_request_error", + ProxyError::ProviderUnavailable => "provider_unavailable", + ProxyError::Bridge(b) => b.error_type(), + } + } + + pub fn envelope(&self) -> ErrorEnvelope { + ErrorEnvelope::new(self.to_string(), self.kind()) + } +} + +impl IntoResponse for ProxyError { + fn into_response(self) -> Response { + let status = self.status(); + let body = self.envelope(); + (status, Json(body)).into_response() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn missing_auth_maps_to_401_invalid_api_key() { + let e = ProxyError::MissingAuth; + assert_eq!(e.status(), StatusCode::UNAUTHORIZED); + assert_eq!(e.kind(), "invalid_api_key"); + } + + #[test] + fn model_forbidden_is_403_permission_denied() { + let e = ProxyError::ModelForbidden("gpt-4o".into()); + assert_eq!(e.status(), StatusCode::FORBIDDEN); + assert_eq!(e.kind(), "permission_denied"); + } + + #[test] + fn bridge_error_inherits_status_and_type() { + let bridge_err = BridgeError::UpstreamStatus { + status: 429, + message: "rate limited".into(), + }; + let wrapped = ProxyError::Bridge(bridge_err); + assert_eq!(wrapped.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(wrapped.kind(), "upstream_error"); + } + + #[test] + fn bridge_5xx_collapses_via_bridge_error_mapping() { + let bridge_err = BridgeError::UpstreamStatus { + status: 503, + message: "busy".into(), + }; + let wrapped = ProxyError::Bridge(bridge_err); + assert_eq!(wrapped.status(), StatusCode::BAD_GATEWAY); + } + + #[test] + fn envelope_omits_null_param_and_code_on_wire() { + let env = ProxyError::ModelNotFound("x".into()).envelope(); + let json = serde_json::to_value(&env).unwrap(); + assert_eq!(json["error"]["type"], "model_not_found"); + assert!(json["error"].get("param").is_none()); + assert!(json["error"].get("code").is_none()); + } +} diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 7da74529..e81354f5 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -1,43 +1,49 @@ -//! aisix-proxy — client-facing proxy router (:3000). +//! aisix-proxy — client-facing proxy router (`:3000`). //! -//! This crate will host the full `/v1/*` OpenAI-compatible surface. For -//! PR #5 only the startup-sequence building block is here: [`build_router`] -//! returns a minimal router with `/health` so the server binary can bind, -//! serve, and be exercised by the startup-sequence e2e test. +//! Mounts the OpenAI-compatible surface: +//! - `GET /health` +//! - `POST /v1/chat/completions` (streaming + non-streaming) //! -//! The full routes, middleware stack, and streaming bridges land in their -//! own follow-up PRs. +//! Handlers run behind the [`AuthenticatedKey`] extractor which reads +//! the Bearer token (or `x-api-key` fallback) and looks the key up in +//! the current [`AisixSnapshot`]. Model authorisation is enforced per +//! request against `ApiKey::allowed_models`. Upstream calls are +//! dispatched through the [`aisix_gateway::Hub`] to the registered +//! `Bridge` for the Model's provider. +//! +//! Errors surface as OpenAI-style envelopes: +//! +//! ```json +//! {"error":{"message":"…","type":"…"}} +//! ``` +//! +//! Status codes follow [`crate::error::ProxyError::status`] — spec §3 auth +//! rules (401/403), `Bridge` mapping preserves upstream 4xx and collapses +//! upstream 5xx to 502. #![forbid(unsafe_code)] #![deny(rust_2018_idioms)] -use aisix_core::snapshot::SnapshotHandle; -use aisix_core::{AisixSnapshot, ProxyConfig}; -use axum::{http::StatusCode, routing::get, Json, Router}; -use serde_json::json; +mod auth; +mod chat; +mod error; +mod render; +mod state; -/// Runtime state the proxy router hands to handlers. Cheap to clone -/// (contains only an `Arc`-backed `SnapshotHandle`). -#[derive(Clone)] -pub struct ProxyState { - pub snapshot: SnapshotHandle, - pub request_body_limit_bytes: usize, -} +pub use auth::AuthenticatedKey; +pub use error::{ErrorEnvelope, ProxyError}; +pub use state::ProxyState; -impl ProxyState { - pub fn new(snapshot: SnapshotHandle, cfg: &ProxyConfig) -> Self { - Self { - snapshot, - request_body_limit_bytes: cfg.request_body_limit_bytes, - } - } -} +use axum::routing::{get, post}; +use axum::{http::StatusCode, Json, Router}; +use serde_json::json; -/// Build the proxy router. For PR #5 this is just `/health`; subsequent -/// PRs will mount `/v1/chat/completions`, `/v1/models`, etc. +/// Build the proxy router. Mounts `/health` plus the +/// OpenAI-compatible chat-completions surface. pub fn build_router(state: ProxyState) -> Router { Router::new() .route("/health", get(health)) + .route("/v1/chat/completions", post(chat::chat_completions)) .with_state(state) } @@ -51,6 +57,7 @@ async fn health( "status": "ok", "models": snap.models.len(), "apikeys": snap.apikeys.len(), + "providers": state.hub.len(), })), ) } @@ -58,40 +65,318 @@ async fn health( #[cfg(test)] mod tests { use super::*; + use aisix_core::models::Provider; + use aisix_core::resource::ResourceEntry; use aisix_core::snapshot::SnapshotHandle; - use axum::body::to_bytes; - use axum::http::Request; + use aisix_core::{AisixSnapshot, ApiKey, Model, ProxyConfig}; + use aisix_gateway::{Hub, SseDecoder, SseEvent}; + use aisix_provider_openai::OpenAiBridge; + use axum::body::{to_bytes, Body}; + use axum::http::{Request, StatusCode}; + use futures::StreamExt; + use std::sync::Arc; use tower::ServiceExt; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; fn cfg() -> ProxyConfig { ProxyConfig { addr: "127.0.0.1:0".into(), - request_body_limit_bytes: 1024, + request_body_limit_bytes: 1_048_576, tls: None, } } + fn build_state(snapshot: AisixSnapshot, hub: Arc) -> ProxyState { + let handle = SnapshotHandle::new(snapshot); + ProxyState::new(handle, hub, &cfg()) + } + + fn model_entry(name: &str, api_base: &str) -> ResourceEntry { + let cfg = format!( + r#"{{ + "name": "{name}", + "model": "openai/gpt-4o", + "provider_config": {{"api_key": "sk-upstream", "api_base": "{api_base}"}} + }}"# + ); + let model: Model = serde_json::from_str(&cfg).unwrap(); + ResourceEntry::new("model-id-1", model, 1) + } + + fn apikey_entry(key: &str, allowed: &[&str]) -> ResourceEntry { + let allowed_json = serde_json::to_string(&allowed).unwrap(); + let cfg = format!(r#"{{"key": "{key}", "allowed_models": {allowed_json}}}"#); + let apikey: ApiKey = serde_json::from_str(&cfg).unwrap(); + ResourceEntry::new("key-id-1", apikey, 1) + } + + fn seed_snapshot(model: &str, allowed: &[&str], api_base: &str) -> AisixSnapshot { + let snap = AisixSnapshot::new(); + snap.models.insert(model_entry(model, api_base)); + snap.apikeys.insert(apikey_entry("sk-caller", allowed)); + snap + } + + async fn run(app: Router, req: Request) -> axum::http::Response { + app.oneshot(req).await.unwrap() + } + #[tokio::test] - async fn health_returns_snapshot_counts() { - let handle = SnapshotHandle::new(AisixSnapshot::new()); - let state = ProxyState::new(handle, &cfg()); - let app = build_router(state); - - let resp = app - .oneshot( - Request::builder() - .uri("/health") - .body(axum::body::Body::empty()) - .unwrap(), - ) - .await + async fn non_streaming_happy_path_returns_openai_shaped_json() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-upstream", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3} + }))) + .mount(&upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(OpenAiBridge::new())); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], &upstream.uri()); + let app = build_router(build_state(snap, hub)); + + let body = serde_json::json!({ + "model": "my-gpt4", + "messages": [{"role": "user", "content": "hello"}] + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) .unwrap(); + + let resp = run(app, req).await; assert_eq!(resp.status(), StatusCode::OK); + let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); + let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(v["object"], "chat.completion"); + assert_eq!(v["choices"][0]["message"]["content"], "hi"); + assert_eq!(v["usage"]["total_tokens"], 3); + } + + #[tokio::test] + async fn missing_authorization_returns_401_envelope() { + let hub = Arc::new(Hub::new()); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], "http://unused"); + let app = build_router(build_state(snap, hub)); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from(r#"{"model":"my-gpt4","messages":[]}"#)) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + let bytes = to_bytes(resp.into_body(), 1024).await.unwrap(); + let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(v["error"]["type"], "invalid_api_key"); + } + + #[tokio::test] + async fn unknown_api_key_returns_401() { + let hub = Arc::new(Hub::new()); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], "http://unused"); + let app = build_router(build_state(snap, hub)); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-does-not-exist") + .header("content-type", "application/json") + .body(Body::from(r#"{"model":"my-gpt4","messages":[]}"#)) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + } + + #[tokio::test] + async fn model_not_in_allowed_list_returns_403() { + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(OpenAiBridge::new())); + // ApiKey allows only "other-model", the caller asks for "my-gpt4". + let snap = seed_snapshot("my-gpt4", &["other-model"], "http://unused"); + let app = build_router(build_state(snap, hub)); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"my-gpt4","messages":[{"role":"user","content":"hi"}]}"#, + )) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + let bytes = to_bytes(resp.into_body(), 1024).await.unwrap(); + let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(v["error"]["type"], "permission_denied"); + } + + #[tokio::test] + async fn unknown_model_returns_404_envelope() { + let hub = Arc::new(Hub::new()); + let snap = seed_snapshot("my-gpt4", &["*"], "http://unused"); + let app = build_router(build_state(snap, hub)); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"no-such-model","messages":[{"role":"user","content":"hi"}]}"#, + )) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + let bytes = to_bytes(resp.into_body(), 1024).await.unwrap(); + let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(v["error"]["type"], "model_not_found"); + } + + #[tokio::test] + async fn empty_messages_returns_400() { + let hub = Arc::new(Hub::new()); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], "http://unused"); + let app = build_router(build_state(snap, hub)); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(r#"{"model":"my-gpt4","messages":[]}"#)) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + } + + #[tokio::test] + async fn upstream_429_passes_through_with_openai_envelope() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(429).set_body_string("rate limited")) + .mount(&upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(OpenAiBridge::new())); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], &upstream.uri()); + let app = build_router(build_state(snap, hub)); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"my-gpt4","messages":[{"role":"user","content":"hi"}]}"#, + )) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS); let bytes = to_bytes(resp.into_body(), 1024).await.unwrap(); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); - assert_eq!(v["status"], "ok"); - assert_eq!(v["models"], 0); - assert_eq!(v["apikeys"], 0); + assert_eq!(v["error"]["type"], "upstream_error"); + } + + #[tokio::test] + async fn provider_without_registered_bridge_returns_503() { + // Snapshot has a Model targeting openai, but the Hub is empty. + let hub = Arc::new(Hub::new()); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], "http://unused"); + let app = build_router(build_state(snap, hub)); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from( + r#"{"model":"my-gpt4","messages":[{"role":"user","content":"hi"}]}"#, + )) + .unwrap(); + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE); + } + + #[tokio::test] + async fn streaming_response_emits_sse_then_done_sentinel() { + let upstream = MockServer::start().await; + let sse = "\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"},\"finish_reason\":null}]}\n\n\ +data: {\"id\":\"up-1\",\"model\":\"gpt-4o\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n\ +data: [DONE]\n\n"; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(sse), + ) + .mount(&upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(OpenAiBridge::new())); + let snap = seed_snapshot("my-gpt4", &["my-gpt4"], &upstream.uri()); + let app = build_router(build_state(snap, hub)); + + let body = serde_json::json!({ + "model": "my-gpt4", + "messages": [{"role": "user", "content": "hi"}], + "stream": true + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::OK); + assert!(resp + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or_default() + .contains("text/event-stream")); + + // Drain the body, decode SSE events, assert we got at least one + // delta chunk plus the terminating [DONE]. + let mut body_stream = resp.into_body().into_data_stream(); + let mut decoder = SseDecoder::new(); + let mut events = Vec::new(); + while let Some(chunk) = body_stream.next().await { + let bytes = chunk.unwrap(); + events.extend(decoder.feed(bytes.as_ref())); + } + assert!(events.contains(&SseEvent::Done), "missing [DONE] sentinel"); + let data_count = events + .iter() + .filter(|e| matches!(e, SseEvent::Data(_))) + .count(); + assert!( + data_count >= 2, + "expected at least two chat chunks, got {data_count}" + ); } } diff --git a/crates/aisix-proxy/src/render.rs b/crates/aisix-proxy/src/render.rs new file mode 100644 index 00000000..07191077 --- /dev/null +++ b/crates/aisix-proxy/src/render.rs @@ -0,0 +1,193 @@ +//! Render the gateway's normalised `ChatResponse` / `ChatChunk` into the +//! OpenAI wire shape that clients expect on `/v1/chat/completions`. +//! +//! The structure is intentionally independent from the provider crates' +//! upstream types — those describe what we *received*, while these +//! describe what we *emit*. Keeping them separate means a client-facing +//! schema change doesn't ripple into every provider adapter. + +use aisix_gateway::{ChatChunk, ChatResponse, FinishReason, Role}; +use serde::Serialize; + +#[derive(Debug, Serialize)] +pub struct ChatCompletion { + pub id: String, + pub object: &'static str, + pub created: i64, + pub model: String, + pub choices: Vec, + pub usage: Usage, +} + +#[derive(Debug, Serialize)] +pub struct NonStreamChoice { + pub index: u32, + pub message: RenderedMessage, + pub finish_reason: String, +} + +#[derive(Debug, Serialize)] +pub struct RenderedMessage { + pub role: &'static str, + pub content: String, +} + +#[derive(Debug, Serialize, Default)] +pub struct Usage { + pub prompt_tokens: u32, + pub completion_tokens: u32, + pub total_tokens: u32, +} + +#[derive(Debug, Serialize)] +pub struct ChatCompletionChunk { + pub id: String, + pub object: &'static str, + pub created: i64, + pub model: String, + pub choices: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub usage: Option, +} + +#[derive(Debug, Serialize)] +pub struct StreamChoice { + pub index: u32, + pub delta: RenderedDelta, + #[serde(skip_serializing_if = "Option::is_none")] + pub finish_reason: Option, +} + +#[derive(Debug, Serialize, Default)] +pub struct RenderedDelta { + #[serde(skip_serializing_if = "Option::is_none")] + pub role: Option<&'static str>, + #[serde(skip_serializing_if = "Option::is_none")] + pub content: Option, +} + +pub fn render_response(created_unix_ts: i64, resp: ChatResponse) -> ChatCompletion { + ChatCompletion { + id: resp.id, + object: "chat.completion", + created: created_unix_ts, + model: resp.model, + choices: vec![NonStreamChoice { + index: 0, + message: RenderedMessage { + role: role_to_str(resp.message.role), + content: resp.message.content, + }, + finish_reason: finish_to_str(&resp.finish_reason).to_string(), + }], + usage: Usage { + prompt_tokens: resp.usage.prompt_tokens, + completion_tokens: resp.usage.completion_tokens, + total_tokens: resp.usage.total_tokens, + }, + } +} + +pub fn render_chunk(created_unix_ts: i64, chunk: ChatChunk) -> ChatCompletionChunk { + ChatCompletionChunk { + id: chunk.id, + object: "chat.completion.chunk", + created: created_unix_ts, + model: chunk.model, + choices: vec![StreamChoice { + index: 0, + delta: RenderedDelta { + role: chunk.delta.role.map(role_to_str), + content: chunk.delta.content, + }, + finish_reason: chunk + .finish_reason + .as_ref() + .map(|f| finish_to_str(f).to_string()), + }], + usage: chunk.usage.map(|u| Usage { + prompt_tokens: u.prompt_tokens, + completion_tokens: u.completion_tokens, + total_tokens: u.total_tokens, + }), + } +} + +fn role_to_str(role: Role) -> &'static str { + match role { + Role::System => "system", + Role::User => "user", + Role::Assistant => "assistant", + Role::Tool => "tool", + } +} + +fn finish_to_str(f: &FinishReason) -> &str { + match f { + FinishReason::Stop => "stop", + FinishReason::Length => "length", + FinishReason::ContentFilter => "content_filter", + FinishReason::ToolCalls => "tool_calls", + FinishReason::Other(s) => s.as_str(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aisix_gateway::{ChatMessage, UsageStats}; + + #[test] + fn render_response_matches_openai_shape() { + let r = ChatResponse { + id: "cmpl-1".into(), + model: "m".into(), + message: ChatMessage::assistant("hello"), + finish_reason: FinishReason::Stop, + usage: UsageStats::new(3, 2), + }; + let out = render_response(42, r); + let json = serde_json::to_value(&out).unwrap(); + assert_eq!(json["object"], "chat.completion"); + assert_eq!(json["created"], 42); + assert_eq!(json["choices"][0]["finish_reason"], "stop"); + assert_eq!(json["choices"][0]["message"]["role"], "assistant"); + assert_eq!(json["choices"][0]["message"]["content"], "hello"); + assert_eq!(json["usage"]["total_tokens"], 5); + } + + #[test] + fn render_chunk_omits_finish_reason_when_absent() { + let chunk = ChatChunk { + id: "c".into(), + model: "m".into(), + delta: aisix_gateway::ChatDelta { + role: None, + content: Some("hi".into()), + }, + finish_reason: None, + usage: None, + }; + let out = render_chunk(1, chunk); + let json = serde_json::to_value(&out).unwrap(); + assert_eq!(json["object"], "chat.completion.chunk"); + assert_eq!(json["choices"][0]["delta"]["content"], "hi"); + // finish_reason / usage must be absent (not null). + assert!(json["choices"][0].get("finish_reason").is_none()); + assert!(json.get("usage").is_none()); + } + + #[test] + fn finish_reason_other_serialises_verbatim() { + let r = ChatResponse { + id: "cmpl".into(), + model: "m".into(), + message: ChatMessage::assistant(""), + finish_reason: FinishReason::Other("weird".into()), + usage: UsageStats::default(), + }; + let out = render_response(0, r); + let json = serde_json::to_value(&out).unwrap(); + assert_eq!(json["choices"][0]["finish_reason"], "weird"); + } +} diff --git a/crates/aisix-proxy/src/state.rs b/crates/aisix-proxy/src/state.rs new file mode 100644 index 00000000..e79e60de --- /dev/null +++ b/crates/aisix-proxy/src/state.rs @@ -0,0 +1,31 @@ +//! Axum state shared across every proxy handler. +//! +//! `ProxyState` holds: +//! - the lock-free `SnapshotHandle` for looking up +//! Models and ApiKeys on every request +//! - the `Hub` for resolving a `Provider` to the Bridge that serves it +//! - the configured request-body size limit +//! +//! Cheap to clone: every field is either an `Arc` or a small Copy scalar. + +use aisix_core::snapshot::SnapshotHandle; +use aisix_core::{AisixSnapshot, ProxyConfig}; +use aisix_gateway::Hub; +use std::sync::Arc; + +#[derive(Clone)] +pub struct ProxyState { + pub snapshot: SnapshotHandle, + pub hub: Arc, + pub request_body_limit_bytes: usize, +} + +impl ProxyState { + pub fn new(snapshot: SnapshotHandle, hub: Arc, cfg: &ProxyConfig) -> Self { + Self { + snapshot, + hub, + request_body_limit_bytes: cfg.request_body_limit_bytes, + } + } +} diff --git a/crates/aisix-server/Cargo.toml b/crates/aisix-server/Cargo.toml index 428de45b..27bef32a 100644 --- a/crates/aisix-server/Cargo.toml +++ b/crates/aisix-server/Cargo.toml @@ -19,6 +19,10 @@ aisix-gateway = { path = "../aisix-gateway" } aisix-obs = { path = "../aisix-obs" } aisix-proxy = { path = "../aisix-proxy" } aisix-admin = { path = "../aisix-admin" } +aisix-provider-openai = { path = "../aisix-provider-openai" } +aisix-provider-anthropic = { path = "../aisix-provider-anthropic" } +aisix-provider-gemini = { path = "../aisix-provider-gemini" } +aisix-provider-deepseek = { path = "../aisix-provider-deepseek" } tokio.workspace = true axum.workspace = true axum-server.workspace = true diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index d18e49e7..2bb49e3c 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -16,9 +16,15 @@ use std::path::PathBuf; use std::sync::Arc; use aisix_admin::AdminState; +use aisix_core::models::Provider; use aisix_core::Config; use aisix_etcd::{EtcdConfigProvider, Supervisor}; +use aisix_gateway::Hub; use aisix_obs::init_tracing; +use aisix_provider_anthropic::AnthropicBridge; +use aisix_provider_deepseek::deepseek_bridge; +use aisix_provider_gemini::gemini_bridge; +use aisix_provider_openai::OpenAiBridge; use aisix_proxy::ProxyState; use clap::Parser; use tokio::sync::watch; @@ -60,9 +66,13 @@ async fn run(cfg: Config) -> anyhow::Result<()> { let (cancel_tx, cancel_rx) = watch::channel(false); let watch_task = tokio::spawn(supervisor.clone().run(cancel_rx.clone())); - // Steps 7-8: routers. - let proxy_router = - aisix_proxy::build_router(ProxyState::new(snapshot_handle.clone(), &cfg.proxy)); + // Steps 7-8: build Hub, then routers. + let hub = Arc::new(build_hub()); + let proxy_router = aisix_proxy::build_router(ProxyState::new( + snapshot_handle.clone(), + hub.clone(), + &cfg.proxy, + )); let admin_router = aisix_admin::build_router(AdminState::new(snapshot_handle.clone(), &cfg.admin)); @@ -94,6 +104,18 @@ async fn run(cfg: Config) -> anyhow::Result<()> { Ok(()) } +/// Register all four provider bridges on a fresh Hub. The Hub is +/// created once at startup; future dynamic reload lands behind the +/// same `register()` call. +fn build_hub() -> Hub { + let hub = Hub::new(); + hub.register(Provider::Openai, Arc::new(OpenAiBridge::new())); + hub.register(Provider::Anthropic, Arc::new(AnthropicBridge::new())); + hub.register(Provider::Gemini, Arc::new(gemini_bridge())); + hub.register(Provider::Deepseek, Arc::new(deepseek_bridge())); + hub +} + /// Completes when the process receives SIGINT or SIGTERM (best-effort on /// Windows — Ctrl+C only) OR when another part of the system has already /// flipped the cancel channel.