feat(proxy): POST /v1/chat/completions + bearer auth + Hub dispatch

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

crates/aisix-proxy/src/auth.rs

@@ -0,0 +1,139 @@
+//! Bearer-token authentication for the proxy surface.
+//!
+//! The extractor [`AuthenticatedKey`] parses `Authorization: Bearer <key>`
+//! (or `x-api-key: <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<ResourceEntry<ApiKey>>,
+}
+
+impl AuthenticatedKey {
+    pub fn key(&self) -> &ApiKey {
+        &self.entry.value
+    }
+}
+
+#[axum::async_trait]
+impl<S> FromRequestParts<S> for AuthenticatedKey
+where
+    S: Send + Sync,
+    ProxyState: FromRef<S>,
+{
+    type Rejection = ProxyError;
+
+    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
+        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<String, ProxyError> {
+    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)
+        ));
+    }
+}

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<ProxyState>,
+    auth: AuthenticatedKey,
+    Json(req): Json<ChatFormat>,
+) -> Result<Response, ProxyError> {
+    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<Item = Result<Event, Infallible>> {
+    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]"));
+    }
+}