Share concurrent request limits across the fleet

crates/api/src/lib.rs

@@ -424,14 +424,30 @@ pub async fn init_domain_services_with_pool(
         as Arc<dyn services::completions::ports::OrganizationConcurrentLimitRepository>;
 
     // Create completion service with usage tracking (needs usage_service)
-    let completion_service = Arc::new(services::CompletionServiceImpl::new(
+    let mut completion_service_impl = services::CompletionServiceImpl::new(
         inference_provider_pool.clone(),
         attestation_service.clone(),
         usage_service.clone(),
         metrics_service.clone(),
         models_repo.clone() as Arc<dyn services::models::ModelsRepository>,
         org_limit_repository,
-    ));
+    );
+
+    if config.fleet_concurrency.mode != config::FleetConcurrencyMode::Off {
+        let lease_repository = Arc::new(
+            database::repositories::concurrency_lease::PostgresConcurrencyLeaseRepository::new(
+                database.pool().clone(),
+            ),
+        );
+        completion_service_impl = completion_service_impl.with_fleet_concurrency(
+            lease_repository,
+            config.fleet_concurrency.instance_id.clone(),
+            std::time::Duration::from_secs(config.fleet_concurrency.lease_ttl_seconds),
+            config.fleet_concurrency.mode == config::FleetConcurrencyMode::Enforce,
+        );
+    }
+
+    let completion_service = Arc::new(completion_service_impl);
 
     let brave_search_provider =
         Arc::new(services::responses::tools::brave::BraveWebSearchProvider::new());
@@ -2759,6 +2775,7 @@ mod tests {
             aml: config::AmlConfig::default(),
             usage_reporting: config::UsageReportingConfig::default(),
             ita: config::ItaAttestationConfig::default(),
+            fleet_concurrency: config::FleetConcurrencyConfig::default(),
         };
 
         // Initialize services
@@ -2869,6 +2886,7 @@ mod tests {
             aml: config::AmlConfig::default(),
             usage_reporting: config::UsageReportingConfig::default(),
             ita: config::ItaAttestationConfig::default(),
+            fleet_concurrency: config::FleetConcurrencyConfig::default(),
         };
 
         let auth_components = init_auth_services(database.clone(), &config);

crates/api/tests/common/mod.rs

@@ -131,6 +131,7 @@ pub fn test_config() -> ApiConfig {
             ..config::UsageReportingConfig::default()
         },
         ita: config::ItaAttestationConfig::default(),
+        fleet_concurrency: config::FleetConcurrencyConfig::default(),
     }
 }
 
@@ -434,6 +435,36 @@ where
     (server, database)
 }
 
