a3s_code_core/llm/
admission.rs

1//! Typed admission control for model-generation transactions.
2//!
3//! Providers differ in how many generations they can actively serve for one
4//! client/account. Callers must not infer that capacity from model names,
5//! endpoint URLs, languages, or observed response text. The provider reports a
6//! typed concurrency contract, and orchestration code turns it into a shared,
7//! cancellation-safe admission gate.
8
9use std::num::NonZeroUsize;
10use std::sync::Arc;
11use std::time::{Duration, Instant};
12use tokio::sync::{OwnedSemaphorePermit, Semaphore};
13use tokio_util::sync::CancellationToken;
14
15/// Bounded active model-generation capacity reported by an
16/// [`LlmClient`](super::LlmClient).
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct ModelGenerationConcurrency {
19    max_concurrency: NonZeroUsize,
20}
21
22impl ModelGenerationConcurrency {
23    /// Conservative contract for providers that do not explicitly advertise
24    /// safe parallel generation.
25    pub const fn single_flight() -> Self {
26        Self {
27            max_concurrency: NonZeroUsize::MIN,
28        }
29    }
30
31    pub const fn bounded(max_concurrency: NonZeroUsize) -> Self {
32        Self { max_concurrency }
33    }
34
35    pub const fn max_concurrency(self) -> NonZeroUsize {
36        self.max_concurrency
37    }
38}
39
40impl Default for ModelGenerationConcurrency {
41    fn default() -> Self {
42        Self::single_flight()
43    }
44}
45
46#[derive(Debug)]
47struct BoundedAdmission {
48    max_concurrency: NonZeroUsize,
49    semaphore: Arc<Semaphore>,
50}
51
52/// Shared admission gate derived from a typed provider concurrency contract.
53///
54/// Clones share the same semaphore. A permit is owned and releases capacity on
55/// every exit path, including future cancellation and task abortion.
56#[derive(Debug, Clone)]
57pub struct ModelGenerationAdmission {
58    bounded: Arc<BoundedAdmission>,
59}
60
61impl ModelGenerationAdmission {
62    pub fn new(concurrency: ModelGenerationConcurrency) -> Self {
63        let max_concurrency = concurrency.max_concurrency();
64        let bounded = Arc::new(BoundedAdmission {
65            max_concurrency,
66            semaphore: Arc::new(Semaphore::new(max_concurrency.get())),
67        });
68        Self { bounded }
69    }
70
71    pub fn concurrency(&self) -> ModelGenerationConcurrency {
72        ModelGenerationConcurrency::bounded(self.bounded.max_concurrency)
73    }
74
75    /// Wait for active-generation capacity without applying an active
76    /// generation deadline to the queue wait.
77    pub async fn acquire(
78        &self,
79        cancellation: &CancellationToken,
80    ) -> Result<ModelGenerationPermit, ModelGenerationAdmissionError> {
81        let queued_at = Instant::now();
82        let acquire = Arc::clone(&self.bounded.semaphore).acquire_owned();
83        tokio::pin!(acquire);
84        let permit = tokio::select! {
85            biased;
86            _ = cancellation.cancelled() => {
87                return Err(ModelGenerationAdmissionError::Cancelled);
88            }
89            permit = &mut acquire => permit.map_err(|_| {
90                ModelGenerationAdmissionError::Closed
91            })?,
92        };
93        Ok(ModelGenerationPermit {
94            admission: Arc::clone(&self.bounded),
95            _bounded: permit,
96            queue_wait: queued_at.elapsed(),
97        })
98    }
99
100    pub(crate) fn owns(&self, permit: &ModelGenerationPermit) -> bool {
101        Arc::ptr_eq(&self.bounded, &permit.admission)
102    }
103
104    #[cfg(test)]
105    fn available_permits(&self) -> usize {
106        self.bounded.semaphore.available_permits()
107    }
108}
109
110impl Default for ModelGenerationAdmission {
111    fn default() -> Self {
112        Self::new(ModelGenerationConcurrency::default())
113    }
114}
115
116/// Owned capacity for one active model-generation transaction.
117#[derive(Debug)]
118#[must_use = "dropping the permit releases model-generation capacity"]
119pub struct ModelGenerationPermit {
120    admission: Arc<BoundedAdmission>,
121    _bounded: OwnedSemaphorePermit,
122    queue_wait: Duration,
123}
124
125impl ModelGenerationPermit {
126    pub fn queue_wait(&self) -> Duration {
127        self.queue_wait
128    }
129}
130
131#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
132pub enum ModelGenerationAdmissionError {
133    #[error("model-generation admission cancelled by caller")]
134    Cancelled,
135    #[error("model-generation admission gate closed")]
136    Closed,
137    #[error("model-generation permit belongs to a different admission gate")]
138    ForeignPermit,
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use std::time::Duration;
145
146    #[tokio::test]
147    async fn cancelling_a_queued_waiter_does_not_consume_capacity() {
148        let admission = ModelGenerationAdmission::new(ModelGenerationConcurrency::single_flight());
149        let holder = admission
150            .acquire(&CancellationToken::new())
151            .await
152            .expect("first permit");
153        assert_eq!(admission.available_permits(), 0);
154
155        let cancellation = CancellationToken::new();
156        let waiter = tokio::spawn({
157            let admission = admission.clone();
158            let cancellation = cancellation.clone();
159            async move { admission.acquire(&cancellation).await }
160        });
161        tokio::task::yield_now().await;
162        cancellation.cancel();
163        assert!(matches!(
164            waiter.await.expect("waiter join"),
165            Err(ModelGenerationAdmissionError::Cancelled)
166        ));
167
168        drop(holder);
169        let replacement = tokio::time::timeout(
170            Duration::from_millis(100),
171            admission.acquire(&CancellationToken::new()),
172        )
173        .await
174        .expect("cancelled waiter must release its queue position")
175        .expect("replacement permit");
176        drop(replacement);
177    }
178}
a3s_code_core/llm/admission.rs

a3s_code_core/llm/
admission.rs
