Skip to main content

rig_core/
id.rs

1//! Non-cryptographic random identifiers, process-local run counters, and
2//! caller-supplied conversation keys. Do not use generated IDs as secrets.
3//!
4//! ```
5//! let id = rig_core::id::generate();
6//! assert_eq!(id.len(), 21);
7//! ```
8
9/// URL-safe ASCII alphabet for random identifiers.
10const ALPHABET: &[u8; 64] = b"_-0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
11
12/// Number of characters returned by [`generate`].
13const DEFAULT_LEN: usize = 21;
14
15/// Generate a 21-character URL-safe random identifier. Not cryptographically
16/// secure; collisions are possible.
17pub fn generate() -> String {
18    generate_with_len(DEFAULT_LEN)
19}
20
21/// Define a nonzero process-local counter with transparent integer serde and
22/// decimal parsing/display. Values increase until wraparound; raw construction
23/// and deserialization do not reserve counter values.
24macro_rules! counter_id {
25    ($(#[$meta:meta])* $name:ident) => {
26        $(#[$meta])*
27        #[derive(
28            Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
29        )]
30        #[serde(transparent)]
31        pub struct $name(std::num::NonZeroU64);
32
33        impl $name {
34            fn counter() -> &'static std::sync::atomic::AtomicU64 {
35                static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
36                &NEXT
37            }
38
39            /// Mint the next nonzero process-local counter value. Values may repeat
40            /// after wraparound or collide with raw or deserialized IDs.
41            #[allow(clippy::new_without_default)]
42            pub fn new() -> Self {
43                // Zero is reserved even after counter wraparound.
44                loop {
45                    if let Some(raw) = std::num::NonZeroU64::new(
46                        Self::counter().fetch_add(1, std::sync::atomic::Ordering::Relaxed),
47                    ) {
48                        return Self(raw);
49                    }
50                }
51            }
52
53            /// Build an id from its raw value. `None` for zero, which is never
54            /// a valid id.
55            pub const fn from_raw(raw: u64) -> Option<Self> {
56                match std::num::NonZeroU64::new(raw) {
57                    Some(raw) => Some(Self(raw)),
58                    None => None,
59                }
60            }
61
62            /// The id's raw value (never zero).
63            pub const fn to_raw(self) -> u64 {
64                self.0.get()
65            }
66        }
67
68        impl From<std::num::NonZeroU64> for $name {
69            fn from(raw: std::num::NonZeroU64) -> Self {
70                Self(raw)
71            }
72        }
73
74        impl From<$name> for std::num::NonZeroU64 {
75            fn from(id: $name) -> Self {
76                id.0
77            }
78        }
79
80        impl std::fmt::Display for $name {
81            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82                std::fmt::Display::fmt(&self.0, f)
83            }
84        }
85
86        impl std::fmt::Debug for $name {
87            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88                write!(f, "{}({})", stringify!($name), self.0)
89            }
90        }
91
92        impl std::str::FromStr for $name {
93            type Err = ParseIdError;
94
95            fn from_str(s: &str) -> Result<Self, Self::Err> {
96                s.parse::<std::num::NonZeroU64>()
97                    .map(Self)
98                    .map_err(|_| ParseIdError(s.to_owned()))
99            }
100        }
101
102        const _: () = assert!(
103            std::mem::size_of::<Option<$name>>() == std::mem::size_of::<u64>()
104        );
105    };
106}
107
108/// Error from parsing a counter id: the text was not a non-zero decimal `u64`.
109#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
110#[error("invalid id: expected a non-zero integer, got {0:?}")]
111pub struct ParseIdError(String);
112
113counter_id! {
114    /// Process-local run identifier shared by that run's hook events.
115    /// Not secret; add a host or session scope for cross-process correlation.
116    RunId
117}
118
119/// Caller-supplied key scoping [`ConversationMemory`] and provider conversation
120/// state across runs. Wraps the string without validation.
121///
122/// [`ConversationMemory`]: crate::memory::ConversationMemory
123#[derive(
124    Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
125)]
126#[serde(transparent)]
127pub struct ConversationId(String);
128
129impl ConversationId {
130    /// Wrap a caller-supplied conversation key.
131    pub fn new(id: impl Into<String>) -> Self {
132        Self(id.into())
133    }
134
135    /// The id as a string slice.
136    pub fn as_str(&self) -> &str {
137        &self.0
138    }
139
140    /// Unwrap into the underlying string.
141    pub fn into_string(self) -> String {
142        self.0
143    }
144}
145
146impl std::fmt::Display for ConversationId {
147    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
148        f.write_str(&self.0)
149    }
150}
151
152impl From<String> for ConversationId {
153    fn from(id: String) -> Self {
154        Self(id)
155    }
156}
157
158impl From<&str> for ConversationId {
159    fn from(id: &str) -> Self {
160        Self(id.to_owned())
161    }
162}
163
164impl AsRef<str> for ConversationId {
165    fn as_ref(&self) -> &str {
166        &self.0
167    }
168}
169
170/// Generate a `len`-character, URL-safe, non-cryptographic identifier.
171pub(crate) fn generate_with_len(len: usize) -> String {
172    std::iter::repeat_with(|| {
173        let idx = fastrand::usize(..ALPHABET.len());
174        ALPHABET.get(idx).copied().unwrap_or(b'_') as char
175    })
176    .take(len)
177    .collect()
178}
179
180#[cfg(test)]
181mod tests;