hara_native/kernel/
session_snapshot.rs

1//! Immutable snapshot-backed sessions for production embeddings.
2
3use crate::core::{self, Value};
4use crate::kernel::{ResolvedSecrets, SecretCatalog};
5use crate::snapshot::{Digest, ResolvedSnapshot};
6use std::collections::{BTreeMap, BTreeSet};
7use std::rc::Rc;
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
10pub enum SessionMode {
11    Sealed,
12    Overlay,
13}
14
15#[derive(Clone, Debug)]
16pub struct SnapshotSessionDefinition {
17    pub id: String,
18    pub snapshot: Digest,
19    pub mode: SessionMode,
20    pub grants: BTreeSet<String>,
21}
22
23#[derive(Clone, Debug)]
24pub struct FrozenSession {
25    id: String,
26    snapshot: Rc<ResolvedSnapshot>,
27    mode: SessionMode,
28    grants: BTreeSet<String>,
29    secrets: ResolvedSecrets,
30}
31
32impl FrozenSession {
33    pub fn id(&self) -> &str {
34        &self.id
35    }
36
37    pub fn snapshot(&self) -> &Rc<ResolvedSnapshot> {
38        &self.snapshot
39    }
40
41    pub fn mode(&self) -> SessionMode {
42        self.mode
43    }
44
45    pub fn grants(&self) -> &BTreeSet<String> {
46        &self.grants
47    }
48
49    pub fn secrets(&self) -> &ResolvedSecrets {
50        &self.secrets
51    }
52
53    pub fn entrypoint(&self, name: &str) -> Option<&str> {
54        self.snapshot
55            .manifest
56            .entrypoints
57            .get(name)
58            .map(String::as_str)
59    }
60
61    pub fn require_mutable_overlay(&self) -> Result<(), String> {
62        match self.mode {
63            SessionMode::Overlay => Ok(()),
64            SessionMode::Sealed => Err("session/sealed-mutation-denied".into()),
65        }
66    }
67}
68
69#[derive(Default)]
70pub struct SnapshotRegistry {
71    values: BTreeMap<Digest, Rc<ResolvedSnapshot>>,
72}
73
74impl SnapshotRegistry {
75    pub fn insert(&mut self, snapshot: ResolvedSnapshot) -> Rc<ResolvedSnapshot> {
76        self.values
77            .entry(snapshot.digest)
78            .or_insert_with(|| Rc::new(snapshot))
79            .clone()
80    }
81
82    pub fn get(&self, digest: &Digest) -> Option<Rc<ResolvedSnapshot>> {
83        self.values.get(digest).cloned()
84    }
85
86    pub fn len(&self) -> usize {
87        self.values.len()
88    }
89
90    pub fn is_empty(&self) -> bool {
91        self.values.is_empty()
92    }
93}
94
95#[derive(Clone, Debug, PartialEq)]
96pub struct SharedStateCell {
97    revision: u64,
98    value: Value,
99}
100
101impl SharedStateCell {
102    pub fn revision(&self) -> u64 {
103        self.revision
104    }
105
106    pub fn value(&self) -> &Value {
107        &self.value
108    }
109}
110
111#[derive(Default)]
112pub struct SnapshotKernel {
113    snapshots: SnapshotRegistry,
114    sessions: BTreeMap<String, Rc<FrozenSession>>,
115    shared_state: BTreeMap<String, SharedStateCell>,
116}
117
118impl SnapshotKernel {
119    pub fn snapshots(&self) -> &SnapshotRegistry {
120        &self.snapshots
121    }
122
123    pub fn register_snapshot(&mut self, snapshot: ResolvedSnapshot) -> Rc<ResolvedSnapshot> {
124        self.snapshots.insert(snapshot)
125    }
126
127    pub fn initialize_state_from(&mut self, digest: &Digest) -> Result<(), String> {
128        let snapshot = self
129            .snapshots
130            .get(digest)
131            .ok_or_else(|| format!("snapshot/not-registered: {}", crate::snapshot::hex(digest)))?;
132        let mut state = BTreeMap::new();
133        for (name, value) in &snapshot.manifest.initial_state {
134            if !core::session_transferable(value) {
135                return Err(format!("snapshot/non-transferable-state: {name}"));
136            }
137            state.insert(
138                name.clone(),
139                SharedStateCell {
140                    revision: 0,
141                    value: value.clone(),
142                },
143            );
144        }
145        self.shared_state = state;
146        Ok(())
147    }
148
149    /// Builds all sessions privately and publishes them only after every
150    /// snapshot, capability, and secret requirement validates.
151    pub fn load_sessions(
152        &mut self,
153        definitions: &[SnapshotSessionDefinition],
154        secrets: &dyn SecretCatalog,
155    ) -> Result<(), String> {
156        let mut candidate = BTreeMap::new();
157        for definition in definitions {
158            if definition.id.is_empty() || candidate.contains_key(&definition.id) {
159                return Err(format!("session/duplicate-or-empty: {}", definition.id));
160            }
161            let snapshot = self.snapshots.get(&definition.snapshot).ok_or_else(|| {
162                format!(
163                    "session/snapshot-not-registered: {}",
164                    crate::snapshot::hex(&definition.snapshot)
165                )
166            })?;
167            for capability in &snapshot.manifest.capabilities {
168                if !definition.grants.contains(capability) {
169                    return Err(format!(
170                        "session/capability-not-granted: {} requires {capability}",
171                        definition.id
172                    ));
173                }
174            }
175            let resolved_secrets = ResolvedSecrets::resolve(&snapshot.manifest.secrets, secrets)?;
176            candidate.insert(
177                definition.id.clone(),
178                Rc::new(FrozenSession {
179                    id: definition.id.clone(),
180                    snapshot,
181                    mode: definition.mode,
182                    grants: definition.grants.clone(),
183                    secrets: resolved_secrets,
184                }),
185            );
186        }
187        self.sessions = candidate;
188        Ok(())
189    }
190
191    pub fn session(&self, id: &str) -> Option<Rc<FrozenSession>> {
192        self.sessions.get(id).cloned()
193    }
194
195    pub fn session_ids(&self) -> impl Iterator<Item = &str> {
196        self.sessions.keys().map(String::as_str)
197    }
198
199    pub fn state(&self, name: &str) -> Option<&SharedStateCell> {
200        self.shared_state.get(name)
201    }
202
203    pub fn state_compare_and_set(
204        &mut self,
205        name: &str,
206        expected_revision: u64,
207        value: Value,
208    ) -> Result<bool, String> {
209        if !core::session_transferable(&value) {
210            return Err(format!("kernel/non-transferable-state: {name}"));
211        }
212        let cell = self
213            .shared_state
214            .get_mut(name)
215            .ok_or_else(|| format!("kernel/unknown-state: {name}"))?;
216        if cell.revision != expected_revision {
217            return Ok(false);
218        }
219        cell.revision = cell
220            .revision
221            .checked_add(1)
222            .ok_or_else(|| format!("kernel/state-revision-exhausted: {name}"))?;
223        cell.value = value;
224        Ok(true)
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use crate::kernel::SecretDescriptor;
232    use crate::snapshot::{SecretRequirement, SnapshotArtifact, SnapshotManifest};
233
234    struct Catalog(BTreeMap<String, SecretDescriptor>);
235
236    impl SecretCatalog for Catalog {
237        fn describe(&self, id: &str) -> Result<Option<SecretDescriptor>, String> {
238            Ok(self.0.get(id).cloned())
239        }
240    }
241
242    fn resolved() -> ResolvedSnapshot {
243        SnapshotArtifact {
244            base: None,
245            manifest: SnapshotManifest {
246                language_version: "0.1".into(),
247                dependency_lock_digest: [0; 32],
248                libraries: vec![],
249                namespaces: vec![],
250                entrypoints: BTreeMap::from([("api".into(), "app/handle".into())]),
251                initial_state: BTreeMap::from([("counter".into(), Value::Number(0))]),
252                capabilities: BTreeSet::from(["nginx/timer".into()]),
253                secrets: vec![SecretRequirement {
254                    id: "key".into(),
255                    purpose: "sign".into(),
256                    required: true,
257                    version: Some("1".into()),
258                }],
259                accelerators: vec![],
260            },
261        }
262        .resolve(None)
263        .unwrap()
264    }
265
266    fn catalog() -> Catalog {
267        Catalog(BTreeMap::from([(
268            "key".into(),
269            SecretDescriptor {
270                id: "key".into(),
271                provider: "test".into(),
272                version: Some("1".into()),
273            },
274        )]))
275    }
276
277    #[test]
278    fn publishes_all_sessions_transactionally_and_shares_snapshot_memory() {
279        let mut kernel = SnapshotKernel::default();
280        let snapshot = kernel.register_snapshot(resolved());
281        kernel.initialize_state_from(&snapshot.digest).unwrap();
282        let definitions = ["primary", "sample"].map(|id| SnapshotSessionDefinition {
283            id: id.into(),
284            snapshot: snapshot.digest,
285            mode: SessionMode::Sealed,
286            grants: BTreeSet::from(["nginx/timer".into()]),
287        });
288        kernel.load_sessions(&definitions, &catalog()).unwrap();
289        let primary = kernel.session("primary").unwrap();
290        let sample = kernel.session("sample").unwrap();
291        assert!(Rc::ptr_eq(primary.snapshot(), sample.snapshot()));
292        assert_eq!(primary.entrypoint("api"), Some("app/handle"));
293        assert!(primary.require_mutable_overlay().is_err());
294        assert_eq!(kernel.state("counter").unwrap().revision(), 0);
295        assert!(kernel
296            .state_compare_and_set("counter", 0, Value::Number(1))
297            .unwrap());
298        assert_eq!(kernel.state("counter").unwrap().revision(), 1);
299    }
300
301    #[test]
302    fn failed_candidate_does_not_replace_published_sessions() {
303        let mut kernel = SnapshotKernel::default();
304        let snapshot = kernel.register_snapshot(resolved());
305        let valid = SnapshotSessionDefinition {
306            id: "primary".into(),
307            snapshot: snapshot.digest,
308            mode: SessionMode::Sealed,
309            grants: BTreeSet::from(["nginx/timer".into()]),
310        };
311        kernel.load_sessions(&[valid], &catalog()).unwrap();
312        let invalid = SnapshotSessionDefinition {
313            id: "candidate".into(),
314            snapshot: snapshot.digest,
315            mode: SessionMode::Sealed,
316            grants: BTreeSet::new(),
317        };
318        assert!(kernel.load_sessions(&[invalid], &catalog()).is_err());
319        assert!(kernel.session("primary").is_some());
320        assert!(kernel.session("candidate").is_none());
321    }
322}
hara_native/kernel/session_snapshot.rs

hara_native/kernel/
session_snapshot.rs
