refactor ai serving modules and crates

apps/aether-gateway/src/ai_serving/planner/state/auth.rs

@@ -0,0 +1,18 @@
+use super::PlannerAppState;
+pub(crate) use crate::data::auth::GatewayAuthApiKeySnapshot;
+use crate::GatewayError;
+
+impl<'a> PlannerAppState<'a> {
+    pub(crate) async fn read_auth_api_key_snapshot(
+        self,
+        user_id: &str,
+        api_key_id: &str,
+        now_unix_secs: u64,
+    ) -> Result<Option<GatewayAuthApiKeySnapshot>, GatewayError> {
+        self.app()
+            .data
+            .read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
+            .await
+            .map_err(|err| GatewayError::Internal(err.to_string()))
+    }
+}

apps/aether-gateway/src/ai_serving/planner/state/candidate_runtime.rs

@@ -0,0 +1,89 @@
+use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
+use serde_json::Value;
+
+use super::PlannerAppState;
+
+impl<'a> PlannerAppState<'a> {
+    pub(crate) async fn resolve_request_candidate_required_capabilities(
+        self,
+        user_id: &str,
+        api_key_id: &str,
+        requested_model: Option<&str>,
+        explicit_required_capabilities: Option<&Value>,
+    ) -> Option<Value> {
+        crate::request_candidate_runtime::resolve_request_candidate_required_capabilities(
+            self.app(),
+            user_id,
+            api_key_id,
+            requested_model,
+            explicit_required_capabilities,
+        )
+        .await
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) async fn persist_available_local_candidate(
+        self,
+        trace_id: &str,
+        user_id: &str,
+        api_key_id: &str,
+        candidate: &SchedulerMinimalCandidateSelectionCandidate,
+        candidate_index: u32,
+        retry_index: u32,
+        candidate_id: &str,
+        required_capabilities: Option<&Value>,
+        extra_data: Option<Value>,
+        created_at_unix_ms: u64,
+        error_context: &'static str,
+    ) -> String {
+        crate::request_candidate_runtime::persist_available_local_candidate(
+            self.app(),
+            trace_id,
+            user_id,
+            api_key_id,
+            candidate,
+            candidate_index,
+            retry_index,
+            candidate_id,
+            required_capabilities,
+            extra_data,
+            created_at_unix_ms,
+            error_context,
+        )
+        .await
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) async fn persist_skipped_local_candidate(
+        self,
+        trace_id: &str,
+        user_id: &str,
+        api_key_id: &str,
+        candidate: &SchedulerMinimalCandidateSelectionCandidate,
+        candidate_index: u32,
+        retry_index: u32,
+        candidate_id: &str,
+        required_capabilities: Option<&Value>,
+        skip_reason: &str,
+        extra_data: Option<Value>,
+        finished_at_unix_ms: u64,
+        error_context: &'static str,
+    ) {
+        crate::request_candidate_runtime::persist_skipped_local_candidate(
+            self.app(),
+            trace_id,
+            user_id,
+            api_key_id,
+            candidate,
+            candidate_index,
+            retry_index,
+            candidate_id,
+            required_capabilities,
+            skip_reason,
+            extra_data,
+            finished_at_unix_ms,
+            error_context,
+        )
+        .await
+    }
+}

apps/aether-gateway/src/ai_serving/planner/state/executor.rs

@@ -0,0 +1,24 @@
+use aether_contracts::ExecutionPlan;
+use serde_json::Value;
+
+use super::PlannerAppState;
+
+impl<'a> PlannerAppState<'a> {
+    pub(crate) async fn mark_unused_local_candidate_items<T, FPlan, FContext>(
+        self,
+        remaining: Vec<T>,
+        plan: FPlan,
+        report_context: FContext,
+    ) where
+        FPlan: Fn(&T) -> &ExecutionPlan,
+        FContext: Fn(&T) -> Option<&Value>,
+    {
+        crate::executor::mark_unused_local_candidate_items(
+            self.app(),
+            remaining,
+            plan,
+            report_context,
+        )
+        .await
+    }
+}

apps/aether-gateway/src/ai_serving/planner/state/mod.rs

@@ -0,0 +1,25 @@
+use crate::AppState;
+
+mod auth;
+mod candidate_runtime;
+mod executor;
+mod scheduler;
+mod transport;
+
+pub(crate) use self::auth::GatewayAuthApiKeySnapshot;
+pub(crate) use self::transport::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth};
+
+#[derive(Clone, Copy)]
+pub(crate) struct PlannerAppState<'a> {
+    app: &'a AppState,
+}
+
+impl<'a> PlannerAppState<'a> {
+    pub(crate) fn new(app: &'a AppState) -> Self {
+        Self { app }
+    }
+
+    pub(crate) fn app(self) -> &'a AppState {
+        self.app
+    }
+}

apps/aether-gateway/src/ai_serving/planner/state/scheduler.rs

