use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, OnceLock, RwLock};
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
use super::types::{OclaError, OclaResult};
use crate::core::ocla_bus::{self, OclaEvent};
static NEXT_FORK_ID: AtomicU64 = AtomicU64::new(1);
static NEXT_SNAPSHOT_ID: AtomicU64 = AtomicU64::new(1);
static GLOBAL_CAPSULE_STORE: OnceLock<CapsuleStore> = OnceLock::new();
#[must_use]
pub fn global_capsule_store() -> &'static CapsuleStore {
GLOBAL_CAPSULE_STORE.get_or_init(CapsuleStore::new)
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Delta {
pub offset: usize,
pub data: Vec<u8>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct CapsuleSnapshot {
pub snapshot_id: String,
pub capsule_ref: String,
pub content: Vec<u8>,
pub delta_count: usize,
pub created_at_ms: u64,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct HandoffDelta {
pub from_ref: String,
pub to_ref: String,
pub operations: Vec<DeltaOp>,
pub compressed_size: usize,
pub original_size: usize,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum DeltaOp {
Keep { offset: usize, len: usize },
Insert { offset: usize, data: Vec<u8> },
Delete { offset: usize, len: usize },
}
#[derive(Clone, Debug)]
pub struct CapsuleEntry {
pub parent_ref: Option<String>,
pub data: Vec<u8>,
pub deltas: Vec<Delta>,
pub budget_tokens: u64,
pub created_at: Instant,
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct CapsuleStats {
pub total_entries: usize,
pub total_bytes: usize,
pub max_depth: usize,
}
#[derive(Clone, Debug, Default)]
struct CapsuleStoreState {
entries: HashMap<String, CapsuleEntry>,
snapshots: Vec<CapsuleSnapshot>,
}
#[derive(Clone, Debug, Default)]
pub struct CapsuleStore {
state: Arc<RwLock<CapsuleStoreState>>,
}
impl CapsuleStore {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn register(&self, data: &[u8]) -> String {
let capsule_ref = format!("capsule:{}", blake3::hash(data).to_hex());
let entry = CapsuleEntry {
parent_ref: None,
data: data.to_vec(),
deltas: Vec::new(),
budget_tokens: 0,
created_at: Instant::now(),
};
if let Ok(mut state) = self.state.write() {
state.entries.insert(capsule_ref.clone(), entry);
}
capsule_ref
}
pub fn fork(&self, parent_ref: &str, budget_tokens: u64) -> OclaResult<String> {
let mut state = self
.state
.write()
.map_err(|_| invalid("capsule store lock poisoned"))?;
if !state.entries.contains_key(parent_ref) {
return Err(invalid(format!("unknown parent capsule: {parent_ref}")));
}
let fork_id = NEXT_FORK_ID.fetch_add(1, Ordering::Relaxed);
let identity = format!("{parent_ref}\0{budget_tokens}\0{fork_id}");
let capsule_ref = format!("capsule:{}", blake3::hash(identity.as_bytes()).to_hex());
state.entries.insert(
capsule_ref.clone(),
CapsuleEntry {
parent_ref: Some(parent_ref.to_string()),
data: Vec::new(),
deltas: Vec::new(),
budget_tokens,
created_at: Instant::now(),
},
);
Ok(capsule_ref)
}
pub fn resolve(&self, capsule_ref: &str) -> OclaResult<Vec<u8>> {
let state = self
.state
.read()
.map_err(|_| invalid("capsule store lock poisoned"))?;
resolve_entries(&state.entries, capsule_ref)
}
#[cfg(test)]
pub(crate) fn budget_tokens(&self, capsule_ref: &str) -> OclaResult<u64> {
let state = self
.state
.read()
.map_err(|_| invalid("capsule store lock poisoned"))?;
state
.entries
.get(capsule_ref)
.map(|entry| entry.budget_tokens)
.ok_or_else(|| invalid(format!("unknown capsule: {capsule_ref}")))
}
pub fn apply_delta(&self, capsule_ref: &str, delta: Delta) -> OclaResult<()> {
let mut state = self
.state
.write()
.map_err(|_| invalid("capsule store lock poisoned"))?;
let current = resolve_entries(&state.entries, capsule_ref)?;
let entry = state
.entries
.get_mut(capsule_ref)
.ok_or_else(|| invalid(format!("unknown capsule: {capsule_ref}")))?;
if entry.parent_ref.is_none() {
return Err(invalid("deltas can only be applied to forked capsules"));
}
if delta.offset > current.len() {
return Err(invalid("capsule delta starts beyond materialized content"));
}
entry.deltas.push(delta);
Ok(())
}
pub fn merge_back(&self, child_ref: &str) -> OclaResult<()> {
let mut state = self
.state
.write()
.map_err(|_| invalid("capsule store lock poisoned"))?;
resolve_entries(&state.entries, child_ref)?;
let (parent_ref, deltas) = {
let child = state
.entries
.get(child_ref)
.ok_or_else(|| invalid(format!("unknown capsule: {child_ref}")))?;
(
child
.parent_ref
.clone()
.ok_or_else(|| invalid("root capsules cannot merge back"))?,
child.deltas.clone(),
)
};
let parent = state
.entries
.get_mut(&parent_ref)
.ok_or_else(|| invalid(format!("unknown parent capsule: {parent_ref}")))?;
parent.deltas.extend(deltas);
state
.entries
.get_mut(child_ref)
.ok_or_else(|| invalid(format!("unknown capsule: {child_ref}")))?
.deltas
.clear();
Ok(())
}
pub fn snapshot(&self, capsule_ref: &str) -> OclaResult<CapsuleSnapshot> {
let mut state = self
.state
.write()
.map_err(|_| invalid("capsule store lock poisoned"))?;
let content = resolve_entries(&state.entries, capsule_ref)?;
let delta_count = state
.entries
.get(capsule_ref)
.ok_or_else(|| invalid(format!("unknown capsule: {capsule_ref}")))?
.deltas
.len();
let snapshot_number = NEXT_SNAPSHOT_ID.fetch_add(1, Ordering::Relaxed);
let snapshot_id = format!(
"snapshot:{}",
blake3::hash(format!("{capsule_ref}\\0{snapshot_number}").as_bytes()).to_hex()
);
let snapshot = CapsuleSnapshot {
snapshot_id,
capsule_ref: capsule_ref.to_string(),
content,
delta_count,
created_at_ms: unix_time_ms(),
};
state.snapshots.push(snapshot.clone());
Ok(snapshot)
}
pub fn rollback(&self, capsule_ref: &str, snapshot_id: &str) -> OclaResult<()> {
let mut state = self
.state
.write()
.map_err(|_| invalid("capsule store lock poisoned"))?;
let snapshot = state
.snapshots
.iter()
.find(|snapshot| {
snapshot.snapshot_id == snapshot_id && snapshot.capsule_ref == capsule_ref
})
.cloned()
.ok_or_else(|| invalid(format!("unknown capsule snapshot: {snapshot_id}")))?;
let entry = state
.entries
.get_mut(capsule_ref)
.ok_or_else(|| invalid(format!("unknown capsule: {capsule_ref}")))?;
entry.parent_ref = None;
entry.data = snapshot.content;
entry.deltas.clear();
ocla_bus::emit(OclaEvent::AgentChainEvent {
agent_id: capsule_ref.to_string(),
action: "capsule_rollback".to_string(),
parent_agent: Some(snapshot_id.to_string()),
});
Ok(())
}
pub fn compute_handoff_delta(&self, from_ref: &str, to_ref: &str) -> OclaResult<HandoffDelta> {
let state = self
.state
.read()
.map_err(|_| invalid("capsule store lock poisoned"))?;
let from = resolve_entries(&state.entries, from_ref)?;
let to = resolve_entries(&state.entries, to_ref)?;
let operations = compute_delta_operations(&from, &to);
Ok(HandoffDelta {
from_ref: from_ref.to_string(),
to_ref: to_ref.to_string(),
compressed_size: handoff_delta_size(&operations),
original_size: to.len(),
operations,
})
}
pub fn apply_handoff_delta(&self, base_ref: &str, delta: &HandoffDelta) -> OclaResult<String> {
if delta.from_ref != base_ref {
return Err(invalid("handoff delta base reference does not match"));
}
let mut state = self
.state
.write()
.map_err(|_| invalid("capsule store lock poisoned"))?;
let base = resolve_entries(&state.entries, base_ref)?;
let content = apply_handoff_operations(&base, &delta.operations)?;
let handoff_id = NEXT_FORK_ID.fetch_add(1, Ordering::Relaxed);
let identity = format!(
"handoff\\0{base_ref}\\0{}\\0{handoff_id}",
blake3::hash(&content)
);
let capsule_ref = format!("capsule:{}", blake3::hash(identity.as_bytes()).to_hex());
state.entries.insert(
capsule_ref.clone(),
CapsuleEntry {
parent_ref: None,
data: content,
deltas: Vec::new(),
budget_tokens: 0,
created_at: Instant::now(),
},
);
Ok(capsule_ref)
}
#[must_use]
pub fn stats(&self) -> CapsuleStats {
let Ok(state) = self.state.read() else {
return CapsuleStats::default();
};
let total_bytes = state.entries.values().fold(0_usize, |total, entry| {
total.saturating_add(entry.data.len()).saturating_add(
entry
.deltas
.iter()
.map(|delta| delta.data.len())
.sum::<usize>(),
)
});
let max_depth = state
.entries
.keys()
.map(|capsule_ref| depth_of(&state.entries, capsule_ref))
.max()
.unwrap_or(0);
CapsuleStats {
total_entries: state.entries.len(),
total_bytes,
max_depth,
}
}
}
fn resolve_entries(
entries: &HashMap<String, CapsuleEntry>,
capsule_ref: &str,
) -> OclaResult<Vec<u8>> {
let mut current = capsule_ref;
let mut visited = HashSet::new();
let mut layers = Vec::new();
loop {
if !visited.insert(current) {
return Err(invalid("capsule parent cycle detected"));
}
let entry = entries
.get(current)
.ok_or_else(|| invalid(format!("unknown capsule: {capsule_ref}")))?;
layers.push(entry.deltas.clone());
if let Some(parent_ref) = entry.parent_ref.as_deref() {
current = parent_ref;
} else {
let mut data = entry.data.clone();
for layer in layers.iter().rev() {
for delta in layer {
apply_patch(&mut data, delta)?;
}
}
return Ok(data);
}
}
}
fn apply_patch(data: &mut Vec<u8>, delta: &Delta) -> OclaResult<()> {
let end = delta
.offset
.checked_add(delta.data.len())
.ok_or_else(|| invalid("capsule delta range overflow"))?;
if delta.offset > data.len() {
return Err(invalid("capsule delta starts beyond materialized content"));
}
if end > data.len() {
data.resize(end, 0);
}
data[delta.offset..end].copy_from_slice(&delta.data);
Ok(())
}
fn compute_delta_operations(from: &[u8], to: &[u8]) -> Vec<DeltaOp> {
if from == to {
return Vec::new();
}
let prefix_len = from
.iter()
.zip(to)
.take_while(|(left, right)| left == right)
.count();
let suffix_len = from[prefix_len..]
.iter()
.rev()
.zip(to[prefix_len..].iter().rev())
.take_while(|(left, right)| left == right)
.count();
let from_middle_end = from.len() - suffix_len;
let to_middle_end = to.len() - suffix_len;
let mut operations = Vec::with_capacity(4);
if prefix_len > 0 {
operations.push(DeltaOp::Keep {
offset: 0,
len: prefix_len,
});
}
if from_middle_end > prefix_len {
operations.push(DeltaOp::Delete {
offset: prefix_len,
len: from_middle_end - prefix_len,
});
}
if to_middle_end > prefix_len {
operations.push(DeltaOp::Insert {
offset: from_middle_end,
data: to[prefix_len..to_middle_end].to_vec(),
});
}
if suffix_len > 0 {
operations.push(DeltaOp::Keep {
offset: from_middle_end,
len: suffix_len,
});
}
operations
}
fn apply_handoff_operations(base: &[u8], operations: &[DeltaOp]) -> OclaResult<Vec<u8>> {
if operations.is_empty() {
return Ok(base.to_vec());
}
let mut cursor = 0;
let mut output = Vec::with_capacity(base.len());
for operation in operations {
match operation {
DeltaOp::Keep { offset, len } => {
validate_handoff_offset(*offset, cursor)?;
let end = cursor
.checked_add(*len)
.ok_or_else(|| invalid("handoff keep range overflow"))?;
let bytes = base
.get(cursor..end)
.ok_or_else(|| invalid("handoff keep extends beyond base content"))?;
output.extend_from_slice(bytes);
cursor = end;
}
DeltaOp::Insert { offset, data } => {
validate_handoff_offset(*offset, cursor)?;
output.extend_from_slice(data);
}
DeltaOp::Delete { offset, len } => {
validate_handoff_offset(*offset, cursor)?;
cursor = cursor
.checked_add(*len)
.ok_or_else(|| invalid("handoff delete range overflow"))?;
if cursor > base.len() {
return Err(invalid("handoff delete extends beyond base content"));
}
}
}
}
if cursor != base.len() {
return Err(invalid("handoff delta does not consume base content"));
}
Ok(output)
}
fn validate_handoff_offset(offset: usize, cursor: usize) -> OclaResult<()> {
if offset != cursor {
return Err(invalid("handoff operation offset is out of order"));
}
Ok(())
}
fn handoff_delta_size(operations: &[DeltaOp]) -> usize {
operations.iter().fold(0_usize, |size, operation| {
let operation_size = match operation {
DeltaOp::Keep { .. } | DeltaOp::Delete { .. } => 2 * std::mem::size_of::<usize>(),
DeltaOp::Insert { data, .. } => 2 * std::mem::size_of::<usize>() + data.len(),
};
size.saturating_add(operation_size)
})
}
fn unix_time_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis()
.try_into()
.unwrap_or(u64::MAX)
}
fn depth_of(entries: &HashMap<String, CapsuleEntry>, capsule_ref: &str) -> usize {
let mut depth = 0;
let mut current = capsule_ref;
let mut visited = HashSet::new();
while visited.insert(current) {
let Some(entry) = entries.get(current) else {
break;
};
let Some(parent_ref) = entry.parent_ref.as_deref() else {
break;
};
depth += 1;
current = parent_ref;
}
depth
}
fn invalid(message: impl Into<String>) -> OclaError {
OclaError::InvalidRequest(message.into())
}
#[cfg(test)]
mod tests {
use super::{CapsuleStore, Delta, DeltaOp, global_capsule_store};
fn test_delta() -> Delta {
Delta {
offset: 1,
data: vec![b'a'],
}
}
#[test]
fn global_store_registers_capsule() {
let capsule_ref = global_capsule_store().register(b"global capsule");
assert_eq!(
global_capsule_store()
.resolve(&capsule_ref)
.expect("global resolves"),
b"global capsule"
);
}
#[test]
fn register_resolves_original_data() {
let store = CapsuleStore::new();
let capsule_ref = store.register(b"hello");
assert_eq!(
store.resolve(&capsule_ref).expect("root resolves"),
b"hello"
);
}
#[test]
fn fork_resolves_parent_data() {
let store = CapsuleStore::new();
let parent_ref = store.register(b"hello");
let child_ref = store.fork(&parent_ref, 100).expect("fork succeeds");
assert_eq!(store.resolve(&child_ref).expect("child resolves"), b"hello");
assert_eq!(store.budget_tokens(&child_ref).expect("budget exists"), 100);
}
#[test]
fn fork_delta_resolves_patched_data() {
let store = CapsuleStore::new();
let parent_ref = store.register(b"hello");
let child_ref = store.fork(&parent_ref, 100).expect("fork succeeds");
store
.apply_delta(&child_ref, test_delta())
.expect("delta applies");
assert_eq!(store.resolve(&child_ref).expect("child resolves"), b"hallo");
}
#[test]
fn later_overlapping_delta_wins_within_layer() {
let store = CapsuleStore::new();
let parent_ref = store.register(b"hello");
let fork_ref = store.fork(&parent_ref, 100).expect("fork succeeds");
store
.apply_delta(
&fork_ref,
Delta {
offset: 1,
data: vec![b'a'],
},
)
.expect("first delta applies");
store
.apply_delta(
&fork_ref,
Delta {
offset: 1,
data: vec![b'u'],
},
)
.expect("second delta applies");
assert_eq!(store.resolve(&fork_ref).expect("fork resolves"), b"hullo");
}
#[test]
fn child_delta_overrides_parent_delta() {
let store = CapsuleStore::new();
let root_ref = store.register(b"hello");
let parent_ref = store.fork(&root_ref, 100).expect("parent fork succeeds");
store
.apply_delta(
&parent_ref,
Delta {
offset: 1,
data: vec![b'a'],
},
)
.expect("parent delta applies");
let child_ref = store.fork(&parent_ref, 100).expect("child fork succeeds");
store
.apply_delta(
&child_ref,
Delta {
offset: 1,
data: vec![b'u'],
},
)
.expect("child delta applies");
assert_eq!(store.resolve(&child_ref).expect("child resolves"), b"hullo");
}
#[test]
fn merge_back_preserves_chronological_delta_order() {
let store = CapsuleStore::new();
let root_ref = store.register(b"hello");
let parent_ref = store.fork(&root_ref, 100).expect("parent fork succeeds");
store
.apply_delta(
&parent_ref,
Delta {
offset: 1,
data: vec![b'a'],
},
)
.expect("parent delta applies");
let child_ref = store.fork(&parent_ref, 100).expect("child fork succeeds");
store
.apply_delta(
&child_ref,
Delta {
offset: 1,
data: vec![b'u'],
},
)
.expect("child delta applies");
store.merge_back(&child_ref).expect("merge succeeds");
assert_eq!(
store.resolve(&parent_ref).expect("parent resolves"),
b"hullo"
);
assert_eq!(store.resolve(&child_ref).expect("child resolves"), b"hullo");
}
#[test]
fn merge_back_projects_deltas_to_parent() {
let store = CapsuleStore::new();
let parent_ref = store.register(b"hello");
let child_ref = store.fork(&parent_ref, 100).expect("fork succeeds");
store
.apply_delta(&child_ref, test_delta())
.expect("delta applies");
store.merge_back(&child_ref).expect("merge succeeds");
assert_eq!(
store.resolve(&parent_ref).expect("parent resolves"),
b"hallo"
);
assert_eq!(store.resolve(&child_ref).expect("child resolves"), b"hallo");
}
#[test]
fn stats_report_entries_storage_and_depth() {
let store = CapsuleStore::new();
let parent_ref = store.register(b"hello");
let child_ref = store.fork(&parent_ref, 100).expect("fork succeeds");
store
.apply_delta(&child_ref, test_delta())
.expect("delta applies");
let stats = store.stats();
assert_eq!(stats.total_entries, 2);
assert_eq!(stats.total_bytes, 6);
assert_eq!(stats.max_depth, 1);
}
#[test]
fn snapshot_rollback_restores_snapshot_content() {
let store = CapsuleStore::new();
let root_ref = store.register(b"hello");
let child_ref = store.fork(&root_ref, 100).expect("fork succeeds");
store
.apply_delta(
&child_ref,
Delta {
offset: 1,
data: vec![b'a'],
},
)
.expect("first delta applies");
let snapshot = store.snapshot(&child_ref).expect("snapshot succeeds");
store
.apply_delta(
&child_ref,
Delta {
offset: 1,
data: vec![b'u'],
},
)
.expect("second delta applies");
store
.rollback(&child_ref, &snapshot.snapshot_id)
.expect("rollback succeeds");
assert_eq!(store.resolve(&child_ref).expect("child resolves"), b"hallo");
}
#[test]
fn handoff_delta_for_identical_capsules_is_empty() {
let store = CapsuleStore::new();
let first_ref = store.register(b"unchanged");
let second_ref = store.register(b"unchanged");
let delta = store
.compute_handoff_delta(&first_ref, &second_ref)
.expect("delta computes");
assert!(delta.operations.is_empty());
assert_eq!(delta.compressed_size, 0);
}
#[test]
fn handoff_delta_for_different_capsules_has_insert_and_delete() {
let store = CapsuleStore::new();
let from_ref = store.register(b"hello");
let to_ref = store.register(b"halo!");
let delta = store
.compute_handoff_delta(&from_ref, &to_ref)
.expect("delta computes");
assert!(
delta
.operations
.iter()
.any(|operation| matches!(operation, DeltaOp::Insert { .. }))
);
assert!(
delta
.operations
.iter()
.any(|operation| matches!(operation, DeltaOp::Delete { .. }))
);
}
#[test]
fn apply_handoff_delta_round_trip_matches_target() {
let store = CapsuleStore::new();
let from_ref = store.register(b"goodbye");
let to_ref = store.register(b"good day!");
let delta = store
.compute_handoff_delta(&from_ref, &to_ref)
.expect("delta computes");
let applied_ref = store
.apply_handoff_delta(&from_ref, &delta)
.expect("delta applies");
assert_eq!(
store.resolve(&applied_ref).expect("applied resolves"),
store.resolve(&to_ref).expect("target resolves")
);
}
#[test]
fn rollback_with_invalid_snapshot_id_errors() {
let store = CapsuleStore::new();
let capsule_ref = store.register(b"hello");
assert!(store.rollback(&capsule_ref, "snapshot:missing").is_err());
}
}