Skip to main content

theway_core/agent/runtime_extensions/
scope.rs

1use std::sync::Arc;
2use std::sync::atomic::{AtomicU64, Ordering};
3
4use thiserror::Error;
5
6#[derive(Clone, Copy, Debug, PartialEq, Eq)]
7pub enum RuntimeExtensionScopeKind {
8    Run,
9    Turn,
10    Request,
11    Message,
12    ToolCall,
13}
14
15impl RuntimeExtensionScopeKind {
16    const fn as_str(self) -> &'static str {
17        match self {
18            Self::Run => "run",
19            Self::Turn => "turn",
20            Self::Request => "request",
21            Self::Message => "message",
22            Self::ToolCall => "tool-call",
23        }
24    }
25}
26
27#[derive(Clone, Debug)]
28pub struct RuntimeExtensionScopeAllocator {
29    session_id: Arc<str>,
30    next_scope: Arc<AtomicU64>,
31    next_sequence: Arc<AtomicU64>,
32}
33
34impl RuntimeExtensionScopeAllocator {
35    pub fn new(session_id: impl Into<String>) -> Result<Self, ScopeAllocationError> {
36        let session_id = session_id.into();
37        if session_id.trim().is_empty() {
38            return Err(ScopeAllocationError::EmptySessionId);
39        }
40        Ok(Self {
41            session_id: Arc::from(session_id),
42            next_scope: Arc::new(AtomicU64::new(1)),
43            next_sequence: Arc::new(AtomicU64::new(1)),
44        })
45    }
46
47    pub fn session_id(&self) -> &str {
48        &self.session_id
49    }
50
51    pub fn allocate(
52        &self,
53        kind: RuntimeExtensionScopeKind,
54    ) -> Result<String, ScopeAllocationError> {
55        let ordinal = take_next(&self.next_scope)?;
56        Ok(format!("{}:{}:{ordinal}", self.session_id, kind.as_str()))
57    }
58
59    pub fn next_sequence(&self) -> Result<u64, ScopeAllocationError> {
60        take_next(&self.next_sequence)
61    }
62}
63
64fn take_next(counter: &AtomicU64) -> Result<u64, ScopeAllocationError> {
65    counter
66        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
67            current.checked_add(1)
68        })
69        .map_err(|_| ScopeAllocationError::Exhausted)
70}
71
72#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)]
73pub enum ScopeAllocationError {
74    #[error("runtime extension session id must not be empty")]
75    EmptySessionId,
76    #[error("runtime extension scope or lifecycle sequence space is exhausted")]
77    Exhausted,
78}