+/// Independent servers sharing one database, standing in for fleet replicas.
+/// Each gets its own services, so agreement has to travel through the database.
+pub async fn setup_test_fleet<F>(
+    instances: usize,
+    mutate: F,
+) -> (
+    Vec<axum_test::TestServer>,
+    Vec<Arc<inference_providers::mock::MockProvider>>,
+    Arc<Database>,
+)
+where
+    F: Fn(&mut config::ApiConfig),
+{
+    let infra = setup_test_infrastructure().await;
+    let database = infra.database.clone();
+
+    let mut servers = Vec::with_capacity(instances);
+    let mut mocks = Vec::with_capacity(instances);
+    for _ in 0..instances {
+        let mut config = infra.config.clone();
+        mutate(&mut config);
+        let (server, _pool, mock, _router) =
+            build_test_server_components(database.clone(), config).await;
+        servers.push(server);
+        mocks.push(mock);
+    }
+
+    (servers, mocks, database)
+}
+
 pub async fn setup_test_server_with_database() -> (axum_test::TestServer, Arc<Database>) {
     let (server, _, _, database) = setup_test_server_with_pool().await;
     (server, database)

crates/api/tests/e2e_all/fleet_concurrency.rs

@@ -0,0 +1,134 @@
+use crate::common;
+use inference_providers::mock::ResponseTemplate;
+use std::time::Duration;
+use uuid::Uuid;
+
+const INSTANCES: usize = 4;
+const LIMIT: i32 = 2;
+const ATTEMPTS: usize = 8;
+const HOLD: Duration = Duration::from_secs(2);
+const MODEL: &str = "Qwen/Qwen3-30B-A3B-Instruct-2507";
+
+async fn set_rate_limit(database: &database::Database, organization_id: &str, limit: i32) {
+    let organization_id = Uuid::parse_str(organization_id).expect("organization id is a uuid");
+    database
+        .pool()
+        .get()
+        .await
+        .expect("database connection")
+        .execute(
+            "UPDATE organizations SET rate_limit = $1 WHERE id = $2",
+            &[&limit, &organization_id],
+        )
+        .await
+        .expect("rate limit update");
+}
+
+async fn live_lease_count(database: &database::Database, organization_id: &str) -> i64 {
+    let organization_id = Uuid::parse_str(organization_id).expect("organization id is a uuid");
+    database
+        .pool()
+        .get()
+        .await
+        .expect("database connection")
+        .query_one(
+            "SELECT COUNT(*) FROM concurrency_leases
+             WHERE organization_id = $1 AND expires_at > NOW()",
+            &[&organization_id],
+        )
+        .await
+        .expect("lease count")
+        .get(0)
+}
+
+/// Four replicas sharing one database must not admit more than the
+/// organization's limit between them. Without fleet-wide accounting each
+/// replica counts only itself, so this admits INSTANCES * LIMIT.
+#[tokio::test]
+async fn fleet_limit_is_shared_across_replicas() {
+    let (servers, mocks, database) = common::setup_test_fleet(INSTANCES, |config| {
+        config.fleet_concurrency.mode = config::FleetConcurrencyMode::Enforce;
+    })
+    .await;
+
+    for mock in &mocks {
+        mock.set_default_response(ResponseTemplate::new("held").with_hold(HOLD))
+            .await;
+    }
+
+    let org = common::setup_org_with_credits(&servers[0], 100_000_000_000).await;
+    let api_key = common::get_api_key_for_org(&servers[0], org.id.clone()).await;
+    set_rate_limit(&database, &org.id, LIMIT).await;
+
+    let attempts = (0..ATTEMPTS).map(|attempt| {
+        let server = &servers[attempt % INSTANCES];
+        let api_key = api_key.clone();
+        async move {
+            server
+                .post("/v1/chat/completions")
+                .add_header("Authorization", format!("Bearer {api_key}"))
+                .json(&serde_json::json!({
+                    "model": MODEL,
+                    "messages": [{ "role": "user", "content": "hello" }]
+                }))
+                .await
+                .status_code()
+                .as_u16()
+        }
+    });
+
+    let statuses = futures::future::join_all(attempts).await;
+    let admitted = statuses.iter().filter(|status| **status == 200).count();
+    let rejected = statuses.iter().filter(|status| **status == 429).count();
+
+    assert_eq!(
+        admitted, LIMIT as usize,
+        "expected the fleet to admit exactly {LIMIT} of {ATTEMPTS}, got {statuses:?}"
+    );
+    assert_eq!(
+        rejected,
+        ATTEMPTS - LIMIT as usize,
+        "every attempt over the limit should be rejected, got {statuses:?}"
+    );
+}
+
+/// A finished request must give its slot back, or the organization stays
+/// locked out until the leases expire.
+#[tokio::test]
+async fn completed_requests_release_their_lease() {
+    let (servers, mocks, database) = common::setup_test_fleet(1, |config| {
+        config.fleet_concurrency.mode = config::FleetConcurrencyMode::Enforce;
+    })
+    .await;
+
+    mocks[0]
+        .set_default_response(ResponseTemplate::new("done"))
+        .await;
+
+    let org = common::setup_org_with_credits(&servers[0], 100_000_000_000).await;
+    let api_key = common::get_api_key_for_org(&servers[0], org.id.clone()).await;
+    set_rate_limit(&database, &org.id, LIMIT).await;
+
+    for _ in 0..(ATTEMPTS as i32 + LIMIT) {
+        let status = servers[0]
+            .post("/v1/chat/completions")
+            .add_header("Authorization", format!("Bearer {api_key}"))
+            .json(&serde_json::json!({
+                "model": MODEL,
+                "messages": [{ "role": "user", "content": "hello" }]
+            }))
+            .await
+            .status_code();
+        assert_eq!(
+            status, 200,
+            "a released slot should let the next request straight through"
+        );
+    }
+
+    tokio::time::sleep(Duration::from_millis(500)).await;
+    assert_eq!(
+        live_lease_count(&database, &org.id).await,
+        0,
+        "leases outlived the requests that took them"
+    );
+}

crates/api/tests/e2e_all/main.rs

@@ -41,6 +41,7 @@ mod error_msg;
 mod external_providers;
 mod feature_requests;
 mod files;
+mod fleet_concurrency;
 mod function_tools;
 mod general;
 mod glm52_tier_routing;

crates/config/src/types.rs

@@ -28,6 +28,73 @@ pub struct ApiConfig {
     pub aml: AmlConfig,
     pub usage_reporting: UsageReportingConfig,
     pub ita: ItaAttestationConfig,
+    pub fleet_concurrency: FleetConcurrencyConfig,
+}
+
+/// How concurrent-request limits are enforced across replicas.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum FleetConcurrencyMode {
+    /// Each replica counts only its own in-flight requests.
+    #[default]
+    Off,
+    /// Leases are counted fleet-wide but rejection stays per-replica.
+    Shadow,
+    /// Admission is decided by the fleet-wide lease count.
+    Enforce,
+}
+
+#[derive(Debug, Clone)]
+pub struct FleetConcurrencyConfig {
+    pub mode: FleetConcurrencyMode,
+    pub lease_ttl_seconds: u64,
+    pub instance_id: String,
+}
+
+impl Default for FleetConcurrencyConfig {
+    fn default() -> Self {
+        Self {
+            mode: FleetConcurrencyMode::Off,
+            lease_ttl_seconds: 60,
+            instance_id: String::new(),
+        }
+    }
+}
+
+impl FleetConcurrencyConfig {
+    pub fn from_env() -> Self {
+        let requested = env::var("FLEET_CONCURRENCY_MODE").unwrap_or_default();
+        let mode = match requested.to_ascii_lowercase().as_str() {
+            "enforce" => FleetConcurrencyMode::Enforce,
+            "shadow" => FleetConcurrencyMode::Shadow,
+            "" | "off" => FleetConcurrencyMode::Off,
+            _ => {
+                // eprintln for the reason given on the CHUTES_MODELS warning
+                // below. Silence here reads as enforcement being on when it is
+                // not.
+                eprintln!(
+                    "WARN: unrecognised FLEET_CONCURRENCY_MODE '{requested}', \
+                     falling back to off"
+                );
+                FleetConcurrencyMode::Off
+            }
+        };
+
+        Self {
+            mode,
+            lease_ttl_seconds: env::var("FLEET_CONCURRENCY_LEASE_TTL_SECONDS")
+                .ok()
+                .and_then(|value| value.parse::<u64>().ok())
+                .filter(|seconds| *seconds > 0)
+                .unwrap_or(60),
+            instance_id: env::var("FLEET_CONCURRENCY_INSTANCE_ID")
+                .ok()
+                .filter(|value| !value.is_empty())
+                .unwrap_or_else(|| {
+                    let host = env::var("HOSTNAME").unwrap_or_else(|_| "unknown".to_string());
+                    format!("{host}-{}", std::process::id())
+                }),
+        }
+    }
 }
 
 impl ApiConfig {
@@ -60,6 +127,7 @@ impl ApiConfig {
             infra: InfraConfig::from_env(),
             aml: AmlConfig::from_env()?,
             ita: ItaAttestationConfig::from_env()?,
+            fleet_concurrency: FleetConcurrencyConfig::from_env(),
             usage_reporting: UsageReportingConfig::from_env()?,
         })
     }

crates/database/src/migrations/sql/V0075__add_concurrency_leases.sql

@@ -0,0 +1,14 @@
+CREATE TABLE concurrency_leases (
+    id UUID PRIMARY KEY,
+    organization_id UUID NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
+    model_id UUID NOT NULL REFERENCES models(id) ON DELETE CASCADE,
+    instance_id TEXT NOT NULL,
+    acquired_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+    expires_at TIMESTAMPTZ NOT NULL
+);
+
+CREATE INDEX idx_concurrency_leases_org_model
+    ON concurrency_leases (organization_id, model_id, expires_at);
+
+CREATE INDEX idx_concurrency_leases_expires_at
+    ON concurrency_leases (expires_at);