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 thiserror::Error;
13use tokio::sync::{OwnedSemaphorePermit, Semaphore};
14use tokio_util::sync::CancellationToken;
15
16const MODEL_GENERATION_POOL_MAX_COMPONENT_BYTES: usize = 256;
17
18/// Errors returned while deriving a provider/model capacity pool identity.
19#[derive(Debug, Clone, PartialEq, Eq, Error)]
20pub enum ModelGenerationPoolError {
21    #[error("model-generation pool {field} is empty or exceeds {limit} bytes")]
22    InvalidComponent { field: &'static str, limit: usize },
23    #[error("model-generation pool {field} contains a control character")]
24    ControlCharacter { field: &'static str },
25    #[error("model-generation pool endpoint is invalid or has no host")]
26    InvalidEndpoint,
27    #[error("model-generation pool endpoint must not contain credentials")]
28    EndpointCredentials,
29    #[error("model-generation pool identity is invalid: {0}")]
30    InvalidIdentity(String),
31    #[error("model-generation pool concurrency must be greater than zero")]
32    InvalidConcurrency,
33}
34
35/// Digest-only provider/model capacity metadata.
36///
37/// The pool is intentionally a descriptor, not an executor or semaphore. Its
38/// identity binds the non-secret routing facts that share a provider capacity
39/// budget; the scheduler or a local admission gate owns the live reservation.
40#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
41#[serde(rename_all = "camelCase", deny_unknown_fields)]
42pub struct ModelGenerationPool {
43    /// Domain-separated digest of provider, model, endpoint origin, and
44    /// optional non-secret account scope.
45    pub identity: crate::execution_identity::ExecutionIdentityV1,
46    /// Maximum active generations for this pool.
47    pub max_concurrency: NonZeroUsize,
48}
49
50impl ModelGenerationPool {
51    /// Build a pool from an already-derived identity.
52    pub fn new(
53        identity: crate::execution_identity::ExecutionIdentityV1,
54        max_concurrency: NonZeroUsize,
55    ) -> Result<Self, ModelGenerationPoolError> {
56        identity
57            .validate()
58            .map_err(|error| ModelGenerationPoolError::InvalidIdentity(error.to_string()))?;
59        Ok(Self {
60            identity,
61            max_concurrency,
62        })
63    }
64
65    /// Derive a stable pool from non-secret provider routing metadata.
66    ///
67    /// Endpoint paths, queries, fragments, and credentials are deliberately
68    /// discarded. Two clients that address the same origin/model therefore
69    /// share one capacity key even when their API path configuration differs.
70    pub fn for_client(
71        provider: &str,
72        model: &str,
73        endpoint: Option<&str>,
74        account_id: Option<&str>,
75        concurrency: ModelGenerationConcurrency,
76    ) -> Result<Self, ModelGenerationPoolError> {
77        let provider = bounded_component("provider", provider)?;
78        let model = bounded_component("model", model)?;
79        let endpoint_origin = endpoint.map(endpoint_origin).transpose()?;
80        let account_id = account_id
81            .map(|value| bounded_component("accountId", value))
82            .transpose()?;
83        let identity = crate::execution_identity::ExecutionIdentityV1::derive(
84            crate::execution_identity::MODEL_GENERATION_POOL_IDENTITY_DOMAIN_V1,
85            &serde_json::json!({
86                "provider": provider,
87                "model": model,
88                "endpoint_origin": endpoint_origin,
89                "account_id": account_id,
90            }),
91        )
92        .map_err(|error| ModelGenerationPoolError::InvalidIdentity(error.to_string()))?;
93        Self::new(identity, concurrency.max_concurrency())
94    }
95
96    /// Convenience constructor for clients with a concrete endpoint URL.
97    pub fn for_endpoint(
98        provider: &str,
99        model: &str,
100        endpoint: &str,
101        concurrency: ModelGenerationConcurrency,
102    ) -> Result<Self, ModelGenerationPoolError> {
103        Self::for_client(provider, model, Some(endpoint), None, concurrency)
104    }
105
106    /// Convenience constructor for account-scoped clients.
107    pub fn for_account_endpoint(
108        provider: &str,
109        model: &str,
110        endpoint: &str,
111        account_id: &str,
112        concurrency: ModelGenerationConcurrency,
113    ) -> Result<Self, ModelGenerationPoolError> {
114        Self::for_client(
115            provider,
116            model,
117            Some(endpoint),
118            Some(account_id),
119            concurrency,
120        )
121    }
122
123    pub fn identity(&self) -> &crate::execution_identity::ExecutionIdentityV1 {
124        &self.identity
125    }
126
127    pub const fn max_concurrency(&self) -> NonZeroUsize {
128        self.max_concurrency
129    }
130
131    pub fn validate(&self) -> Result<(), ModelGenerationPoolError> {
132        self.identity
133            .validate()
134            .map_err(|error| ModelGenerationPoolError::InvalidIdentity(error.to_string()))?;
135        if self.max_concurrency.get() == 0 {
136            return Err(ModelGenerationPoolError::InvalidConcurrency);
137        }
138        Ok(())
139    }
140}
141
142fn bounded_component(field: &'static str, value: &str) -> Result<String, ModelGenerationPoolError> {
143    let value = value.trim();
144    if value.is_empty() || value.len() > MODEL_GENERATION_POOL_MAX_COMPONENT_BYTES {
145        return Err(ModelGenerationPoolError::InvalidComponent {
146            field,
147            limit: MODEL_GENERATION_POOL_MAX_COMPONENT_BYTES,
148        });
149    }
150    if value
151        .chars()
152        .any(|character| character.is_control() || matches!(character, '\u{2028}' | '\u{2029}'))
153    {
154        return Err(ModelGenerationPoolError::ControlCharacter { field });
155    }
156    Ok(value.to_string())
157}
158
159fn endpoint_origin(value: &str) -> Result<String, ModelGenerationPoolError> {
160    let parsed =
161        url::Url::parse(value.trim()).map_err(|_| ModelGenerationPoolError::InvalidEndpoint)?;
162    if parsed.host_str().is_none() || parsed.scheme().is_empty() {
163        return Err(ModelGenerationPoolError::InvalidEndpoint);
164    }
165    if !parsed.username().is_empty() || parsed.password().is_some() {
166        return Err(ModelGenerationPoolError::EndpointCredentials);
167    }
168    let origin = parsed.origin().ascii_serialization();
169    if origin == "null" {
170        return Err(ModelGenerationPoolError::InvalidEndpoint);
171    }
172    Ok(origin)
173}
174
175/// Bounded active model-generation capacity reported by an
176/// [`LlmClient`](super::LlmClient).
177#[derive(Debug, Clone, Copy, PartialEq, Eq)]
178pub struct ModelGenerationConcurrency {
179    max_concurrency: NonZeroUsize,
180}
181
182impl ModelGenerationConcurrency {
183    /// Conservative contract for providers that do not explicitly advertise
184    /// safe parallel generation.
185    pub const fn single_flight() -> Self {
186        Self {
187            max_concurrency: NonZeroUsize::MIN,
188        }
189    }
190
191    pub const fn bounded(max_concurrency: NonZeroUsize) -> Self {
192        Self { max_concurrency }
193    }
194
195    pub const fn max_concurrency(self) -> NonZeroUsize {
196        self.max_concurrency
197    }
198}
199
200impl Default for ModelGenerationConcurrency {
201    fn default() -> Self {
202        Self::single_flight()
203    }
204}
205
206#[derive(Debug)]
207struct BoundedAdmission {
208    max_concurrency: NonZeroUsize,
209    semaphore: Arc<Semaphore>,
210    scheduler: Option<Arc<SchedulerBinding>>,
211    /// Optional immutable provider pool descriptor used by host diagnostics.
212    /// The descriptor contains only a digest identity and numeric capacity.
213    pool: Option<ModelGenerationPool>,
214}
215
216#[derive(Debug)]
217struct SchedulerBinding {
218    scheduler: Arc<crate::task_scheduler::TaskScheduler>,
219    quota: crate::task_scheduler::TaskSchedulerQuota,
220    priority: crate::task_scheduler::TaskPriority,
221    label: String,
222}
223
224/// Shared admission gate derived from a typed provider concurrency contract.
225///
226/// Clones share the same semaphore. A permit is owned and releases capacity on
227/// every exit path, including future cancellation and task abortion.
228#[derive(Debug, Clone)]
229pub struct ModelGenerationAdmission {
230    bounded: Arc<BoundedAdmission>,
231}
232
233/// Read-only, secret-free configuration and occupancy evidence for one model
234/// generation pool.
235///
236/// `local_reserved` includes permits waiting for the shared scheduler after
237/// acquiring the local client gate; it is therefore intentionally distinct
238/// from provider calls that have reached the transport. The optional scheduler
239/// projection carries the same pool identity and retains a bounded recent
240/// health epoch after the final reservation is released.
241#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
242#[serde(rename_all = "camelCase", deny_unknown_fields)]
243pub struct ModelGenerationPoolHealthSnapshot {
244    /// Immutable provider/model capacity descriptor.
245    pub pool: ModelGenerationPool,
246    /// Effective local gate limit for this admission facade.
247    pub local_max_concurrency: usize,
248    /// Local permits currently reserved by this facade.
249    pub local_reserved: usize,
250    /// Local permits immediately available to this facade.
251    pub local_available: usize,
252    /// Shared scheduler health for this pool, when the facade is scheduler-bound.
253    #[serde(skip_serializing_if = "Option::is_none")]
254    pub scheduler: Option<crate::task_scheduler::TaskSchedulerQuotaHealthSnapshot>,
255}
256
257impl ModelGenerationAdmission {
258    pub fn new(concurrency: ModelGenerationConcurrency) -> Self {
259        let max_concurrency = concurrency.max_concurrency();
260        let bounded = Arc::new(BoundedAdmission {
261            max_concurrency,
262            semaphore: Arc::new(Semaphore::new(max_concurrency.get())),
263            scheduler: None,
264            pool: None,
265        });
266        Self { bounded }
267    }
268
269    /// Attach a provider/model quota to this gate without creating another
270    /// queue. Each permit reserves the quota through the existing scheduler
271    /// actor, while the local semaphore retains the client's own contract.
272    pub fn with_scheduler_quota(
273        self,
274        scheduler: Arc<crate::task_scheduler::TaskScheduler>,
275        quota: crate::task_scheduler::TaskSchedulerQuota,
276        priority: crate::task_scheduler::TaskPriority,
277        label: impl Into<String>,
278    ) -> Result<Self, crate::task_scheduler::TaskSchedulerError> {
279        self.with_scheduler_binding(scheduler, quota, priority, label.into(), None)
280    }
281
282    /// Attach a provider pool descriptor and its shared scheduler quota.
283    ///
284    /// The pool is configuration evidence only; live reservations continue to
285    /// be owned by the existing scheduler actor and local semaphore.
286    pub fn with_model_generation_pool(
287        self,
288        scheduler: Arc<crate::task_scheduler::TaskScheduler>,
289        pool: ModelGenerationPool,
290        priority: crate::task_scheduler::TaskPriority,
291        label: impl Into<String>,
292    ) -> Result<Self, crate::task_scheduler::TaskSchedulerError> {
293        pool.validate().map_err(|error| {
294            crate::task_scheduler::TaskSchedulerError::InvalidConfig(error.to_string())
295        })?;
296        let quota = crate::task_scheduler::TaskSchedulerQuota::new(
297            pool.identity.clone(),
298            pool.max_concurrency.get(),
299        )?;
300        self.with_scheduler_binding(scheduler, quota, priority, label.into(), Some(pool))
301    }
302
303    fn with_scheduler_binding(
304        self,
305        scheduler: Arc<crate::task_scheduler::TaskScheduler>,
306        quota: crate::task_scheduler::TaskSchedulerQuota,
307        priority: crate::task_scheduler::TaskPriority,
308        label: String,
309        pool: Option<ModelGenerationPool>,
310    ) -> Result<Self, crate::task_scheduler::TaskSchedulerError> {
311        quota.validate()?;
312        if let Some(pool) = &pool {
313            if pool.identity != quota.identity || pool.max_concurrency.get() != quota.max_active {
314                return Err(crate::task_scheduler::TaskSchedulerError::InvalidConfig(
315                    "model-generation pool and scheduler quota do not match".to_string(),
316                ));
317            }
318        }
319        let bounded = Arc::new(BoundedAdmission {
320            max_concurrency: self.bounded.max_concurrency,
321            semaphore: Arc::clone(&self.bounded.semaphore),
322            scheduler: Some(Arc::new(SchedulerBinding {
323                scheduler,
324                quota,
325                priority,
326                label,
327            })),
328            pool,
329        });
330        Ok(Self { bounded })
331    }
332
333    /// Copy the scheduler-backed provider reservation from another admission
334    /// while retaining this gate's local concurrency. This is used by nested
335    /// workflow steps that impose a tighter local limit but must still count
336    /// against the session/provider pool in the one shared scheduler actor.
337    pub(crate) fn with_scheduler_quota_from(
338        self,
339        source: &Self,
340        label: impl Into<String>,
341    ) -> Result<Self, crate::task_scheduler::TaskSchedulerError> {
342        let Some(binding) = source.bounded.scheduler.as_ref() else {
343            return Ok(self);
344        };
345        self.with_scheduler_binding(
346            Arc::clone(&binding.scheduler),
347            binding.quota.clone(),
348            binding.priority,
349            label.into(),
350            source.bounded.pool.clone(),
351        )
352    }
353
354    pub fn concurrency(&self) -> ModelGenerationConcurrency {
355        ModelGenerationConcurrency::bounded(self.bounded.max_concurrency)
356    }
357
358    /// Whether every permit also reserves a quota through the shared
359    /// scheduler actor.
360    pub(crate) fn has_scheduler_quota(&self) -> bool {
361        self.bounded.scheduler.is_some()
362    }
363
364    /// Whether this admission publishes a typed product
365    /// [`ModelGenerationPool`] (OPT-POOL1). Scheduler quota without a pool is
366    /// not product shared-capacity evidence.
367    #[cfg(test)]
368    pub(crate) fn publishes_model_generation_pool(&self) -> bool {
369        self.bounded.pool.is_some()
370    }
371
372    /// Return secret-free pool configuration and point-in-time health.
373    ///
374    /// `None` means this admission was created for a custom client that did
375    /// not publish a [`ModelGenerationPool`].
376    pub async fn pool_health(
377        &self,
378    ) -> Result<Option<ModelGenerationPoolHealthSnapshot>, crate::task_scheduler::TaskSchedulerError>
379    {
380        let Some(pool) = self.bounded.pool.clone() else {
381            return Ok(None);
382        };
383        let local_max_concurrency = self.bounded.max_concurrency.get();
384        let local_available = self.bounded.semaphore.available_permits();
385        let local_reserved = local_max_concurrency.saturating_sub(local_available);
386        let scheduler = match self.bounded.scheduler.as_ref() {
387            Some(binding) => Some(binding.scheduler.quota_health(&binding.quota).await?),
388            None => None,
389        };
390        Ok(Some(ModelGenerationPoolHealthSnapshot {
391            pool,
392            local_max_concurrency,
393            local_reserved,
394            local_available,
395            scheduler,
396        }))
397    }
398
399    /// Wait for active-generation capacity without applying an active
400    /// generation deadline to the queue wait.
401    pub async fn acquire(
402        &self,
403        cancellation: &CancellationToken,
404    ) -> Result<ModelGenerationPermit, ModelGenerationAdmissionError> {
405        let queued_at = Instant::now();
406        let acquire = Arc::clone(&self.bounded.semaphore).acquire_owned();
407        tokio::pin!(acquire);
408        let permit = tokio::select! {
409            biased;
410            _ = cancellation.cancelled() => {
411                return Err(ModelGenerationAdmissionError::Cancelled);
412            }
413            permit = &mut acquire => permit.map_err(|_| {
414                ModelGenerationAdmissionError::Closed
415            })?,
416        };
417        // Take the local contract first. A scheduler quota-only lease is a
418        // scarcer shared resource; acquiring it second prevents a session
419        // waiting on its own local semaphore from hoarding provider capacity.
420        let scheduler_lease = if let Some(binding) = self.bounded.scheduler.as_ref() {
421            Some(
422                binding
423                    .scheduler
424                    .acquire_quota(
425                        binding.priority,
426                        binding.label.clone(),
427                        &binding.quota,
428                        cancellation,
429                    )
430                    .await
431                    .map_err(ModelGenerationAdmissionError::from_scheduler_error)?,
432            )
433        } else {
434            None
435        };
436        Ok(ModelGenerationPermit {
437            admission: Arc::clone(&self.bounded),
438            _bounded: permit,
439            _scheduler: scheduler_lease,
440            queue_wait: queued_at.elapsed(),
441        })
442    }
443
444    pub(crate) fn owns(&self, permit: &ModelGenerationPermit) -> bool {
445        Arc::ptr_eq(&self.bounded, &permit.admission)
446    }
447
448    #[cfg(test)]
449    fn available_permits(&self) -> usize {
450        self.bounded.semaphore.available_permits()
451    }
452}
453
454impl Default for ModelGenerationAdmission {
455    fn default() -> Self {
456        Self::new(ModelGenerationConcurrency::default())
457    }
458}
459
460/// Owned capacity for one active model-generation transaction.
461#[derive(Debug)]
462#[must_use = "dropping the permit releases model-generation capacity"]
463pub struct ModelGenerationPermit {
464    admission: Arc<BoundedAdmission>,
465    _bounded: OwnedSemaphorePermit,
466    _scheduler: Option<crate::task_scheduler::TaskLease>,
467    queue_wait: Duration,
468}
469
470impl ModelGenerationPermit {
471    pub fn queue_wait(&self) -> Duration {
472        self.queue_wait
473    }
474}
475
476#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
477pub enum ModelGenerationAdmissionError {
478    #[error("model-generation admission cancelled by caller")]
479    Cancelled,
480    #[error("model-generation admission gate closed")]
481    Closed,
482    #[error("model-generation permit belongs to a different admission gate")]
483    ForeignPermit,
484    #[error("scheduler-backed model-generation admission failed: {0}")]
485    Scheduler(String),
486}
487
488impl ModelGenerationAdmissionError {
489    fn from_scheduler_error(error: crate::task_scheduler::TaskSchedulerError) -> Self {
490        match error {
491            crate::task_scheduler::TaskSchedulerError::Cancelled => Self::Cancelled,
492            crate::task_scheduler::TaskSchedulerError::Closed => Self::Closed,
493            crate::task_scheduler::TaskSchedulerError::AtCapacity { limit } => {
494                Self::Scheduler(format!("scheduler admission queue is full (limit {limit})"))
495            }
496            crate::task_scheduler::TaskSchedulerError::InvalidConfig(message) => {
497                Self::Scheduler(message)
498            }
499        }
500    }
501}
502
503#[cfg(test)]
504mod tests {
505    use super::*;
506    use std::time::Duration;
507
508    #[tokio::test]
509    async fn cancelling_a_queued_waiter_does_not_consume_capacity() {
510        let admission = ModelGenerationAdmission::new(ModelGenerationConcurrency::single_flight());
511        let holder = admission
512            .acquire(&CancellationToken::new())
513            .await
514            .expect("first permit");
515        assert_eq!(admission.available_permits(), 0);
516
517        let cancellation = CancellationToken::new();
518        let waiter = tokio::spawn({
519            let admission = admission.clone();
520            let cancellation = cancellation.clone();
521            async move { admission.acquire(&cancellation).await }
522        });
523        tokio::task::yield_now().await;
524        cancellation.cancel();
525        assert!(matches!(
526            waiter.await.expect("waiter join"),
527            Err(ModelGenerationAdmissionError::Cancelled)
528        ));
529
530        drop(holder);
531        let replacement = tokio::time::timeout(
532            Duration::from_millis(100),
533            admission.acquire(&CancellationToken::new()),
534        )
535        .await
536        .expect("cancelled waiter must release its queue position")
537        .expect("replacement permit");
538        drop(replacement);
539    }
540
541    #[test]
542    fn pool_identity_uses_only_non_secret_endpoint_origin() {
543        let concurrency = ModelGenerationConcurrency::bounded(NonZeroUsize::new(2).unwrap());
544        let first = ModelGenerationPool::for_client(
545            "openai-compatible",
546            "model-a",
547            Some("https://example.test:443/v1/chat/completions?key=ignored"),
548            None,
549            concurrency,
550        )
551        .unwrap();
552        let second = ModelGenerationPool::for_client(
553            "openai-compatible",
554            "model-a",
555            Some("https://example.test:443/another-path"),
556            None,
557            concurrency,
558        )
559        .unwrap();
560        assert_eq!(first.identity, second.identity);
561        assert_eq!(first.max_concurrency(), NonZeroUsize::new(2).unwrap());
562        let encoded = serde_json::to_string(&first).unwrap();
563        assert!(!encoded.contains("ignored"));
564        assert!(!encoded.contains("example.test:443/v1"));
565    }
566
567    #[test]
568    fn pool_identity_separates_provider_model_and_account() {
569        let concurrency = ModelGenerationConcurrency::single_flight();
570        let base = ModelGenerationPool::for_endpoint(
571            "provider-a",
572            "model-a",
573            "https://example.test",
574            concurrency,
575        )
576        .unwrap();
577        let provider = ModelGenerationPool::for_endpoint(
578            "provider-b",
579            "model-a",
580            "https://example.test",
581            concurrency,
582        )
583        .unwrap();
584        let model = ModelGenerationPool::for_endpoint(
585            "provider-a",
586            "model-b",
587            "https://example.test",
588            concurrency,
589        )
590        .unwrap();
591        let account = ModelGenerationPool::for_account_endpoint(
592            "provider-a",
593            "model-a",
594            "https://example.test",
595            "account-1",
596            concurrency,
597        )
598        .unwrap();
599        assert_ne!(base.identity, provider.identity);
600        assert_ne!(base.identity, model.identity);
601        assert_ne!(base.identity, account.identity);
602    }
603
604    #[tokio::test]
605    async fn tighter_nested_gate_retains_the_session_scheduler_quota() {
606        let scheduler = Arc::new(
607            crate::task_scheduler::TaskScheduler::new(crate::task_scheduler::TaskSchedulerConfig {
608                max_active: 1,
609                aging_interval_ms: 1_000,
610            })
611            .unwrap(),
612        );
613        let quota =
614            crate::task_scheduler::TaskSchedulerQuota::for_scope("provider-session", 1).unwrap();
615        let session = ModelGenerationAdmission::new(ModelGenerationConcurrency::single_flight())
616            .with_scheduler_quota(
617                Arc::clone(&scheduler),
618                quota,
619                crate::task_scheduler::TaskPriority::Foreground,
620                "session-generation",
621            )
622            .unwrap();
623        let nested = ModelGenerationAdmission::new(ModelGenerationConcurrency::bounded(
624            NonZeroUsize::new(2).unwrap(),
625        ))
626        .with_scheduler_quota_from(&session, "nested-generation")
627        .unwrap();
628
629        let holder = session.acquire(&CancellationToken::new()).await.unwrap();
630        let cancellation = CancellationToken::new();
631        let waiting = tokio::spawn({
632            let nested = nested.clone();
633            let cancellation = cancellation.clone();
634            async move { nested.acquire(&cancellation).await }
635        });
636        tokio::task::yield_now().await;
637        assert!(!waiting.is_finished());
638        drop(holder);
639        let permit = tokio::time::timeout(Duration::from_millis(100), waiting)
640            .await
641            .unwrap()
642            .unwrap()
643            .unwrap();
644        drop(permit);
645        scheduler.shutdown().await;
646    }
647
648    #[test]
649    fn pool_rejects_credentials_and_unbounded_components() {
650        let concurrency = ModelGenerationConcurrency::single_flight();
651        assert!(matches!(
652            ModelGenerationPool::for_endpoint(
653                "provider",
654                "model",
655                "https://user:secret@example.test/v1",
656                concurrency,
657            ),
658            Err(ModelGenerationPoolError::EndpointCredentials)
659        ));
660        assert!(matches!(
661            ModelGenerationPool::for_endpoint(
662                &"p".repeat(MODEL_GENERATION_POOL_MAX_COMPONENT_BYTES + 1),
663                "model",
664                "https://example.test",
665                concurrency,
666            ),
667            Err(ModelGenerationPoolError::InvalidComponent {
668                field: "provider",
669                ..
670            })
671        ));
672    }
673
674    #[tokio::test]
675    async fn scheduler_backed_gates_share_provider_capacity_across_sessions() {
676        let scheduler = Arc::new(
677            crate::task_scheduler::TaskScheduler::new(crate::task_scheduler::TaskSchedulerConfig {
678                max_active: 1,
679                aging_interval_ms: 60_000,
680            })
681            .unwrap(),
682        );
683        let pool = ModelGenerationPool::for_endpoint(
684            "provider",
685            "model",
686            "https://example.test",
687            ModelGenerationConcurrency::single_flight(),
688        )
689        .unwrap();
690        let quota = crate::task_scheduler::TaskSchedulerQuota::new(
691            pool.identity.clone(),
692            pool.max_concurrency().get(),
693        )
694        .unwrap();
695        let admission_a =
696            ModelGenerationAdmission::new(ModelGenerationConcurrency::single_flight())
697                .with_scheduler_quota(
698                    Arc::clone(&scheduler),
699                    quota.clone(),
700                    crate::task_scheduler::TaskPriority::Foreground,
701                    "model-generation:session-a",
702                )
703                .unwrap();
704        let admission_b =
705            ModelGenerationAdmission::new(ModelGenerationConcurrency::single_flight())
706                .with_scheduler_quota(
707                    Arc::clone(&scheduler),
708                    quota.clone(),
709                    crate::task_scheduler::TaskPriority::Foreground,
710                    "model-generation:session-b",
711                )
712                .unwrap();
713        let first = admission_a
714            .acquire(&CancellationToken::new())
715            .await
716            .unwrap();
717        let cancellation = CancellationToken::new();
718        let waiting = tokio::spawn({
719            let admission = admission_b.clone();
720            let cancellation = cancellation.clone();
721            async move { admission.acquire(&cancellation).await }
722        });
723        for _ in 0..100 {
724            if scheduler.quota_snapshot(&quota).await.unwrap().pending == 1 {
725                break;
726            }
727            tokio::task::yield_now().await;
728        }
729        assert_eq!(scheduler.quota_snapshot(&quota).await.unwrap().active, 1);
730        assert_eq!(scheduler.quota_snapshot(&quota).await.unwrap().pending, 1);
731        cancellation.cancel();
732        assert!(matches!(
733            waiting.await.unwrap(),
734            Err(ModelGenerationAdmissionError::Cancelled)
735        ));
736        drop(first);
737        let replacement = tokio::time::timeout(
738            Duration::from_millis(100),
739            admission_b.acquire(&CancellationToken::new()),
740        )
741        .await
742        .expect("provider quota should be released")
743        .unwrap();
744        drop(replacement);
745        scheduler.shutdown().await;
746    }
747
748    #[tokio::test]
749    async fn pool_health_composes_local_and_scheduler_capacity_without_routing_text() {
750        let scheduler = Arc::new(
751            crate::task_scheduler::TaskScheduler::new(crate::task_scheduler::TaskSchedulerConfig {
752                max_active: 1,
753                aging_interval_ms: 60_000,
754            })
755            .unwrap(),
756        );
757        let pool = ModelGenerationPool::for_client(
758            "secret-provider-name",
759            "secret-model-name",
760            Some("https://provider.test/v1/chat?api_key=secret"),
761            Some("private-account"),
762            ModelGenerationConcurrency::single_flight(),
763        )
764        .unwrap();
765        let admission = ModelGenerationAdmission::new(ModelGenerationConcurrency::single_flight())
766            .with_model_generation_pool(
767                Arc::clone(&scheduler),
768                pool.clone(),
769                crate::task_scheduler::TaskPriority::Foreground,
770                "secret-session-label",
771            )
772            .unwrap();
773
774        let configured = admission.pool_health().await.unwrap().unwrap();
775        assert_eq!(configured.pool, pool);
776        assert_eq!(configured.local_max_concurrency, 1);
777        assert_eq!(configured.local_reserved, 0);
778        assert_eq!(configured.local_available, 1);
779        assert!(!configured.scheduler.as_ref().unwrap().observed);
780
781        let permit = admission.acquire(&CancellationToken::new()).await.unwrap();
782        let active = admission.pool_health().await.unwrap().unwrap();
783        assert_eq!(active.local_reserved, 1);
784        assert_eq!(active.local_available, 0);
785        assert!(active.scheduler.as_ref().unwrap().live);
786        assert_eq!(active.scheduler.as_ref().unwrap().active, 1);
787
788        drop(permit);
789        let retained = admission.pool_health().await.unwrap().unwrap();
790        assert_eq!(retained.local_reserved, 0);
791        assert_eq!(retained.scheduler.as_ref().unwrap().released, 1);
792        assert!(!retained.scheduler.as_ref().unwrap().live);
793        let encoded = serde_json::to_string(&retained).unwrap();
794        for secret in [
795            "secret-provider-name",
796            "secret-model-name",
797            "provider.test",
798            "api_key",
799            "private-account",
800            "secret-session-label",
801        ] {
802            assert!(!encoded.contains(secret), "snapshot leaked {secret}");
803        }
804        scheduler.shutdown().await;
805    }
806
807    #[tokio::test]
808    async fn nested_runtime_rebind_retains_the_exact_provider_pool_health() {
809        let scheduler = Arc::new(
810            crate::task_scheduler::TaskScheduler::new(crate::task_scheduler::TaskSchedulerConfig {
811                max_active: 1,
812                aging_interval_ms: 60_000,
813            })
814            .unwrap(),
815        );
816        let pool = ModelGenerationPool::for_endpoint(
817            "provider",
818            "model",
819            "https://provider.test",
820            ModelGenerationConcurrency::bounded(NonZeroUsize::new(2).unwrap()),
821        )
822        .unwrap();
823        let session = ModelGenerationAdmission::new(ModelGenerationConcurrency::bounded(
824            NonZeroUsize::new(2).unwrap(),
825        ))
826        .with_model_generation_pool(
827            Arc::clone(&scheduler),
828            pool.clone(),
829            crate::task_scheduler::TaskPriority::Foreground,
830            "session-generation",
831        )
832        .unwrap();
833        let nested = ModelGenerationAdmission::new(ModelGenerationConcurrency::single_flight())
834            .with_scheduler_quota_from(&session, "nested-generation")
835            .unwrap();
836
837        let health = nested.pool_health().await.unwrap().unwrap();
838        assert_eq!(health.pool, pool);
839        assert_eq!(health.local_max_concurrency, 1);
840        assert_eq!(health.scheduler.as_ref().unwrap().max_active, 2);
841        let permit = nested.acquire(&CancellationToken::new()).await.unwrap();
842        assert_eq!(
843            session
844                .pool_health()
845                .await
846                .unwrap()
847                .unwrap()
848                .scheduler
849                .unwrap()
850                .active,
851            1
852        );
853        drop(permit);
854        scheduler.shutdown().await;
855    }
856
857    #[tokio::test]
858    async fn scheduler_quota_without_typed_pool_is_not_product_pool_health() {
859        let scheduler = Arc::new(
860            crate::task_scheduler::TaskScheduler::new(crate::task_scheduler::TaskSchedulerConfig {
861                max_active: 1,
862                aging_interval_ms: 60_000,
863            })
864            .unwrap(),
865        );
866        let quota = crate::task_scheduler::TaskSchedulerQuota::for_scope("compat-only", 1).unwrap();
867        let admission = ModelGenerationAdmission::new(ModelGenerationConcurrency::single_flight())
868            .with_scheduler_quota(
869                Arc::clone(&scheduler),
870                quota,
871                crate::task_scheduler::TaskPriority::Foreground,
872                "compat-generation",
873            )
874            .unwrap();
875        assert!(admission.has_scheduler_quota());
876        assert!(!admission.publishes_model_generation_pool());
877        assert!(admission.pool_health().await.unwrap().is_none());
878        scheduler.shutdown().await;
879    }
880
881    #[tokio::test]
882    async fn typed_pool_publishes_product_pool_health() {
883        let scheduler = Arc::new(
884            crate::task_scheduler::TaskScheduler::new(crate::task_scheduler::TaskSchedulerConfig {
885                max_active: 1,
886                aging_interval_ms: 60_000,
887            })
888            .unwrap(),
889        );
890        let pool = ModelGenerationPool::for_client(
891            "provider",
892            "model",
893            Some("https://example.test/v1"),
894            None,
895            ModelGenerationConcurrency::single_flight(),
896        )
897        .unwrap();
898        let admission = ModelGenerationAdmission::new(ModelGenerationConcurrency::single_flight())
899            .with_model_generation_pool(
900                Arc::clone(&scheduler),
901                pool,
902                crate::task_scheduler::TaskPriority::Foreground,
903                "product-generation",
904            )
905            .unwrap();
906        assert!(admission.publishes_model_generation_pool());
907        assert!(admission.pool_health().await.unwrap().is_some());
908        scheduler.shutdown().await;
909    }
910}
a3s_code_core/llm/admission.rs

a3s_code_core/llm/
admission.rs
