agentos_runtime/
capability.rs

1//! One generation-aware registry for native and kernel-backed capabilities.
2
3use std::collections::BTreeMap;
4use std::fmt;
5use std::sync::{Arc, Mutex, Weak};
6
7use crate::accounting::{LimitError, Reservation, ResourceClass, ResourceLedger};
8
9pub type CapabilityId = u64;
10/// Capability IDs are never recycled within one VM/session generation, so the
11/// per-capability generation remains one. The session generation is part of
12/// every validated identity and distinguishes separate VM lifetimes.
13pub type CapabilityGeneration = u64;
14pub type SessionGeneration = u64;
15
16const NON_RECYCLING_CAPABILITY_GENERATION: CapabilityGeneration = 1;
17
18#[derive(Clone, Copy, Debug, Eq, PartialEq)]
19pub enum CapabilityKind {
20    TcpSocket,
21    TcpListener,
22    UnixSocket,
23    UnixListener,
24    UdpSocket,
25    TlsTransport,
26    Http2Connection,
27    Http2Stream,
28}
29
30impl CapabilityKind {
31    pub const fn is_socket(self) -> bool {
32        matches!(
33            self,
34            Self::TcpSocket
35                | Self::TcpListener
36                | Self::UnixSocket
37                | Self::UnixListener
38                | Self::UdpSocket
39                | Self::Http2Connection
40        )
41    }
42
43    pub const fn is_connection(self) -> bool {
44        matches!(
45            self,
46            Self::TcpSocket | Self::UnixSocket | Self::Http2Connection
47        )
48    }
49}
50
51#[derive(Clone, Debug, Eq, PartialEq)]
52pub enum CapabilityBackend {
53    Native { local_id: String },
54    Kernel { socket_id: u64 },
55}
56
57#[derive(Clone, Copy, Debug, Eq, PartialEq)]
58pub enum CapabilityLifecycle {
59    Allocating,
60    Open,
61    Closing,
62    Failed,
63    Closed,
64}
65
66#[derive(Clone, Debug, Eq, PartialEq)]
67pub struct CapabilitySnapshot {
68    pub session_generation: SessionGeneration,
69    pub id: CapabilityId,
70    pub generation: CapabilityGeneration,
71    pub kind: CapabilityKind,
72    pub backend: CapabilityBackend,
73    pub lifecycle: CapabilityLifecycle,
74    pub referenced: bool,
75}
76
77#[derive(Debug)]
78pub enum CapabilityError {
79    Limit(LimitError),
80    IdExhausted,
81    RegistryClosed,
82    Stale {
83        id: CapabilityId,
84        supplied_generation: CapabilityGeneration,
85    },
86    WrongSession {
87        id: CapabilityId,
88        expected: SessionGeneration,
89        actual: SessionGeneration,
90    },
91    WrongKind {
92        id: CapabilityId,
93        expected: CapabilityKind,
94        actual: CapabilityKind,
95    },
96    InvalidTransition {
97        id: CapabilityId,
98        from: CapabilityLifecycle,
99        to: CapabilityLifecycle,
100    },
101    Poisoned,
102}
103
104impl fmt::Display for CapabilityError {
105    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
106        match self {
107            Self::Limit(error) => error.fmt(formatter),
108            Self::IdExhausted => formatter.write_str(
109                "ERR_AGENTOS_CAPABILITY_ID_EXHAUSTED: capability id space exhausted",
110            ),
111            Self::RegistryClosed => formatter.write_str(
112                "ERR_AGENTOS_CAPABILITY_REGISTRY_CLOSED: VM capability admission is closed",
113            ),
114            Self::Stale {
115                id,
116                supplied_generation,
117            } => write!(
118                formatter,
119                "ERR_AGENTOS_STALE_CAPABILITY: capability {id} generation {supplied_generation} is not active"
120            ),
121            Self::WrongSession {
122                id,
123                expected,
124                actual,
125            } => write!(
126                formatter,
127                "ERR_AGENTOS_CAPABILITY_SESSION: capability {id} belongs to VM generation {actual}, not {expected}"
128            ),
129            Self::WrongKind {
130                id,
131                expected,
132                actual,
133            } => write!(
134                formatter,
135                "ERR_AGENTOS_CAPABILITY_KIND: capability {id} has kind {actual:?}, expected {expected:?}"
136            ),
137            Self::InvalidTransition { id, from, to } => write!(
138                formatter,
139                "ERR_AGENTOS_CAPABILITY_LIFECYCLE: capability {id} cannot transition from {from:?} to {to:?}"
140            ),
141            Self::Poisoned => formatter.write_str(
142                "ERR_AGENTOS_CAPABILITY_REGISTRY_POISONED: capability registry lock poisoned",
143            ),
144        }
145    }
146}
147
148impl std::error::Error for CapabilityError {}
149
150impl From<LimitError> for CapabilityError {
151    fn from(value: LimitError) -> Self {
152        Self::Limit(value)
153    }
154}
155
156#[derive(Debug)]
157struct CapabilityEntry {
158    snapshot: CapabilitySnapshot,
159}
160
161#[derive(Debug)]
162struct RegistryState {
163    next_id: CapabilityId,
164    open: bool,
165    pending: usize,
166    entries: BTreeMap<CapabilityId, CapabilityEntry>,
167}
168
169#[derive(Debug)]
170struct RegistryInner {
171    session_generation: SessionGeneration,
172    ledger: Arc<ResourceLedger>,
173    state: Mutex<RegistryState>,
174    admission_changed: tokio::sync::Notify,
175    settled: tokio::sync::Notify,
176}
177
178#[derive(Clone, Debug)]
179pub struct CapabilityRegistry {
180    inner: Arc<RegistryInner>,
181}
182
183impl CapabilityRegistry {
184    pub fn new(session_generation: SessionGeneration, ledger: Arc<ResourceLedger>) -> Self {
185        Self {
186            inner: Arc::new(RegistryInner {
187                session_generation,
188                ledger,
189                state: Mutex::new(RegistryState {
190                    next_id: 1,
191                    open: true,
192                    pending: 0,
193                    entries: BTreeMap::new(),
194                }),
195                admission_changed: tokio::sync::Notify::new(),
196                settled: tokio::sync::Notify::new(),
197            }),
198        }
199    }
200
201    pub fn session_generation(&self) -> SessionGeneration {
202        self.inner.session_generation
203    }
204
205    pub fn resources(&self) -> Arc<ResourceLedger> {
206        Arc::clone(&self.inner.ledger)
207    }
208
209    /// Acquire every count permit before allocating the descriptor/backend.
210    pub fn reserve(&self, kind: CapabilityKind) -> Result<PendingCapability, CapabilityError> {
211        {
212            let state = self
213                .inner
214                .state
215                .lock()
216                .map_err(|_| CapabilityError::Poisoned)?;
217            if !state.open {
218                return Err(CapabilityError::RegistryClosed);
219            }
220        }
221        let capability = self.inner.ledger.reserve(ResourceClass::Capabilities, 1)?;
222        // One possible ready-set entry is reserved with every admitted
223        // capability. Readiness for a live handle therefore cannot fail after
224        // the backend has already become observable.
225        let ready_handle = self.inner.ledger.reserve(ResourceClass::ReadyHandles, 1)?;
226        let socket = kind
227            .is_socket()
228            .then(|| self.inner.ledger.reserve(ResourceClass::Sockets, 1))
229            .transpose()?;
230        let connection = kind
231            .is_connection()
232            .then(|| self.inner.ledger.reserve(ResourceClass::Connections, 1))
233            .transpose()?;
234        {
235            let mut state = self
236                .inner
237                .state
238                .lock()
239                .map_err(|_| CapabilityError::Poisoned)?;
240            if !state.open {
241                return Err(CapabilityError::RegistryClosed);
242            }
243            state.pending = state
244                .pending
245                .checked_add(1)
246                .ok_or(CapabilityError::IdExhausted)?;
247        }
248        Ok(PendingCapability {
249            registry: Arc::clone(&self.inner),
250            kind,
251            reservations: vec![Some(capability), Some(ready_handle), socket, connection]
252                .into_iter()
253                .flatten()
254                .collect(),
255            counted: true,
256        })
257    }
258
259    /// Pause an async source at admission instead of accepting/reading and then
260    /// queueing uncharged state. The notification is armed before each retry so
261    /// a concurrent release cannot be missed.
262    pub async fn reserve_when_available(
263        &self,
264        kind: CapabilityKind,
265    ) -> Result<PendingCapability, CapabilityError> {
266        loop {
267            let changed = self.inner.ledger.capacity_changed();
268            let admission_changed = self.inner.admission_changed.notified();
269            match self.reserve(kind) {
270                Ok(pending) => return Ok(pending),
271                Err(CapabilityError::Limit(error)) if error.requested > error.limit => {
272                    return Err(CapabilityError::Limit(error));
273                }
274                Err(CapabilityError::Limit(_)) => {
275                    tokio::select! {
276                        _ = changed => {}
277                        _ = admission_changed => {}
278                    }
279                }
280                Err(error) => return Err(error),
281            }
282        }
283    }
284
285    pub fn snapshot(&self, id: CapabilityId) -> Option<CapabilitySnapshot> {
286        self.inner
287            .state
288            .lock()
289            .ok()?
290            .entries
291            .get(&id)
292            .map(|entry| entry.snapshot.clone())
293    }
294
295    pub fn snapshots(&self) -> Vec<CapabilitySnapshot> {
296        self.inner
297            .state
298            .lock()
299            .map(|state| {
300                state
301                    .entries
302                    .values()
303                    .map(|entry| entry.snapshot.clone())
304                    .collect()
305            })
306            .unwrap_or_else(|_| {
307                eprintln!(
308                    "ERR_AGENTOS_CAPABILITY_REGISTRY_POISONED: failed to snapshot capability registry"
309                );
310                Vec::new()
311            })
312    }
313
314    pub fn close_admission(&self) -> Result<(), CapabilityError> {
315        self.inner
316            .state
317            .lock()
318            .map_err(|_| CapabilityError::Poisoned)?
319            .open = false;
320        self.inner.admission_changed.notify_waiters();
321        Ok(())
322    }
323
324    pub fn active_len(&self) -> usize {
325        self.inner
326            .state
327            .lock()
328            .map(|state| state.entries.len())
329            .unwrap_or(0)
330    }
331
332    pub fn outstanding_len(&self) -> usize {
333        self.inner
334            .state
335            .lock()
336            .map(|state| state.pending.saturating_add(state.entries.len()))
337            .unwrap_or(usize::MAX)
338    }
339
340    pub async fn wait_empty(&self) {
341        loop {
342            let settled = self.inner.settled.notified();
343            if self.outstanding_len() == 0 {
344                return;
345            }
346            settled.await;
347        }
348    }
349}
350
351#[derive(Debug)]
352pub struct PendingCapability {
353    registry: Arc<RegistryInner>,
354    kind: CapabilityKind,
355    reservations: Vec<Reservation>,
356    counted: bool,
357}
358
359impl PendingCapability {
360    /// Commit only after the backend has been allocated. Dropping before commit
361    /// rolls admission back; dropping the returned lease closes the registry row.
362    pub fn commit(
363        mut self,
364        backend: CapabilityBackend,
365    ) -> Result<CapabilityLease, CapabilityError> {
366        let (id, generation) = {
367            let mut state = self
368                .registry
369                .state
370                .lock()
371                .map_err(|_| CapabilityError::Poisoned)?;
372            if !state.open {
373                return Err(CapabilityError::RegistryClosed);
374            }
375            if state.pending == 0 {
376                eprintln!(
377                    "ERR_AGENTOS_CAPABILITY_ACCOUNTING_UNDERFLOW: pending commit without admission"
378                );
379                return Err(CapabilityError::Poisoned);
380            }
381            let id = state.next_id;
382            let next_id = id.checked_add(1).ok_or(CapabilityError::IdExhausted)?;
383            state.pending -= 1;
384            self.counted = false;
385            state.next_id = next_id;
386            let generation = NON_RECYCLING_CAPABILITY_GENERATION;
387            let replaced = state.entries.insert(
388                id,
389                CapabilityEntry {
390                    snapshot: CapabilitySnapshot {
391                        session_generation: self.registry.session_generation,
392                        id,
393                        generation,
394                        kind: self.kind,
395                        backend,
396                        lifecycle: CapabilityLifecycle::Open,
397                        referenced: true,
398                    },
399                },
400            );
401            if let Some(previous) = replaced {
402                eprintln!(
403                    "ERR_AGENTOS_CAPABILITY_ID_REUSED: monotonic capability id {id} was already active"
404                );
405                state.entries.insert(id, previous);
406                return Err(CapabilityError::Poisoned);
407            }
408            (id, generation)
409        };
410        Ok(CapabilityLease {
411            registry: Arc::downgrade(&self.registry),
412            id,
413            generation,
414            reservations: std::mem::take(&mut self.reservations),
415        })
416    }
417}
418
419impl Drop for PendingCapability {
420    fn drop(&mut self) {
421        if self.counted {
422            match self.registry.state.lock() {
423                Ok(mut state) => {
424                    if state.pending == 0 {
425                        eprintln!(
426                            "ERR_AGENTOS_CAPABILITY_ACCOUNTING_UNDERFLOW: pending reservation released at zero"
427                        );
428                    } else {
429                        state.pending -= 1;
430                    }
431                }
432                Err(_) => eprintln!(
433                    "ERR_AGENTOS_CAPABILITY_REGISTRY_POISONED: pending reservation release failed"
434                ),
435            }
436            self.counted = false;
437            self.registry.settled.notify_waiters();
438        }
439    }
440}
441
442#[derive(Debug)]
443pub struct CapabilityLease {
444    registry: Weak<RegistryInner>,
445    id: CapabilityId,
446    generation: CapabilityGeneration,
447    reservations: Vec<Reservation>,
448}
449
450impl CapabilityLease {
451    pub fn id(&self) -> CapabilityId {
452        self.id
453    }
454
455    pub fn generation(&self) -> CapabilityGeneration {
456        self.generation
457    }
458
459    /// Validate a guest-visible alias against the live registry row before an
460    /// operation reaches its backend.
461    pub fn validate(
462        &self,
463        session_generation: SessionGeneration,
464        kind: CapabilityKind,
465    ) -> Result<(), CapabilityError> {
466        let registry = self
467            .registry
468            .upgrade()
469            .ok_or(CapabilityError::RegistryClosed)?;
470        let state = registry
471            .state
472            .lock()
473            .map_err(|_| CapabilityError::Poisoned)?;
474        let entry = state.entries.get(&self.id).ok_or(CapabilityError::Stale {
475            id: self.id,
476            supplied_generation: self.generation,
477        })?;
478        if entry.snapshot.generation != self.generation {
479            return Err(CapabilityError::Stale {
480                id: self.id,
481                supplied_generation: self.generation,
482            });
483        }
484        if entry.snapshot.session_generation != session_generation {
485            return Err(CapabilityError::WrongSession {
486                id: self.id,
487                expected: session_generation,
488                actual: entry.snapshot.session_generation,
489            });
490        }
491        if entry.snapshot.kind != kind {
492            return Err(CapabilityError::WrongKind {
493                id: self.id,
494                expected: kind,
495                actual: entry.snapshot.kind,
496            });
497        }
498        Ok(())
499    }
500
501    pub fn set_referenced(&self, referenced: bool) -> Result<(), CapabilityError> {
502        self.update(|snapshot| snapshot.referenced = referenced)
503    }
504
505    pub fn transition(&self, to: CapabilityLifecycle) -> Result<(), CapabilityError> {
506        let registry = self
507            .registry
508            .upgrade()
509            .ok_or(CapabilityError::RegistryClosed)?;
510        let mut state = registry
511            .state
512            .lock()
513            .map_err(|_| CapabilityError::Poisoned)?;
514        let entry = state
515            .entries
516            .get_mut(&self.id)
517            .ok_or(CapabilityError::Stale {
518                id: self.id,
519                supplied_generation: self.generation,
520            })?;
521        if entry.snapshot.generation != self.generation {
522            return Err(CapabilityError::Stale {
523                id: self.id,
524                supplied_generation: self.generation,
525            });
526        }
527        let from = entry.snapshot.lifecycle;
528        if !valid_transition(from, to) {
529            return Err(CapabilityError::InvalidTransition {
530                id: self.id,
531                from,
532                to,
533            });
534        }
535        entry.snapshot.lifecycle = to;
536        Ok(())
537    }
538
539    fn update(&self, update: impl FnOnce(&mut CapabilitySnapshot)) -> Result<(), CapabilityError> {
540        let registry = self
541            .registry
542            .upgrade()
543            .ok_or(CapabilityError::RegistryClosed)?;
544        let mut state = registry
545            .state
546            .lock()
547            .map_err(|_| CapabilityError::Poisoned)?;
548        let entry = state
549            .entries
550            .get_mut(&self.id)
551            .ok_or(CapabilityError::Stale {
552                id: self.id,
553                supplied_generation: self.generation,
554            })?;
555        if entry.snapshot.generation != self.generation {
556            return Err(CapabilityError::Stale {
557                id: self.id,
558                supplied_generation: self.generation,
559            });
560        }
561        update(&mut entry.snapshot);
562        Ok(())
563    }
564}
565
566fn valid_transition(from: CapabilityLifecycle, to: CapabilityLifecycle) -> bool {
567    from == to
568        || matches!(
569            (from, to),
570            (CapabilityLifecycle::Open, CapabilityLifecycle::Closing)
571                | (CapabilityLifecycle::Open, CapabilityLifecycle::Failed)
572                | (CapabilityLifecycle::Open, CapabilityLifecycle::Closed)
573                | (CapabilityLifecycle::Closing, CapabilityLifecycle::Closed)
574                | (CapabilityLifecycle::Failed, CapabilityLifecycle::Closed)
575        )
576}
577
578impl Drop for CapabilityLease {
579    fn drop(&mut self) {
580        if let Some(registry) = self.registry.upgrade() {
581            match registry.state.lock() {
582                Ok(mut state) => {
583                    let remove = state
584                        .entries
585                        .get(&self.id)
586                        .is_some_and(|entry| entry.snapshot.generation == self.generation);
587                    if remove {
588                        state.entries.remove(&self.id);
589                    } else {
590                        eprintln!(
591                            "ERR_AGENTOS_CAPABILITY_RELEASE_STALE: capability={} generation={}",
592                            self.id, self.generation
593                        );
594                    }
595                }
596                Err(_) => eprintln!(
597                    "ERR_AGENTOS_CAPABILITY_REGISTRY_POISONED: capability={} release failed",
598                    self.id
599                ),
600            }
601            registry.settled.notify_waiters();
602        }
603        self.reservations.clear();
604    }
605}
606
607#[cfg(test)]
608mod tests {
609    use super::*;
610    use crate::accounting::ResourceLimit;
611
612    fn registry(maximum: usize) -> (Arc<ResourceLedger>, CapabilityRegistry) {
613        let ledger = Arc::new(ResourceLedger::root(
614            "vm-1",
615            [
616                (
617                    ResourceClass::Capabilities,
618                    ResourceLimit::new(maximum, "runtime.capabilities.maxPerVm"),
619                ),
620                (
621                    ResourceClass::Sockets,
622                    ResourceLimit::new(maximum, "runtime.resources.maxSockets"),
623                ),
624                (
625                    ResourceClass::Connections,
626                    ResourceLimit::new(maximum, "runtime.resources.maxConnections"),
627                ),
628            ],
629        ));
630        let registry = CapabilityRegistry::new(7, Arc::clone(&ledger));
631        (ledger, registry)
632    }
633
634    #[test]
635    fn admission_precedes_commit_and_drop_reconciles() {
636        let (ledger, registry) = registry(1);
637        let pending = registry.reserve(CapabilityKind::TcpSocket).unwrap();
638        assert_eq!(ledger.usage(ResourceClass::Sockets).used, 1);
639        assert!(registry.reserve(CapabilityKind::TcpSocket).is_err());
640        let lease = pending
641            .commit(CapabilityBackend::Native {
642                local_id: String::from("socket-1"),
643            })
644            .unwrap();
645        assert_eq!(registry.active_len(), 1);
646        assert_eq!(registry.snapshot(lease.id()).unwrap().session_generation, 7);
647        drop(lease);
648        assert_eq!(registry.active_len(), 0);
649        assert!(ledger.is_zero());
650    }
651
652    #[test]
653    fn failed_backend_allocation_rolls_back_pending_reservation() {
654        let (ledger, registry) = registry(1);
655        drop(registry.reserve(CapabilityKind::UdpSocket).unwrap());
656        assert!(ledger.is_zero());
657        assert_eq!(registry.active_len(), 0);
658    }
659
660    #[test]
661    fn alias_validation_rejects_wrong_vm_generation_and_kind() {
662        let (_ledger, registry) = registry(1);
663        let lease = registry
664            .reserve(CapabilityKind::TcpSocket)
665            .expect("reserve TCP capability")
666            .commit(CapabilityBackend::Native {
667                local_id: String::from("socket-1"),
668            })
669            .expect("commit TCP capability");
670
671        lease
672            .validate(7, CapabilityKind::TcpSocket)
673            .expect("matching generation and kind");
674        assert!(matches!(
675            lease.validate(8, CapabilityKind::TcpSocket),
676            Err(CapabilityError::WrongSession {
677                expected: 8,
678                actual: 7,
679                ..
680            })
681        ));
682        assert!(matches!(
683            lease.validate(7, CapabilityKind::UdpSocket),
684            Err(CapabilityError::WrongKind {
685                expected: CapabilityKind::UdpSocket,
686                actual: CapabilityKind::TcpSocket,
687                ..
688            })
689        ));
690    }
691
692    #[test]
693    fn close_admission_rejects_pending_commit_without_leaking() {
694        let (ledger, registry) = registry(1);
695        let pending = registry.reserve(CapabilityKind::TcpListener).unwrap();
696        registry.close_admission().unwrap();
697        assert!(pending
698            .commit(CapabilityBackend::Kernel { socket_id: 4 })
699            .is_err());
700        assert!(ledger.is_zero());
701    }
702
703    #[tokio::test]
704    async fn close_admission_wakes_capacity_waiters_with_typed_error() {
705        let (ledger, registry) = registry(1);
706        let held = registry
707            .reserve(CapabilityKind::TcpSocket)
708            .expect("fill registry");
709        let waiting_registry = registry.clone();
710        let waiter = tokio::spawn(async move {
711            waiting_registry
712                .reserve_when_available(CapabilityKind::TcpSocket)
713                .await
714        });
715        tokio::task::yield_now().await;
716        assert!(!waiter.is_finished());
717        registry.close_admission().expect("close admission");
718        let error = tokio::time::timeout(std::time::Duration::from_secs(1), waiter)
719            .await
720            .expect("close must wake waiter")
721            .expect("waiter task")
722            .expect_err("closed registry must reject admission");
723        assert!(matches!(error, CapabilityError::RegistryClosed));
724        drop(held);
725        assert!(ledger.is_zero());
726    }
727
728    #[tokio::test]
729    async fn non_retryable_zero_limit_returns_without_waiting_or_leaking() {
730        let ledger = Arc::new(ResourceLedger::root(
731            "vm-zero-sockets",
732            [
733                (
734                    ResourceClass::Capabilities,
735                    ResourceLimit::new(1, "runtime.capabilities.maxPerVm"),
736                ),
737                (
738                    ResourceClass::Sockets,
739                    ResourceLimit::new(0, "runtime.resources.maxSockets"),
740                ),
741                (
742                    ResourceClass::Connections,
743                    ResourceLimit::new(1, "runtime.resources.maxConnections"),
744                ),
745            ],
746        ));
747        let registry = CapabilityRegistry::new(7, Arc::clone(&ledger));
748
749        let error = tokio::time::timeout(
750            std::time::Duration::from_millis(100),
751            registry.reserve_when_available(CapabilityKind::TcpSocket),
752        )
753        .await
754        .expect("an impossible admission must not wait for capacity")
755        .expect_err("zero socket limit must reject admission");
756        let CapabilityError::Limit(error) = error else {
757            panic!("expected typed resource limit error");
758        };
759        assert_eq!(error.resource, ResourceClass::Sockets);
760        assert_eq!(error.requested, 1);
761        assert_eq!(error.limit, 0);
762        assert_eq!(error.config_path, "runtime.resources.maxSockets");
763        assert!(ledger.is_zero());
764        assert_eq!(registry.outstanding_len(), 0);
765    }
766
767    #[test]
768    fn capability_ids_are_monotonic_and_never_recycled() {
769        let (_ledger, registry) = registry(1);
770        let first = registry
771            .reserve(CapabilityKind::UdpSocket)
772            .expect("first reservation")
773            .commit(CapabilityBackend::Native {
774                local_id: String::from("udp-1"),
775            })
776            .expect("first capability");
777        assert_eq!(first.id(), 1);
778        assert_eq!(first.generation(), NON_RECYCLING_CAPABILITY_GENERATION);
779        drop(first);
780
781        let second = registry
782            .reserve(CapabilityKind::UdpSocket)
783            .expect("second reservation")
784            .commit(CapabilityBackend::Native {
785                local_id: String::from("udp-2"),
786            })
787            .expect("second capability");
788        assert_eq!(second.id(), 2);
789        assert_eq!(second.generation(), NON_RECYCLING_CAPABILITY_GENERATION);
790        assert!(registry.snapshot(1).is_none());
791    }
792
793    #[tokio::test]
794    async fn wait_empty_includes_pending_and_committed_ownership() {
795        let (ledger, registry) = registry(2);
796        let pending = registry
797            .reserve(CapabilityKind::UdpSocket)
798            .expect("pending capability");
799        let lease = registry
800            .reserve(CapabilityKind::TcpListener)
801            .expect("committed capability")
802            .commit(CapabilityBackend::Kernel { socket_id: 8 })
803            .expect("commit");
804        assert_eq!(registry.outstanding_len(), 2);
805        let waiter = tokio::spawn({
806            let registry = registry.clone();
807            async move { registry.wait_empty().await }
808        });
809        tokio::task::yield_now().await;
810        assert!(!waiter.is_finished());
811        drop(pending);
812        tokio::task::yield_now().await;
813        assert!(!waiter.is_finished());
814        drop(lease);
815        tokio::time::timeout(std::time::Duration::from_secs(1), waiter)
816            .await
817            .expect("empty notification")
818            .expect("wait task");
819        assert!(ledger.is_zero());
820    }
821
822    #[tokio::test]
823    async fn id_exhaustion_reconciles_pending_waiters() {
824        let (_, registry) = registry(1);
825        let pending = registry
826            .reserve(CapabilityKind::TcpSocket)
827            .expect("pending capability");
828        registry.inner.state.lock().expect("registry state").next_id = u64::MAX;
829        let waiter = tokio::spawn({
830            let registry = registry.clone();
831            async move { registry.wait_empty().await }
832        });
833        tokio::task::yield_now().await;
834
835        let error = pending
836            .commit(CapabilityBackend::Native {
837                local_id: String::from("exhausted"),
838            })
839            .expect_err("exhausted id must reject commit");
840        assert!(matches!(error, CapabilityError::IdExhausted));
841        tokio::time::timeout(std::time::Duration::from_secs(1), waiter)
842            .await
843            .expect("pending release must wake waiter")
844            .expect("wait task");
845        assert_eq!(registry.outstanding_len(), 0);
846    }
847}
agentos_runtime/capability.rs

agentos_runtime/
capability.rs
