feat(api): add rate limiting to endpoint passthrough API

crates/tabby/Cargo.toml

@@ -61,6 +61,7 @@ reqwest.workspace = true
 async-openai-alt.workspace = true
 spinners = "4.1.1"
 regex.workspace = true
+leaky-bucket = "1.1.2"
 
 [dependencies.openssl]
 optional = true

crates/tabby/src/routes/agent.rs → crates/tabby/src/routes/endpoint.rs

@@ -10,8 +10,11 @@ use axum::{
 use hyper::StatusCode;
 use serde::Serialize;
 use tabby_common::{axum::MaybeUserExt, config::EndpointConfig};
+use tracing::error;
 
-pub async fn agent_policy(
+use super::rate_limit::EndpointRateLimiters;
+
+pub async fn endpoint_policy(
     State(_config): State<Arc<EndpointConfig>>,
     Extension(MaybeUserExt(_user)): Extension<MaybeUserExt>,
     request: Request<axum::body::Body>,
@@ -39,7 +42,8 @@ pub async fn list_endpoints(State(config): State<Arc<EndpointConfig>>) -> Json<V
 }
 
 pub async fn endpoint(
-    State(config): State<Arc<EndpointConfig>>,
+    State((config, rate_limiters)): State<(Arc<EndpointConfig>, Arc<EndpointRateLimiters>)>,
+    Extension(MaybeUserExt(user)): Extension<MaybeUserExt>,
     Path((name, path)): Path<(String, String)>,
     method: Method,
     uri: Uri,
@@ -54,6 +58,25 @@ pub async fn endpoint(
         return StatusCode::NOT_FOUND.into_response();
     };
 
+    // User must be authenticated to access the endpoint
+    let user = match user {
+        Some(u) => u,
+        None => return StatusCode::UNAUTHORIZED.into_response(),
+    };
+
+    // Apply rate limiting if user_quota is configured
+    if let Some(user_quota) = &endpoint.user_quota {
+        let limiter = rate_limiters.get_or_create(
+            &endpoint.name,
+            &user.id.to_string(),
+            user_quota.requests_per_minute,
+        );
+        // Try to acquire a permit, return 429 if rate limited
+        if !limiter.try_acquire(1) {
+            return StatusCode::TOO_MANY_REQUESTS.into_response();
+        }
+    }
+
     let client = reqwest::Client::new();
     let path = if path.starts_with('/') {
         path
@@ -102,14 +125,46 @@ pub async fn endpoint(
 
     match req.send().await {
         Ok(resp) => {
+            // Forward the response back (including 4xx errors from upstream)
             let mut builder = Response::builder().status(resp.status());
             if let Some(headers) = builder.headers_mut() {
                 *headers = resp.headers().clone();
             }
             builder
                 .body(axum::body::Body::from_stream(resp.bytes_stream()))
-                .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
+                .unwrap_or_else(|e| {
+                    error!(
+                        "Failed to build response body for endpoint '{}': {}",
+                        name, e
+                    );
+                    StatusCode::INTERNAL_SERVER_ERROR.into_response()
+                })
+        }
+        Err(e) => {
+            // Check if the error has an associated status code (e.g., from error_for_status())
+            if let Some(status) = e.status() {
+                error!(
+                    "Endpoint '{}' ({}) returned error status {}: {}",
+                    name, target_url, status, e
+                );
+                return status.into_response();
+            }
+
+            // Handle timeout errors with 504 Gateway Timeout
+            if e.is_timeout() {
+                error!(
+                    "Request to endpoint '{}' ({}) timed out: {}",
+                    name, target_url, e
+                );
+                return StatusCode::GATEWAY_TIMEOUT.into_response();
+            }
+
+            // Other errors (connection errors, etc.) return 500
+            error!(
+                "Failed to forward request to endpoint '{}' ({}): {}",
+                name, target_url, e
+            );
+            StatusCode::INTERNAL_SERVER_ERROR.into_response()
         }
-        Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
     }
 }

crates/tabby/src/routes/mod.rs

@@ -52,18 +52,20 @@ pub async fn run_app(api: Router, ui: Option<Router>, host: IpAddr, port: u16) {
     .unwrap_or_else(|err| fatal!("Error happens during serving: {}", err))
 }
 
-mod agent;
 mod chat;
 mod completions;
+mod endpoint;
 mod events;
 mod health;
 mod models;
+mod rate_limit;
 mod server_setting;
 
-pub use agent::*;
 pub use chat::*;
 pub use completions::*;
+pub use endpoint::*;
 pub use events::*;
 pub use health::*;
 pub use models::*;
+pub use rate_limit::*;
 pub use server_setting::*;

crates/tabby/src/routes/rate_limit.rs

@@ -0,0 +1,58 @@
+use std::{
+    collections::HashMap,
+    sync::{Arc, RwLock},
+};
+
+use leaky_bucket::RateLimiter;
+use tokio::time::Duration;
+
+/// Creates a rate limiter for the given requests per minute.
+fn new_rate_limiter(rpm: u32) -> RateLimiter {
+    let rps = (rpm as f64 / 60.0).ceil() as usize;
+    RateLimiter::builder()
+        .initial(rps)
+        .interval(Duration::from_secs(1))
+        .refill(rps)
+        .build()
+}
+
+/// Key for rate limiting: (endpoint_name, user_id)
+type RateLimitKey = (String, String);
+
+/// Shared state for rate limiters, keyed by (endpoint_name, user_id)
+#[derive(Default)]
+pub struct EndpointRateLimiters {
+    limiters: RwLock<HashMap<RateLimitKey, Arc<RateLimiter>>>,
+}
+
+impl EndpointRateLimiters {
+    pub fn new() -> Self {
+        Self {
+            limiters: RwLock::new(HashMap::new()),
+        }
+    }
+
+    /// Get or create a rate limiter for the given endpoint and user.
+    pub fn get_or_create(&self, endpoint_name: &str, user_id: &str, rpm: u32) -> Arc<RateLimiter> {
+        let key = (endpoint_name.to_string(), user_id.to_string());
+
+        // Try read lock first
+        {
+            let limiters = self.limiters.read().unwrap();
+            if let Some(limiter) = limiters.get(&key) {
+                return limiter.clone();
+            }
+        }
+
+        // Need to create a new limiter, acquire write lock
+        let mut limiters = self.limiters.write().unwrap();
+        // Double-check in case another thread created it
+        if let Some(limiter) = limiters.get(&key) {
+            return limiter.clone();
+        }
+
+        let limiter = Arc::new(new_rate_limiter(rpm));
+        limiters.insert(key, limiter.clone());
+        limiter
+    }
+}

crates/tabby/src/serve.rs

@@ -287,21 +287,25 @@ async fn api_router(
     });
 
     if !config.endpoints.is_empty() {
-        let agent_state = Arc::new(tabby_common::config::EndpointConfig {
+        let endpoint_config = Arc::new(tabby_common::config::EndpointConfig {
             endpoints: config.endpoints.clone(),
         });
+        let endpoint_rate_limiters = Arc::new(routes::EndpointRateLimiters::new());
         routers.push(
             Router::new()
                 .route(
                     "/v2/endpoints/{:name}/{*path}",
-                    routing::any(routes::endpoint),
+                    routing::any(routes::endpoint)
+                        .with_state((endpoint_config.clone(), endpoint_rate_limiters)),
+                )
+                .route(
+                    "/v2/endpoints",
+                    routing::get(routes::list_endpoints).with_state(endpoint_config.clone()),
                 )
-                .route("/v2/endpoints", routing::get(routes::list_endpoints))
                 .layer(axum::middleware::from_fn_with_state(
-                    agent_state.clone(),
-                    routes::agent_policy,
-                ))
-                .with_state(agent_state),
+                    endpoint_config.clone(),
+                    routes::endpoint_policy,
+                )),
         );
     };