@@ -0,0 +1,125 @@
+use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
+use std::time::Duration;
+use tokio::time::Instant;
+
+use super::{GatewayAuthApiKeySnapshot, PlannerAppState};
+use crate::clock::current_unix_secs;
+use crate::constants::{
+    API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS, API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS,
+};
+use crate::scheduler::candidate::SchedulerSkippedCandidate;
+use crate::GatewayError;
+
+impl<'a> PlannerAppState<'a> {
+    pub(crate) async fn list_selectable_candidates(
+        self,
+        api_format: &str,
+        global_model_name: &str,
+        require_streaming: bool,
+        required_capabilities: Option<&serde_json::Value>,
+        auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+        now_unix_secs: u64,
+    ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
+        crate::scheduler::candidate::list_selectable_candidates(
+            self.app().data.as_ref(),
+            self.app(),
+            api_format,
+            global_model_name,
+            require_streaming,
+            required_capabilities,
+            auth_snapshot,
+            now_unix_secs,
+        )
+        .await
+    }
+
+    pub(crate) async fn list_selectable_candidates_with_skip_reasons(
+        self,
+        api_format: &str,
+        global_model_name: &str,
+        require_streaming: bool,
+        required_capabilities: Option<&serde_json::Value>,
+        auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+        now_unix_secs: u64,
+    ) -> Result<
+        (
+            Vec<SchedulerMinimalCandidateSelectionCandidate>,
+            Vec<SchedulerSkippedCandidate>,
+        ),
+        GatewayError,
+    > {
+        let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
+        let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
+        let wait_deadline = Instant::now() + wait_timeout;
+        let mut attempt_now_unix_secs = now_unix_secs;
+
+        loop {
+            let result = crate::scheduler::candidate::list_selectable_candidates_with_skip_reasons(
+                self.app().data.as_ref(),
+                self.app(),
+                api_format,
+                global_model_name,
+                require_streaming,
+                required_capabilities,
+                auth_snapshot,
+                attempt_now_unix_secs,
+            )
+            .await?;
+
+            if !crate::scheduler::candidate::is_exact_all_skipped_by_auth_limit(
+                &result.0, &result.1,
+            ) {
+                return Ok(result);
+            }
+
+            let now = Instant::now();
+            if now >= wait_deadline {
+                return Ok(result);
+            }
+
+            let remaining = wait_deadline.duration_since(now);
+            tokio::time::sleep(wait_interval.min(remaining)).await;
+            attempt_now_unix_secs = current_unix_secs();
+        }
+    }
+
+    pub(crate) async fn list_selectable_candidates_for_required_capability_without_requested_model(
+        self,
+        candidate_api_format: &str,
+        required_capability: &str,
+        require_streaming: bool,
+        auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
+        now_unix_secs: u64,
+    ) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
+        let wait_timeout = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS);
+        let wait_interval = Duration::from_millis(API_KEY_CONCURRENCY_WAIT_POLL_INTERVAL_MS.max(1));
+        let wait_deadline = Instant::now() + wait_timeout;
+        let mut attempt_now_unix_secs = now_unix_secs;
+
+        loop {
+            let (result, auth_limit_blocked) = crate::scheduler::candidate::list_selectable_candidates_for_required_capability_without_requested_model_with_auth_limit_signal(
+                self.app().data.as_ref(),
+                self.app(),
+                candidate_api_format,
+                required_capability,
+                require_streaming,
+                auth_snapshot,
+                attempt_now_unix_secs,
+            )
+            .await?;
+
+            if !auth_limit_blocked {
+                return Ok(result);
+            }
+
+            let now = Instant::now();
+            if now >= wait_deadline {
+                return Ok(result);
+            }
+
+            let remaining = wait_deadline.duration_since(now);
+            tokio::time::sleep(wait_interval.min(remaining)).await;
+            attempt_now_unix_secs = current_unix_secs();
+        }
+    }
+}

apps/aether-gateway/src/ai_serving/planner/state/transport.rs

@@ -0,0 +1,25 @@
+use super::PlannerAppState;
+pub(crate) use crate::ai_serving::transport::{
+    GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
+};
+use crate::GatewayError;
+
+impl<'a> PlannerAppState<'a> {
+    pub(crate) async fn read_provider_transport_snapshot(
+        self,
+        provider_id: &str,
+        endpoint_id: &str,
+        key_id: &str,
+    ) -> Result<Option<GatewayProviderTransportSnapshot>, GatewayError> {
+        self.app()
+            .read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
+            .await
+    }
+
+    pub(crate) async fn resolve_local_oauth_request_auth(
+        self,
+        transport: &GatewayProviderTransportSnapshot,
+    ) -> Result<Option<LocalResolvedOAuthRequestAuth>, GatewayError> {
+        self.app().resolve_local_oauth_request_auth(transport).await
+    }
+}