refactor: isolate dispatch scheduling core

crates/aether-dispatch-core/src/candidate.rs

@@ -0,0 +1,68 @@
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub struct ProviderEndpointRef {
+    pub provider_id: String,
+    pub endpoint_id: String,
+    pub model_id: String,
+    pub selected_provider_model_name: String,
+    pub api_format: String,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub struct KeyRef {
+    pub provider_id: String,
+    pub endpoint_id: String,
+    pub key_id: String,
+    pub model_id: String,
+    pub selected_provider_model_name: String,
+    pub api_format: String,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub struct PoolRef {
+    pub provider_id: String,
+    pub endpoint_id: String,
+    pub model_id: String,
+    pub selected_provider_model_name: String,
+    pub api_format: String,
+    pub pool_group_id: String,
+}
+
+#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub struct DispatchRankFacts {
+    pub provider_priority: i32,
+    pub key_priority: Option<i32>,
+    pub ranking_reason: Option<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub enum DispatchCandidateRef {
+    SingleKey {
+        key: KeyRef,
+        rank: DispatchRankFacts,
+    },
+    PoolRef {
+        pool: PoolRef,
+        rank: DispatchRankFacts,
+    },
+}
+
+impl DispatchCandidateRef {
+    pub fn provider_endpoint(&self) -> ProviderEndpointRef {
+        match self {
+            Self::SingleKey { key, .. } => ProviderEndpointRef {
+                provider_id: key.provider_id.clone(),
+                endpoint_id: key.endpoint_id.clone(),
+                model_id: key.model_id.clone(),
+                selected_provider_model_name: key.selected_provider_model_name.clone(),
+                api_format: key.api_format.clone(),
+            },
+            Self::PoolRef { pool, .. } => ProviderEndpointRef {
+                provider_id: pool.provider_id.clone(),
+                endpoint_id: pool.endpoint_id.clone(),
+                model_id: pool.model_id.clone(),
+                selected_provider_model_name: pool.selected_provider_model_name.clone(),
+                api_format: pool.api_format.clone(),
+            },
+        }
+    }
+}

crates/aether-dispatch-core/src/effects.rs

@@ -0,0 +1,16 @@
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub enum DispatchEffectKind {
+    CandidateFailed,
+    RateLimited,
+    Succeeded,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub struct DispatchEffect {
+    pub kind: DispatchEffectKind,
+    pub provider_id: String,
+    pub endpoint_id: String,
+    pub key_id: Option<String>,
+    pub candidate_index: u32,
+    pub reason: Option<String>,
+}

crates/aether-dispatch-core/src/lib.rs

@@ -0,0 +1,15 @@
+pub mod candidate;
+pub mod effects;
+pub mod pool;
+pub mod sequence;
+
+pub use candidate::{
+    DispatchCandidateRef, DispatchRankFacts, KeyRef, PoolRef, ProviderEndpointRef,
+};
+pub use effects::{DispatchEffect, DispatchEffectKind};
+pub use pool::{
+    run_pool_dispatch_cursor, PoolDispatchCursorOutcome, PoolDispatchError, PoolDispatchPort,
+    PoolDispatchWindow, PoolWindowConfig, DEFAULT_POOL_MAX_SCAN, DEFAULT_POOL_PAGE_SIZE,
+    DEFAULT_POOL_WINDOW_SIZE,
+};
+pub use sequence::{DispatchSequence, DispatchSequenceItem, DispatchSequenceMark};

crates/aether-dispatch-core/src/pool.rs

@@ -0,0 +1,216 @@
+use async_trait::async_trait;
+
+pub const DEFAULT_POOL_WINDOW_SIZE: u32 = 16;
+pub const DEFAULT_POOL_PAGE_SIZE: u32 = 64;
+pub const DEFAULT_POOL_MAX_SCAN: u32 = 512;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub struct PoolWindowConfig {
+    pub window_size: u32,
+    pub page_size: u32,
+    pub max_scan: u32,
+}
+
+impl Default for PoolWindowConfig {
+    fn default() -> Self {
+        Self {
+            window_size: DEFAULT_POOL_WINDOW_SIZE,
+            page_size: DEFAULT_POOL_PAGE_SIZE,
+            max_scan: DEFAULT_POOL_MAX_SCAN,
+        }
+    }
+}
+
+impl PoolWindowConfig {
+    pub fn normalized(self) -> Self {
+        let page_size = self.page_size.max(1);
+        let window_size = self.window_size.max(1).min(page_size);
+        let max_scan = self.max_scan.max(window_size);
+        Self {
+            window_size,
+            page_size,
+            max_scan,
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct PoolDispatchWindow<Candidate> {
+    pub candidates: Vec<Candidate>,
+    pub scanned_count: u32,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct PoolDispatchCursorOutcome<Candidate> {
+    pub candidates: Vec<Candidate>,
+    pub scanned_count: u32,
+    pub exhausted: bool,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
+pub enum PoolDispatchError<Error> {
+    #[error("pool dispatch port failed")]
+    Port(Error),
+}
+
+#[async_trait]
+pub trait PoolDispatchPort {
+    type Candidate: Send;
+    type Error: Send;
+
+    async fn read_page(
+        &mut self,
+        offset: u32,
+        limit: u32,
+    ) -> Result<Vec<Self::Candidate>, Self::Error>;
+
+    async fn rank_and_filter_window(
+        &mut self,
+        candidates: Vec<Self::Candidate>,
+        window_size: u32,
+    ) -> Result<PoolDispatchWindow<Self::Candidate>, Self::Error>;
+}
+
+pub async fn run_pool_dispatch_cursor<Port>(
+    port: &mut Port,
+    config: PoolWindowConfig,
+) -> Result<PoolDispatchCursorOutcome<Port::Candidate>, PoolDispatchError<Port::Error>>
+where
+    Port: PoolDispatchPort + Send,
+{
+    let config = config.normalized();
+    let mut offset = 0_u32;
+    let mut scanned_count = 0_u32;
+
+    while scanned_count < config.max_scan {
+        let limit = config.page_size.min(config.max_scan - scanned_count);
+        let page = port
+            .read_page(offset, limit)
+            .await
+            .map_err(PoolDispatchError::Port)?;
+        if page.is_empty() {
+            return Ok(PoolDispatchCursorOutcome {
+                candidates: Vec::new(),
+                scanned_count,
+                exhausted: true,
+            });
+        }
+
+        let page_len = u32::try_from(page.len()).unwrap_or(u32::MAX);
+        offset = offset.saturating_add(page_len);
+        scanned_count = scanned_count.saturating_add(page_len);
+
+        let window = port
+            .rank_and_filter_window(page, config.window_size)
+            .await
+            .map_err(PoolDispatchError::Port)?;
+        if !window.candidates.is_empty() {
+            return Ok(PoolDispatchCursorOutcome {
+                candidates: window.candidates,
+                scanned_count,
+                exhausted: false,
+            });
+        }
+    }
+
+    Ok(PoolDispatchCursorOutcome {
+        candidates: Vec::new(),
+        scanned_count,
+        exhausted: true,
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::VecDeque;
+
+    use async_trait::async_trait;
+
+    use super::{run_pool_dispatch_cursor, PoolDispatchPort, PoolDispatchWindow, PoolWindowConfig};
+
+    #[derive(Default)]
+    struct TestPort {
+        pages: VecDeque<Vec<u32>>,
+        read_limits: Vec<u32>,
+    }
+
+    #[async_trait]
+    impl PoolDispatchPort for TestPort {
+        type Candidate = u32;
+        type Error = std::convert::Infallible;
+
+        async fn read_page(
+            &mut self,
+            _offset: u32,
+            limit: u32,
+        ) -> Result<Vec<Self::Candidate>, Self::Error> {
+            self.read_limits.push(limit);
+            Ok(self.pages.pop_front().unwrap_or_default())
+        }
+
+        async fn rank_and_filter_window(
+            &mut self,
+            mut candidates: Vec<Self::Candidate>,
+            window_size: u32,
+        ) -> Result<PoolDispatchWindow<Self::Candidate>, Self::Error> {
+            candidates.sort();
+            candidates.truncate(window_size as usize);
+            let scanned_count = u32::try_from(candidates.len()).unwrap_or(u32::MAX);
+            Ok(PoolDispatchWindow {
+                candidates,
+                scanned_count,
+            })
+        }
+    }
+
+    #[tokio::test]
+    async fn small_pool_is_returned_in_one_frozen_window() {
+        let mut port = TestPort {
+            pages: VecDeque::from([vec![3, 1, 2]]),
+            read_limits: Vec::new(),
+        };
+
+        let outcome = run_pool_dispatch_cursor(&mut port, PoolWindowConfig::default())
+            .await
+            .unwrap();
+
+        assert_eq!(outcome.candidates, [1, 2, 3]);
+        assert_eq!(outcome.scanned_count, 3);
+        assert_eq!(port.read_limits, [64]);
+    }
+
+    #[tokio::test]
+    async fn large_pool_returns_bounded_window() {
+        let mut port = TestPort {
+            pages: VecDeque::from([(0..100).rev().collect::<Vec<_>>()]),
+            read_limits: Vec::new(),
+        };
+
+        let outcome = run_pool_dispatch_cursor(&mut port, PoolWindowConfig::default())
+            .await
+            .unwrap();
+
+        assert_eq!(outcome.candidates.len(), 16);
+        assert_eq!(outcome.candidates[0], 0);
+        assert_eq!(outcome.candidates[15], 15);
+        assert_eq!(port.read_limits, [64]);
+    }
+
+    #[tokio::test]
+    async fn max_scan_caps_page_reads() {
+        let mut port = TestPort {
+            pages: VecDeque::from([Vec::new()]),
+            read_limits: Vec::new(),
+        };
+        let config = PoolWindowConfig {
+            window_size: 16,
+            page_size: 64,
+            max_scan: 32,
+        };
+
+        let outcome = run_pool_dispatch_cursor(&mut port, config).await.unwrap();
+
+        assert!(outcome.exhausted);
+        assert_eq!(port.read_limits, [32]);
+    }
+}

crates/aether-dispatch-core/src/sequence.rs

@@ -0,0 +1,110 @@
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub struct DispatchSequenceItem<Candidate> {
+    pub candidate_index: u32,
+    pub retry_index: u32,
+    pub candidate: Candidate,
+    pub mark: DispatchSequenceMark,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub enum DispatchSequenceMark {
+    Pending,
+    Failed,
+    Succeeded,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct DispatchSequence<Candidate> {
+    items: Vec<DispatchSequenceItem<Candidate>>,
+    cursor: usize,
+}
+
+impl<Candidate> DispatchSequence<Candidate> {
+    pub fn new(items: Vec<DispatchSequenceItem<Candidate>>) -> Self {
+        Self { items, cursor: 0 }
+    }
+
+    pub fn from_candidates(candidates: Vec<Candidate>) -> Self {
+        let items = candidates
+            .into_iter()
+            .enumerate()
+            .map(|(index, candidate)| DispatchSequenceItem {
+                candidate_index: u32::try_from(index).unwrap_or(u32::MAX),
+                retry_index: 0,
+                candidate,
+                mark: DispatchSequenceMark::Pending,
+            })
+            .collect();
+        Self::new(items)
+    }
+
+    pub fn peek_current(&self) -> Option<&DispatchSequenceItem<Candidate>> {
+        self.items.get(self.cursor)
+    }
+
+    #[allow(clippy::should_implement_trait)]
+    pub fn next(&mut self) -> Option<&DispatchSequenceItem<Candidate>> {
+        while self
+            .items
+            .get(self.cursor)
+            .is_some_and(|item| item.mark != DispatchSequenceMark::Pending)
+        {
+            self.cursor = self.cursor.saturating_add(1);
+        }
+        self.items.get(self.cursor)
+    }
+
+    pub fn mark_failed(&mut self) -> Option<&DispatchSequenceItem<Candidate>> {
+        self.mark_current(DispatchSequenceMark::Failed)
+    }
+
+    pub fn mark_succeeded(&mut self) -> Option<&DispatchSequenceItem<Candidate>> {
+        self.mark_current(DispatchSequenceMark::Succeeded)
+    }
+
+    pub fn len(&self) -> usize {
+        self.items.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.items.is_empty()
+    }
+
+    pub fn cursor(&self) -> usize {
+        self.cursor
+    }
+
+    pub fn items(&self) -> &[DispatchSequenceItem<Candidate>] {
+        &self.items
+    }
+
+    fn mark_current(
+        &mut self,
+        mark: DispatchSequenceMark,
+    ) -> Option<&DispatchSequenceItem<Candidate>> {
+        let item = self.items.get_mut(self.cursor)?;
+        item.mark = mark;
+        self.cursor = self.cursor.saturating_add(1);
+        self.items.get(self.cursor)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{DispatchSequence, DispatchSequenceMark};
+
+    #[test]
+    fn mark_failed_advances_without_reordering() {
+        let mut sequence = DispatchSequence::from_candidates(vec!["a", "b", "c"]);
+
+        assert_eq!(sequence.next().map(|item| item.candidate), Some("a"));
+        assert_eq!(sequence.mark_failed().map(|item| item.candidate), Some("b"));
+        assert_eq!(sequence.next().map(|item| item.candidate), Some("b"));
+        assert_eq!(sequence.mark_failed().map(|item| item.candidate), Some("c"));
+        assert_eq!(sequence.next().map(|item| item.candidate), Some("c"));
+
+        assert_eq!(sequence.items()[0].mark, DispatchSequenceMark::Failed);
+        assert_eq!(sequence.items()[1].mark, DispatchSequenceMark::Failed);
+        assert_eq!(sequence.items()[2].mark, DispatchSequenceMark::Pending);
+    }
+}