use anyhow::{anyhow, Context, Result};
use bincode;
use chrono::{DateTime, Utc};
use rocksdb::{
ColumnFamily, ColumnFamilyDescriptor, IteratorMode, Options, WriteBatch, WriteOptions, DB,
};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use super::types::*;
trait LogErrors<T> {
fn log_errors(self) -> impl Iterator<Item = T>;
}
impl<I, T, E> LogErrors<T> for I
where
I: Iterator<Item = Result<T, E>>,
E: std::fmt::Display,
{
fn log_errors(self) -> impl Iterator<Item = T> {
self.filter_map(|r| match r {
Ok(v) => Some(v),
Err(e) => {
tracing::warn!("RocksDB iterator error (continuing): {}", e);
None
}
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WriteMode {
Sync,
Async,
}
impl Default for WriteMode {
fn default() -> Self {
match std::env::var("SHODH_WRITE_MODE") {
Ok(mode) if mode.to_lowercase() == "sync" => WriteMode::Sync,
_ => WriteMode::Async,
}
}
}
pub(crate) const STORAGE_MAGIC: &[u8; 3] = b"SHO";
use std::collections::HashMap;
fn default_legacy_experience_type() -> ExperienceType {
ExperienceType::Observation
}
#[derive(Deserialize)]
struct MinimalMemory {
id: MemoryId,
content: String,
}
#[derive(Deserialize)]
struct MemoryWithTypePrefix {
id: MemoryId,
_unknown_field: u8, experience_type: u8, content: String, }
#[derive(Deserialize)]
struct MemoryWith3ByteHeader {
id: MemoryId,
_header1: u8, _header2: u8, _header3: u8, content: String, }
impl MemoryWith3ByteHeader {
fn into_memory(self) -> Memory {
let now = Utc::now();
let experience = Experience {
experience_type: ExperienceType::Observation,
content: self.content,
..Default::default()
};
Memory::from_legacy(
self.id,
experience,
0.5,
0,
now,
now,
false,
MemoryTier::LongTerm,
Vec::new(),
1.0,
None,
None,
None,
None,
0.0,
None,
None,
1,
Vec::new(),
Vec::new(),
)
}
}
fn is_plausible_msgpack_struct(data: &[u8]) -> bool {
if data.is_empty() {
return false;
}
let first = data[0];
matches!(first, 0x80..=0x9f | 0xdc..=0xdf)
}
const MAX_MSGPACK_DESER_SIZE: usize = 1024 * 1024;
fn try_msgpack_deserialize<T: serde::de::DeserializeOwned>(
data: &[u8],
) -> Option<Result<T, String>> {
if !is_plausible_msgpack_struct(data) {
return None;
}
if data.len() > MAX_MSGPACK_DESER_SIZE {
return Some(Err(format!(
"msgpack data too large ({} bytes, max {})",
data.len(),
MAX_MSGPACK_DESER_SIZE
)));
}
match rmp_serde::from_slice::<T>(data) {
Ok(val) => Some(Ok(val)),
Err(e) => Some(Err(e.to_string())),
}
}
fn try_raw_memory_parse(data: &[u8]) -> Option<Memory> {
if data.len() < 20 {
return None;
}
let uuid_bytes: [u8; 16] = data[0..16].try_into().ok()?;
let id = MemoryId(uuid::Uuid::from_bytes(uuid_bytes));
for header_skip in [2, 3, 4, 5, 6] {
let content_start = 16 + header_skip;
if content_start >= data.len() {
continue;
}
if let Ok(content) = std::str::from_utf8(&data[content_start..]) {
if !content.is_empty()
&& content
.chars()
.next()
.map(|c| c.is_ascii_graphic())
.unwrap_or(false)
{
let now = Utc::now();
let experience = Experience {
experience_type: ExperienceType::Observation,
content: content.to_string(),
..Default::default()
};
tracing::debug!(
"Recovered memory with raw parsing (header_skip={}, content_len={})",
header_skip,
content.len()
);
return Some(Memory::from_legacy(
id,
experience,
0.5,
0,
now,
now,
false,
MemoryTier::LongTerm,
Vec::new(),
1.0,
None,
None,
None,
None,
0.0,
None,
None,
1,
Vec::new(),
Vec::new(),
));
}
}
}
None
}
impl MemoryWithTypePrefix {
fn into_memory(self) -> Memory {
let now = Utc::now();
let exp_type = match self.experience_type {
0 => ExperienceType::Observation,
1 => ExperienceType::Decision,
2 => ExperienceType::Learning,
3 => ExperienceType::Error,
4 => ExperienceType::Discovery,
5 => ExperienceType::Pattern,
6 => ExperienceType::Context,
7 => ExperienceType::Task,
8 => ExperienceType::CodeEdit,
9 => ExperienceType::FileAccess,
10 => ExperienceType::Search,
11 => ExperienceType::Command,
12 => ExperienceType::Conversation,
_ => ExperienceType::Observation,
};
let experience = Experience {
experience_type: exp_type,
content: self.content,
..Default::default()
};
Memory::from_legacy(
self.id,
experience,
0.5,
0,
now,
now,
false,
MemoryTier::LongTerm,
Vec::new(),
1.0,
None,
None,
None,
None,
0.0,
None,
None,
1,
Vec::new(),
Vec::new(),
)
}
}
impl MinimalMemory {
fn into_memory(self) -> Memory {
let now = Utc::now();
let experience = Experience {
experience_type: ExperienceType::Observation,
content: self.content,
..Default::default()
};
Memory::from_legacy(
self.id,
experience,
0.5, 0,
now,
now,
false,
MemoryTier::LongTerm,
Vec::new(),
1.0,
None,
None,
None,
None,
0.0,
None,
None,
1,
Vec::new(),
Vec::new(),
)
}
}
#[derive(Deserialize)]
struct SimpleLegacyMemory {
id: MemoryId,
content: String, #[serde(default)]
importance: f32,
#[serde(default)]
access_count: u32,
#[serde(default)]
created_at: Option<DateTime<Utc>>,
#[serde(default)]
last_accessed: Option<DateTime<Utc>>,
#[serde(default)]
compressed: bool,
#[serde(default)]
agent_id: Option<String>,
#[serde(default)]
run_id: Option<String>,
#[serde(default)]
actor_id: Option<String>,
#[serde(default)]
temporal_relevance: f32,
#[serde(default)]
score: Option<f32>,
}
impl SimpleLegacyMemory {
fn into_memory(self) -> Memory {
let now = Utc::now();
let experience = Experience {
experience_type: ExperienceType::Observation,
content: self.content,
..Default::default()
};
Memory::from_legacy(
self.id,
experience,
if self.importance > 0.0 {
self.importance
} else {
0.5
},
self.access_count,
self.created_at.unwrap_or(now),
self.last_accessed.unwrap_or(now),
self.compressed,
MemoryTier::LongTerm,
Vec::new(),
1.0,
None,
self.agent_id,
self.run_id,
self.actor_id,
self.temporal_relevance,
self.score,
None,
1,
Vec::new(),
Vec::new(),
)
}
}
#[derive(Deserialize)]
struct LegacyExperienceV1 {
#[serde(default = "default_legacy_experience_type")]
experience_type: ExperienceType,
content: String,
#[serde(default)]
context: Option<RichContext>,
#[serde(default)]
entities: Vec<String>,
#[serde(default)]
metadata: HashMap<String, String>,
#[serde(default)]
embeddings: Option<Vec<f32>>,
#[serde(default)]
related_memories: Vec<MemoryId>,
#[serde(default)]
causal_chain: Vec<MemoryId>,
#[serde(default)]
outcomes: Vec<String>,
#[serde(default)]
robot_id: Option<String>,
#[serde(default)]
mission_id: Option<String>,
#[serde(default)]
geo_location: Option<[f64; 3]>,
#[serde(default)]
local_position: Option<[f32; 3]>,
#[serde(default)]
heading: Option<f32>,
#[serde(default)]
action_type: Option<String>,
#[serde(default)]
reward: Option<f32>,
#[serde(default)]
sensor_data: HashMap<String, f64>,
#[serde(default)]
decision_context: Option<HashMap<String, String>>,
#[serde(default)]
action_params: Option<HashMap<String, String>>,
#[serde(default)]
outcome_type: Option<String>,
#[serde(default)]
outcome_details: Option<String>,
#[serde(default)]
confidence: Option<f32>,
#[serde(default)]
alternatives_considered: Vec<String>,
#[serde(default)]
weather: Option<HashMap<String, String>>,
#[serde(default)]
terrain_type: Option<String>,
#[serde(default)]
lighting: Option<String>,
#[serde(default)]
nearby_agents: Vec<HashMap<String, String>>,
#[serde(default)]
is_failure: bool,
#[serde(default)]
is_anomaly: bool,
#[serde(default)]
severity: Option<String>,
#[serde(default)]
recovery_action: Option<String>,
#[serde(default)]
root_cause: Option<String>,
#[serde(default)]
pattern_id: Option<String>,
#[serde(default)]
predicted_outcome: Option<String>,
#[serde(default)]
prediction_accurate: Option<bool>,
#[serde(default)]
tags: Vec<String>,
}
impl LegacyExperienceV1 {
fn into_experience(self) -> Experience {
Experience {
experience_type: self.experience_type,
content: self.content,
context: self.context,
entities: self.entities,
metadata: self.metadata,
embeddings: self.embeddings,
image_embeddings: None,
audio_embeddings: None,
video_embeddings: None,
media_refs: Vec::new(),
related_memories: self.related_memories,
causal_chain: self.causal_chain,
outcomes: self.outcomes,
robot_id: self.robot_id,
mission_id: self.mission_id,
geo_location: self.geo_location,
local_position: self.local_position,
heading: self.heading,
action_type: self.action_type,
reward: self.reward,
sensor_data: self.sensor_data,
decision_context: self.decision_context,
action_params: self.action_params,
outcome_type: self.outcome_type,
outcome_details: self.outcome_details,
confidence: self.confidence,
alternatives_considered: self.alternatives_considered,
weather: self.weather,
terrain_type: self.terrain_type,
lighting: self.lighting,
nearby_agents: self.nearby_agents,
is_failure: self.is_failure,
is_anomaly: self.is_anomaly,
severity: self.severity,
recovery_action: self.recovery_action,
root_cause: self.root_cause,
pattern_id: self.pattern_id,
predicted_outcome: self.predicted_outcome,
prediction_accurate: self.prediction_accurate,
tags: self.tags,
temporal_refs: Vec::new(),
ner_entities: Vec::new(),
cooccurrence_pairs: Vec::new(),
importance_override: None,
}
}
}
#[derive(Deserialize)]
struct LegacyMemoryV1Full {
#[serde(rename = "memory_id")]
id: MemoryId,
experience: LegacyExperienceV1,
importance: f32,
access_count: u32,
created_at: DateTime<Utc>,
last_accessed: DateTime<Utc>,
compressed: bool,
agent_id: Option<String>,
run_id: Option<String>,
actor_id: Option<String>,
temporal_relevance: f32,
score: Option<f32>,
}
impl LegacyMemoryV1Full {
fn into_memory(self) -> Memory {
Memory::from_legacy(
self.id,
self.experience.into_experience(),
self.importance,
self.access_count,
self.created_at,
self.last_accessed,
self.compressed,
MemoryTier::LongTerm,
Vec::new(),
1.0,
None,
self.agent_id,
self.run_id,
self.actor_id,
self.temporal_relevance,
self.score,
None,
1,
Vec::new(),
Vec::new(),
)
}
}
#[derive(Deserialize)]
struct LegacyMemoryV1 {
#[serde(rename = "memory_id")]
id: MemoryId,
experience: LegacyExperienceV1, importance: f32,
access_count: u32,
created_at: DateTime<Utc>,
last_accessed: DateTime<Utc>,
compressed: bool,
agent_id: Option<String>,
run_id: Option<String>,
actor_id: Option<String>,
temporal_relevance: f32,
score: Option<f32>,
}
impl LegacyMemoryV1 {
fn into_memory(self) -> Memory {
Memory::from_legacy(
self.id,
self.experience.into_experience(),
self.importance,
self.access_count,
self.created_at,
self.last_accessed,
self.compressed,
MemoryTier::LongTerm,
Vec::new(),
1.0,
None,
self.agent_id,
self.run_id,
self.actor_id,
self.temporal_relevance,
self.score,
None,
1,
Vec::new(),
Vec::new(),
)
}
}
#[derive(Deserialize)]
struct LegacyMemoryV2 {
id: MemoryId,
experience: LegacyExperienceV1, importance: f32,
access_count: u32,
created_at: DateTime<Utc>,
last_accessed: DateTime<Utc>,
compressed: bool,
tier: MemoryTier,
entity_refs: Vec<EntityRef>,
activation: f32,
last_retrieval_id: Option<uuid::Uuid>,
agent_id: Option<String>,
run_id: Option<String>,
actor_id: Option<String>,
temporal_relevance: f32,
score: Option<f32>,
}
impl LegacyMemoryV2 {
fn into_memory(self) -> Memory {
Memory::from_legacy(
self.id,
self.experience.into_experience(),
self.importance,
self.access_count,
self.created_at,
self.last_accessed,
self.compressed,
self.tier,
self.entity_refs,
self.activation,
self.last_retrieval_id,
self.agent_id,
self.run_id,
self.actor_id,
self.temporal_relevance,
self.score,
None, 1, Vec::new(), Vec::new(), )
}
}
fn deserialize_memory(data: &[u8]) -> Result<(Memory, bool)> {
use crate::serialization::{SHO_VERSION_BINCODE2, SHO_VERSION_POSTCARD};
if let Some((version, payload)) = crate::serialization::unwrap_sho(data) {
match version {
SHO_VERSION_POSTCARD => {
let memory: Memory = crate::serialization::decode_raw(payload)
.map_err(|e| anyhow!("SHO v2 postcard decode failed: {e}"))?;
Ok((memory, false))
}
SHO_VERSION_BINCODE2 => {
let (memory, _): (Memory, _) =
bincode::serde::decode_from_slice(payload, crate::bincode_safe_config())
.map_err(|e| anyhow!("SHO v1 bincode decode failed: {e}"))?;
Ok((memory, true))
}
_ => {
deserialize_with_fallback(payload)
.map_err(|e| anyhow!("SHO v{version} decode failed: {e}"))
}
}
} else {
deserialize_with_fallback(data)
.map_err(|e| anyhow!("legacy (no SHO header) decode failed: {e}"))
}
}
pub fn deserialize_memory_for_migration(data: &[u8]) -> Result<Memory> {
deserialize_memory(data).map(|(m, _)| m)
}
#[derive(Deserialize)]
struct LegacyMemoryFlatV2 {
id: MemoryId,
experience: LegacyExperienceV1, importance: f32,
access_count: u32,
created_at: DateTime<Utc>,
last_accessed: DateTime<Utc>,
compressed: bool,
tier: MemoryTier,
entity_refs: Vec<EntityRef>,
activation: f32,
last_retrieval_id: Option<uuid::Uuid>,
agent_id: Option<String>,
run_id: Option<String>,
actor_id: Option<String>,
temporal_relevance: f32,
score: Option<f32>,
external_id: Option<String>,
version: u32,
history: Vec<MemoryRevision>,
#[serde(default)]
related_todo_ids: Vec<TodoId>,
}
impl LegacyMemoryFlatV2 {
fn into_memory(self) -> Memory {
Memory::from_legacy(
self.id,
self.experience.into_experience(),
self.importance,
self.access_count,
self.created_at,
self.last_accessed,
self.compressed,
self.tier,
self.entity_refs,
self.activation,
self.last_retrieval_id,
self.agent_id,
self.run_id,
self.actor_id,
self.temporal_relevance,
self.score,
self.external_id,
self.version,
self.history,
self.related_todo_ids,
)
}
}
fn deserialize_with_fallback(data: &[u8]) -> Result<(Memory, bool)> {
fn record_branch(branch: &str) {
crate::metrics::LEGACY_FALLBACK_BRANCH_TOTAL
.with_label_values(&[branch])
.inc();
}
match bincode::serde::decode_from_slice::<Memory, _>(data, crate::bincode_safe_config()) {
Ok((memory, _)) => {
return Ok((memory, false));
} Err(e) => {
return deserialize_legacy_fallback(data, e, record_branch);
}
}
}
fn deserialize_legacy_fallback(
data: &[u8],
first_error: bincode::error::DecodeError,
record_branch: fn(&str),
) -> Result<(Memory, bool)> {
static DEBUG_ENTRY_LOGGED: std::sync::atomic::AtomicBool =
std::sync::atomic::AtomicBool::new(false);
let is_first_failure = !DEBUG_ENTRY_LOGGED.load(std::sync::atomic::Ordering::Relaxed);
let mut errors: Vec<(&str, String)> = Vec::new();
errors.push(("bincode2 Memory", first_error.to_string()));
match bincode::serde::decode_from_slice::<MinimalMemory, _>(data, crate::bincode_safe_config())
{
Ok((minimal, _)) => {
tracing::debug!("Migrated memory from bincode 2.x minimal format");
record_branch("bincode2_minimal");
return Ok((minimal.into_memory(), true));
}
Err(e) => errors.push(("bincode2 MinimalMemory", e.to_string())),
}
match bincode::serde::decode_from_slice::<MemoryWithTypePrefix, _>(
data,
crate::bincode_safe_config(),
) {
Ok((typed, _)) => {
tracing::debug!("Migrated memory from bincode 2.x with type prefix");
record_branch("bincode2_type_prefix");
return Ok((typed.into_memory(), true));
}
Err(e) => errors.push(("bincode2 MemoryWithTypePrefix", e.to_string())),
}
match bincode::serde::decode_from_slice::<LegacyMemoryFlatV2, _>(
data,
crate::bincode_safe_config(),
) {
Ok((legacy, _)) => {
tracing::debug!("Migrated memory from bincode 2.x pre-multimodal format");
record_branch("bincode2_legacy_flat_v2");
return Ok((legacy.into_memory(), true));
}
Err(e) => errors.push(("bincode2 LegacyMemoryFlatV2", e.to_string())),
}
use bincode1::Options;
let bincode1_safe = bincode1::options()
.with_fixint_encoding()
.with_limit(data.len() as u64 + 1024)
.allow_trailing_bytes();
match bincode1_safe.deserialize::<LegacyMemoryV1>(data) {
Ok(legacy) => {
tracing::debug!("Migrated memory from bincode 1.x v0.1.0 format");
record_branch("bincode1_legacy_v1");
return Ok((legacy.into_memory(), true));
}
Err(e) => errors.push(("bincode1 LegacyMemoryV1", e.to_string())),
}
match bincode1_safe.deserialize::<MinimalMemory>(data) {
Ok(minimal) => {
tracing::debug!("Migrated memory from bincode 1.x minimal format");
record_branch("bincode1_minimal");
return Ok((minimal.into_memory(), true));
}
Err(e) => errors.push(("bincode1 MinimalMemory", e.to_string())),
}
match bincode1_safe.deserialize::<SimpleLegacyMemory>(data) {
Ok(legacy) => {
tracing::debug!("Migrated memory from bincode 1.x simple format");
record_branch("bincode1_simple");
return Ok((legacy.into_memory(), true));
}
Err(e) => errors.push(("bincode1 SimpleLegacyMemory", e.to_string())),
}
let fixint_config = bincode1::options()
.with_fixint_encoding()
.with_limit(data.len() as u64 + 1024)
.allow_trailing_bytes();
match fixint_config.deserialize::<MinimalMemory>(data) {
Ok(minimal) => {
tracing::debug!("Migrated memory from bincode 1.x fixint minimal format");
record_branch("bincode1_fixint_minimal");
return Ok((minimal.into_memory(), true));
}
Err(e) => errors.push(("bincode1 fixint MinimalMemory", e.to_string())),
}
match fixint_config.deserialize::<SimpleLegacyMemory>(data) {
Ok(legacy) => {
tracing::debug!("Migrated memory from bincode 1.x fixint simple format");
record_branch("bincode1_fixint_simple");
return Ok((legacy.into_memory(), true));
}
Err(e) => errors.push(("bincode1 fixint SimpleLegacyMemory", e.to_string())),
}
match try_msgpack_deserialize::<MinimalMemory>(data) {
Some(Ok(minimal)) => {
tracing::debug!("Migrated memory from MessagePack minimal format");
record_branch("msgpack_minimal");
return Ok((minimal.into_memory(), true));
}
Some(Err(e)) => errors.push(("msgpack MinimalMemory", e)),
None => errors.push(("msgpack MinimalMemory", "not msgpack format".to_string())),
}
match try_msgpack_deserialize::<SimpleLegacyMemory>(data) {
Some(Ok(legacy)) => {
tracing::debug!("Migrated memory from MessagePack simple format");
record_branch("msgpack_simple");
return Ok((legacy.into_memory(), true));
}
Some(Err(e)) => errors.push(("msgpack SimpleLegacyMemory", e)),
None => errors.push((
"msgpack SimpleLegacyMemory",
"not msgpack format".to_string(),
)),
}
match bincode1_safe.deserialize::<LegacyMemoryV1Full>(data) {
Ok(legacy) => {
tracing::debug!("Migrated memory from bincode 1.x v1 full format");
record_branch("bincode1_legacy_v1_full");
return Ok((legacy.into_memory(), true));
}
Err(e) => errors.push(("bincode1 LegacyMemoryV1Full", e.to_string())),
}
match fixint_config.deserialize::<LegacyMemoryV1Full>(data) {
Ok(legacy) => {
tracing::debug!("Migrated memory from bincode 1.x fixint v1 full format");
record_branch("bincode1_fixint_legacy_v1_full");
return Ok((legacy.into_memory(), true));
}
Err(e) => errors.push(("bincode1 fixint LegacyMemoryV1Full", e.to_string())),
}
match try_msgpack_deserialize::<LegacyMemoryV1Full>(data) {
Some(Ok(legacy)) => {
tracing::debug!("Migrated memory from MessagePack v1 full format");
record_branch("msgpack_legacy_v1_full");
return Ok((legacy.into_memory(), true));
}
Some(Err(e)) => errors.push(("msgpack LegacyMemoryV1Full", e)),
None => errors.push((
"msgpack LegacyMemoryV1Full",
"not msgpack format".to_string(),
)),
}
match bincode1_safe.deserialize::<LegacyMemoryV1>(data) {
Ok(legacy) => {
tracing::debug!("Migrated memory from bincode 1.x format");
record_branch("bincode1_legacy_v1_repeat");
return Ok((legacy.into_memory(), true));
}
Err(_) => {} }
match bincode1_safe.deserialize::<LegacyMemoryV2>(data) {
Ok(legacy) => {
tracing::debug!("Migrated memory from bincode 1.x v2 format");
record_branch("bincode1_legacy_v2");
return Ok((legacy.into_memory(), true));
}
Err(e) => errors.push(("bincode1 LegacyMemoryV2", e.to_string())),
}
match bincode::serde::decode_from_slice::<MemoryWith3ByteHeader, _>(
data,
crate::bincode_safe_config(),
) {
Ok((mem, _)) => {
tracing::debug!("Migrated memory from bincode 2.x with 3-byte header");
record_branch("bincode2_3byte_header");
return Ok((mem.into_memory(), true));
}
Err(e) => errors.push(("bincode2 MemoryWith3ByteHeader", e.to_string())),
}
if let Some(memory) = try_raw_memory_parse(data) {
record_branch("raw_parse");
return Ok((memory, true));
}
errors.push(("raw parse", "no valid UTF-8 content found".to_string()));
if is_first_failure {
DEBUG_ENTRY_LOGGED.store(true, std::sync::atomic::Ordering::Relaxed);
let hex_preview: String = data
.iter()
.take(32)
.map(|b| format!("{:02x}", b))
.collect::<Vec<_>>()
.join(" ");
tracing::debug!(
"Unknown memory format ({} bytes): {}...",
data.len(),
hex_preview
);
}
record_branch("decode_failed");
Err(anyhow!(
"Failed to deserialize memory: incompatible format ({} bytes)",
data.len()
))
}
pub(crate) fn crc32_simple(data: &[u8]) -> u32 {
let mut crc: u32 = 0xFFFFFFFF;
for byte in data {
crc ^= *byte as u32;
for _ in 0..8 {
crc = if crc & 1 != 0 {
(crc >> 1) ^ 0xEDB88320
} else {
crc >> 1
};
}
}
!crc
}
const CF_INDEX: &str = "memory_index";
const WRITE_RETRY_BUFFER_CAPACITY: usize = 100;
pub struct MemoryStorage {
db: Arc<DB>,
storage_path: PathBuf,
write_mode: WriteMode,
write_retry_buffer: parking_lot::Mutex<std::collections::VecDeque<Memory>>,
write_failure_count: std::sync::atomic::AtomicU64,
}
impl MemoryStorage {
fn index_cf(&self) -> &ColumnFamily {
self.db
.cf_handle(CF_INDEX)
.expect("memory_index CF must exist")
}
pub fn new(path: &Path, shared_cache: Option<&rocksdb::Cache>) -> Result<Self> {
use crate::constants::ROCKSDB_MEMORY_WRITE_BUFFER_BYTES;
let storage_path = path.join("storage");
std::fs::create_dir_all(&storage_path)?;
let mut opts = Options::default();
opts.create_if_missing(true);
opts.create_missing_column_families(true);
opts.set_compression_type(rocksdb::DBCompressionType::Lz4);
opts.set_manual_wal_flush(false);
opts.set_max_write_buffer_number(2);
opts.set_write_buffer_size(ROCKSDB_MEMORY_WRITE_BUFFER_BYTES);
opts.set_level_zero_file_num_compaction_trigger(4);
opts.set_target_file_size_base(64 * 1024 * 1024); opts.set_max_bytes_for_level_base(256 * 1024 * 1024); opts.set_max_background_jobs(4);
opts.set_level_compaction_dynamic_level_bytes(true);
use rocksdb::{BlockBasedOptions, Cache};
let mut block_opts = BlockBasedOptions::default();
block_opts.set_bloom_filter(10.0, false); let local_cache;
let cache = match shared_cache {
Some(c) => c,
None => {
local_cache = Cache::new_lru_cache(16 * 1024 * 1024); &local_cache
}
};
block_opts.set_block_cache(cache);
block_opts.set_cache_index_and_filter_blocks(true);
block_opts.set_pin_l0_filter_and_index_blocks_in_cache(true); opts.set_block_based_table_factory(&block_opts);
let main_opts = opts.clone();
let db = Arc::new(Self::open_or_repair_cf(&opts, &storage_path, move || {
vec![
ColumnFamilyDescriptor::new("default", main_opts.clone()),
ColumnFamilyDescriptor::new(CF_INDEX, {
let mut idx_opts = Options::default();
idx_opts.create_if_missing(true);
idx_opts.set_compression_type(rocksdb::DBCompressionType::Lz4);
idx_opts.set_max_write_buffer_number(2);
idx_opts.set_write_buffer_size(ROCKSDB_MEMORY_WRITE_BUFFER_BYTES);
idx_opts
}),
]
})?);
Self::migrate_from_separate_dbs(path, &db)?;
let write_mode = WriteMode::default();
tracing::info!(
"Storage initialized with {:?} write mode (latency: {})",
write_mode,
if write_mode == WriteMode::Sync {
"2-10ms per write"
} else {
"<1ms per write"
}
);
Ok(Self {
db,
storage_path: path.to_path_buf(),
write_mode,
write_retry_buffer: parking_lot::Mutex::new(std::collections::VecDeque::new()),
write_failure_count: std::sync::atomic::AtomicU64::new(0),
})
}
fn open_or_repair_cf<F>(opts: &Options, path: &Path, build_cfs: F) -> Result<DB>
where
F: Fn() -> Vec<ColumnFamilyDescriptor>,
{
match DB::open_cf_descriptors(opts, path, build_cfs()) {
Ok(db) => Ok(db),
Err(open_err) => {
let err_str = open_err.to_string();
if err_str.contains("Corruption")
|| err_str.contains("bad block")
|| err_str.contains("checksum mismatch")
|| err_str.contains("MANIFEST")
|| err_str.contains("CURRENT")
{
tracing::warn!(
error = %open_err,
"RocksDB corruption detected in memory storage, attempting repair"
);
if let Err(repair_err) = DB::repair(opts, path) {
tracing::error!(
error = %repair_err,
"RocksDB repair failed for memory storage"
);
return Err(anyhow::anyhow!(
"Failed to open or repair memory storage: open={open_err}, repair={repair_err}"
));
}
tracing::info!("RocksDB repair succeeded for memory storage, reopening");
DB::open_cf_descriptors(opts, path, build_cfs()).map_err(|e| {
anyhow::anyhow!("Failed to open memory storage after repair: {e}")
})
} else {
Err(anyhow::anyhow!("Failed to open memory storage: {open_err}"))
}
}
}
}
fn migrate_from_separate_dbs(base_path: &Path, db: &DB) -> Result<()> {
let old_memories_dir = base_path.join("memories");
let old_index_dir = base_path.join("memory_index");
let has_old_memories = old_memories_dir.is_dir();
let has_old_index = old_index_dir.is_dir();
if !has_old_memories && !has_old_index {
return Ok(());
}
tracing::info!("Detected old separate-DB layout, migrating to column families...");
let mut total_migrated = 0usize;
if has_old_memories {
let old_opts = Options::default();
match DB::open_for_read_only(&old_opts, &old_memories_dir, false) {
Ok(old_db) => {
let mut batch = WriteBatch::default();
let mut count = 0usize;
for item in old_db.iterator(IteratorMode::Start) {
if let Ok((key, value)) = item {
batch.put(&key, &value);
count += 1;
if count % 10_000 == 0 {
db.write(std::mem::take(&mut batch))?;
tracing::info!(" memories: migrated {count} entries...");
}
}
}
if !batch.is_empty() {
db.write(batch)?;
}
drop(old_db);
total_migrated += count;
tracing::info!(" memories: migrated {count} entries to default CF");
let backup_name = base_path.join("memories.pre_cf_migration");
if backup_name.exists() {
let _ = std::fs::remove_dir_all(&backup_name);
}
if let Err(e) = std::fs::rename(&old_memories_dir, &backup_name) {
tracing::warn!("Could not rename old memories dir: {e}");
}
}
Err(e) => {
tracing::warn!("Could not open old memories DB for migration: {e}");
}
}
}
if has_old_index {
let index_cf = db
.cf_handle(CF_INDEX)
.expect("memory_index CF must exist during migration");
let old_opts = Options::default();
match DB::open_for_read_only(&old_opts, &old_index_dir, false) {
Ok(old_db) => {
let mut batch = WriteBatch::default();
let mut count = 0usize;
for item in old_db.iterator(IteratorMode::Start) {
if let Ok((key, value)) = item {
batch.put_cf(&index_cf, &key, &value);
count += 1;
if count % 10_000 == 0 {
db.write(std::mem::take(&mut batch))?;
tracing::info!(" index: migrated {count} entries...");
}
}
}
if !batch.is_empty() {
db.write(batch)?;
}
drop(old_db);
total_migrated += count;
tracing::info!(" index: migrated {count} entries to {CF_INDEX} CF");
let backup_name = base_path.join("memory_index.pre_cf_migration");
if backup_name.exists() {
let _ = std::fs::remove_dir_all(&backup_name);
}
if let Err(e) = std::fs::rename(&old_index_dir, &backup_name) {
tracing::warn!("Could not rename old memory_index dir: {e}");
}
}
Err(e) => {
tracing::warn!("Could not open old memory_index DB for migration: {e}");
}
}
}
if total_migrated > 0 {
tracing::info!(
"Memory storage migration complete: {total_migrated} total entries migrated"
);
}
Ok(())
}
pub fn path(&self) -> &Path {
&self.storage_path
}
pub fn store(&self, memory: &Memory) -> Result<()> {
self.drain_retry_buffer();
match self.store_inner(memory) {
Ok(()) => Ok(()),
Err(e) => {
let err_str = e.to_string();
if err_str.contains("lock")
|| err_str.contains("LOCK")
|| err_str.contains("disk")
|| err_str.contains("space")
|| err_str.contains("I/O")
{
self.write_failure_count
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let mut buffer = self.write_retry_buffer.lock();
if buffer.len() < WRITE_RETRY_BUFFER_CAPACITY {
tracing::warn!(
memory_id = %memory.id.0,
buffer_len = buffer.len() + 1,
error = %e,
"Write failed, buffered for retry"
);
buffer.push_back(memory.clone());
} else {
tracing::error!(
memory_id = %memory.id.0,
error = %e,
"Write failed and retry buffer full — memory dropped"
);
}
}
Err(e)
}
}
}
fn store_inner(&self, memory: &Memory) -> Result<()> {
let key = memory.id.0.as_bytes();
let value = crate::serialization::encode_sho(memory)
.context(format!("Failed to serialize memory {}", memory.id.0))?;
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db
.put_opt(key, &value, &write_opts)
.context(format!("Failed to put memory {} in RocksDB", memory.id.0))?;
if let Err(e) = self.update_indices(memory) {
if let Err(del_err) = self.db.delete(key) {
tracing::error!(
"Index write failed AND rollback failed for memory {}: index_err={}, delete_err={}",
memory.id.0, e, del_err
);
}
return Err(e.context(format!(
"Failed to update indices for memory {} (rolled back)",
memory.id.0
)));
}
Ok(())
}
pub fn drain_retry_buffer(&self) -> usize {
let mut buffer = self.write_retry_buffer.lock();
if buffer.is_empty() {
return 0;
}
let count = buffer.len();
let mut succeeded = 0;
let mut still_failing = std::collections::VecDeque::new();
for memory in buffer.drain(..) {
match self.store_inner(&memory) {
Ok(()) => {
succeeded += 1;
tracing::info!(
memory_id = %memory.id.0,
"Retried write succeeded"
);
}
Err(e) => {
tracing::debug!(
memory_id = %memory.id.0,
error = %e,
"Retry write still failing"
);
if still_failing.len() < WRITE_RETRY_BUFFER_CAPACITY {
still_failing.push_back(memory);
}
}
}
}
*buffer = still_failing;
if succeeded > 0 || !buffer.is_empty() {
tracing::info!(
"Write retry drain: {}/{} succeeded, {} still pending",
succeeded,
count,
buffer.len()
);
}
succeeded
}
pub fn pending_retry_count(&self) -> usize {
self.write_retry_buffer.lock().len()
}
pub fn total_write_failures(&self) -> u64 {
self.write_failure_count
.load(std::sync::atomic::Ordering::Relaxed)
}
fn update_indices(&self, memory: &Memory) -> Result<()> {
let idx = self.index_cf();
let mut batch = WriteBatch::default();
let date_key = format!(
"date:{}:{}",
memory.created_at.format("%Y%m%d"),
memory.id.0
);
batch.put_cf(idx, date_key.as_bytes(), b"1");
let type_key = format!(
"type:{:?}:{}",
memory.experience.experience_type, memory.id.0
);
batch.put_cf(idx, type_key.as_bytes(), b"1");
let importance_bucket = (memory.importance() * 10.0) as u32;
let importance_key = format!("importance:{}:{}", importance_bucket, memory.id.0);
batch.put_cf(idx, importance_key.as_bytes(), b"1");
for entity in &memory.experience.entities {
let normalized_entity = entity.to_lowercase();
let entity_key = format!("entity:{}:{}", normalized_entity, memory.id.0);
batch.put_cf(idx, entity_key.as_bytes(), b"1");
}
for tag in &memory.experience.tags {
let normalized_tag = tag.to_lowercase();
let tag_key = format!("tag:{}:{}", normalized_tag, memory.id.0);
batch.put_cf(idx, tag_key.as_bytes(), b"1");
}
if let Some(ctx) = &memory.experience.context {
if let Some(episode_id) = &ctx.episode.episode_id {
let episode_key = format!("episode:{}:{}", episode_id, memory.id.0);
batch.put_cf(idx, episode_key.as_bytes(), b"1");
if let Some(seq) = ctx.episode.sequence_number {
let seq_key = format!("episode_seq:{}:{:010}:{}", episode_id, seq, memory.id.0);
batch.put_cf(idx, seq_key.as_bytes(), b"1");
}
}
}
if let Some(ref robot_id) = memory.experience.robot_id {
let robot_key = format!("robot:{}:{}", robot_id, memory.id.0);
batch.put_cf(idx, robot_key.as_bytes(), b"1");
}
if let Some(ref mission_id) = memory.experience.mission_id {
let mission_key = format!("mission:{}:{}", mission_id, memory.id.0);
batch.put_cf(idx, mission_key.as_bytes(), b"1");
}
if let Some(geo) = memory.experience.geo_location {
let lat = geo[0];
let lon = geo[1];
let geohash = super::types::geohash_encode(lat, lon, 10);
let geo_key = format!("geo:{}:{}", geohash, memory.id.0);
batch.put_cf(idx, geo_key.as_bytes(), b"1");
}
if let Some(ref action_type) = memory.experience.action_type {
let action_key = format!("action:{}:{}", action_type, memory.id.0);
batch.put_cf(idx, action_key.as_bytes(), b"1");
}
if let Some(reward) = memory.experience.reward {
let clamped_reward = reward.clamp(-1.0, 1.0);
let reward_bucket = ((clamped_reward + 1.0) * 10.0) as i32;
let reward_key = format!("reward:{}:{}", reward_bucket, memory.id.0);
batch.put_cf(idx, reward_key.as_bytes(), b"1");
}
{
let content_hash = Self::sha256_content_hash(&memory.experience.content);
let hash_key = format!("content_hash:{}", content_hash);
batch.put_cf(idx, hash_key.as_bytes(), memory.id.0.as_bytes());
}
if let Some(ref external_id) = memory.external_id {
let external_key = format!("external:{}:{}", external_id, memory.id.0);
batch.put_cf(idx, external_key.as_bytes(), memory.id.0.as_bytes());
}
if let Some(ref parent_id) = memory.parent_id {
let parent_key = format!("parent:{}:{}", parent_id.0, memory.id.0);
batch.put_cf(idx, parent_key.as_bytes(), b"1");
}
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db.write_opt(batch, &write_opts)?;
Ok(())
}
fn sha256_content_hash(content: &str) -> String {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
hasher.update(content.as_bytes());
hex::encode(hasher.finalize())
}
pub fn get_by_content_hash(&self, content: &str) -> Option<MemoryId> {
let content_hash = Self::sha256_content_hash(content);
let hash_key = format!("content_hash:{}", content_hash);
let idx = self.index_cf();
match self.db.get_cf(idx, hash_key.as_bytes()) {
Ok(Some(value)) if value.len() == 16 => {
let uuid = uuid::Uuid::from_slice(&value).ok()?;
let memory_id = MemoryId(uuid);
match self.get(&memory_id) {
Ok(_) => Some(memory_id),
Err(_) => {
let _ = self.db.delete_cf(idx, hash_key.as_bytes());
None
}
}
}
_ => None,
}
}
pub fn get(&self, id: &MemoryId) -> Result<Memory> {
let key = id.0.as_bytes();
match self.db.get(key)? {
Some(value) => {
let (memory, needs_migration) = deserialize_memory(&value).with_context(|| {
format!(
"Failed to deserialize memory {} ({} bytes)",
id.0,
value.len()
)
})?;
if needs_migration {
if let Err(e) = self.migrate_memory_format(&memory) {
tracing::debug!("Lazy migration skipped for memory {}: {}", memory.id.0, e);
}
}
Ok(memory)
}
None => Err(anyhow!("Memory not found: {id:?}")),
}
}
fn migrate_memory_format(&self, memory: &Memory) -> Result<()> {
let key = memory.id.0.as_bytes();
let value = crate::serialization::encode_sho(memory)
.context("Failed to serialize for migration")?;
let mut write_opts = WriteOptions::default();
write_opts.set_sync(false);
self.db.put_opt(key, &value, &write_opts)?;
tracing::debug!("Migrated memory {} to current format", memory.id.0);
Ok(())
}
pub fn find_by_external_id(&self, external_id: &str) -> Result<Option<Memory>> {
let prefix = format!("external:{external_id}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _value) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
return Ok(Some(self.get(&MemoryId(uuid))?));
}
}
}
Ok(None)
}
pub fn update(&self, memory: &Memory) -> Result<()> {
self.remove_from_indices(&memory.id)?;
self.store(memory)
}
#[allow(unused)] pub fn delete(&self, id: &MemoryId) -> Result<()> {
self.remove_from_indices(id)?;
let key = id.0.as_bytes();
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db.delete_opt(key, &write_opts)?;
let mapping_key = format!("vmapping:{}", id.0);
let _ = self.db.delete_opt(mapping_key.as_bytes(), &write_opts);
Ok(())
}
fn remove_from_indices(&self, id: &MemoryId) -> Result<()> {
let memory = match self.get(id) {
Ok(m) => m,
Err(_) => {
tracing::debug!("Memory {} not found, skipping index cleanup", id.0);
return Ok(());
}
};
let idx = self.index_cf();
let mut batch = WriteBatch::default();
let date_key = format!("date:{}:{}", memory.created_at.format("%Y%m%d"), id.0);
batch.delete_cf(idx, date_key.as_bytes());
let type_key = format!("type:{:?}:{}", memory.experience.experience_type, id.0);
batch.delete_cf(idx, type_key.as_bytes());
let importance_bucket = (memory.importance() * 10.0) as u32;
let importance_key = format!("importance:{}:{}", importance_bucket, id.0);
batch.delete_cf(idx, importance_key.as_bytes());
for entity in &memory.experience.entities {
let normalized_entity = entity.to_lowercase();
let entity_key = format!("entity:{}:{}", normalized_entity, id.0);
batch.delete_cf(idx, entity_key.as_bytes());
}
for tag in &memory.experience.tags {
let normalized_tag = tag.to_lowercase();
let tag_key = format!("tag:{}:{}", normalized_tag, id.0);
batch.delete_cf(idx, tag_key.as_bytes());
}
if let Some(ctx) = &memory.experience.context {
if let Some(episode_id) = &ctx.episode.episode_id {
let episode_key = format!("episode:{}:{}", episode_id, id.0);
batch.delete_cf(idx, episode_key.as_bytes());
if let Some(seq) = ctx.episode.sequence_number {
let seq_key = format!("episode_seq:{}:{:010}:{}", episode_id, seq, id.0);
batch.delete_cf(idx, seq_key.as_bytes());
}
}
}
if let Some(ref robot_id) = memory.experience.robot_id {
let robot_key = format!("robot:{}:{}", robot_id, id.0);
batch.delete_cf(idx, robot_key.as_bytes());
}
if let Some(ref mission_id) = memory.experience.mission_id {
let mission_key = format!("mission:{}:{}", mission_id, id.0);
batch.delete_cf(idx, mission_key.as_bytes());
}
if let Some(geo) = memory.experience.geo_location {
let geohash = super::types::geohash_encode(geo[0], geo[1], 10);
let geo_key = format!("geo:{}:{}", geohash, id.0);
batch.delete_cf(idx, geo_key.as_bytes());
}
if let Some(ref action_type) = memory.experience.action_type {
let action_key = format!("action:{}:{}", action_type, id.0);
batch.delete_cf(idx, action_key.as_bytes());
}
if let Some(reward) = memory.experience.reward {
let clamped_reward = reward.clamp(-1.0, 1.0);
let reward_bucket = ((clamped_reward + 1.0) * 10.0) as i32;
let reward_key = format!("reward:{}:{}", reward_bucket, id.0);
batch.delete_cf(idx, reward_key.as_bytes());
}
{
let content_hash = Self::sha256_content_hash(&memory.experience.content);
let hash_key = format!("content_hash:{}", content_hash);
batch.delete_cf(idx, hash_key.as_bytes());
}
if let Some(ref external_id) = memory.external_id {
let external_key = format!("external:{}:{}", external_id, id.0);
batch.delete_cf(idx, external_key.as_bytes());
}
if let Some(ref parent_id) = memory.parent_id {
let parent_key = format!("parent:{}:{}", parent_id.0, id.0);
batch.delete_cf(idx, parent_key.as_bytes());
}
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db.write_opt(batch, &write_opts)?;
Ok(())
}
pub fn search(&self, criteria: SearchCriteria) -> Result<Vec<Memory>> {
let mut memory_ids = Vec::new();
match criteria {
SearchCriteria::ByDate { start, end } => {
memory_ids = self.search_by_date_range(start, end)?;
}
SearchCriteria::ByType(exp_type) => {
memory_ids = self.search_by_type(exp_type)?;
}
SearchCriteria::ByImportance { min, max } => {
memory_ids = self.search_by_importance(min, max)?;
}
SearchCriteria::ByEntity(entity) => {
memory_ids = self.search_by_entity(&entity)?;
}
SearchCriteria::ByTags(tags) => {
memory_ids = self.search_by_tags(&tags)?;
}
SearchCriteria::ByEpisode(episode_id) => {
memory_ids = self.search_by_episode(&episode_id)?;
}
SearchCriteria::ByEpisodeSequence {
episode_id,
min_sequence,
max_sequence,
} => {
memory_ids =
self.search_by_episode_sequence(&episode_id, min_sequence, max_sequence)?;
}
SearchCriteria::ByRobot(robot_id) => {
memory_ids = self.search_by_robot(&robot_id)?;
}
SearchCriteria::ByMission(mission_id) => {
memory_ids = self.search_by_mission(&mission_id)?;
}
SearchCriteria::ByLocation {
lat,
lon,
radius_meters,
} => {
memory_ids = self.search_by_location(lat, lon, radius_meters)?;
}
SearchCriteria::ByActionType(action_type) => {
memory_ids = self.search_by_action_type(&action_type)?;
}
SearchCriteria::ByReward { min, max } => {
memory_ids = self.search_by_reward(min, max)?;
}
SearchCriteria::Combined(criterias) => {
use std::collections::HashSet;
let mut result_sets: Vec<HashSet<MemoryId>> = Vec::new();
for c in criterias {
result_sets.push(
self.search(c)?
.into_iter()
.map(|m| m.id)
.collect::<HashSet<_>>(),
);
}
if !result_sets.is_empty() {
let first_set = result_sets.remove(0);
memory_ids = first_set
.into_iter()
.filter(|id| result_sets.iter().all(|set| set.contains(id)))
.collect();
}
}
SearchCriteria::ByParent(parent_id) => {
memory_ids = self.search_by_parent(&parent_id)?;
}
SearchCriteria::RootsOnly => {
memory_ids = self.search_roots()?;
}
}
let mut memories = Vec::new();
for id in memory_ids {
if let Ok(memory) = self.get(&id) {
if !memory.is_forgotten() {
memories.push(memory);
}
}
}
Ok(memories)
}
fn search_by_date_range(
&self,
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let start_key = format!("date:{}", start.format("%Y%m%d"));
let end_key = format!("date:{}~", end.format("%Y%m%d"));
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(start_key.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _value) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if &*key_str > end_key.as_str() {
break;
}
if key_str.starts_with("date:") {
let parts: Vec<&str> = key_str.split(':').collect();
if parts.len() >= 3 {
if let Ok(uuid) = uuid::Uuid::parse_str(parts[2]) {
ids.push(MemoryId(uuid));
}
}
}
}
Ok(ids)
}
fn search_by_type(&self, exp_type: ExperienceType) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let prefix = format!("type:{exp_type:?}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
ids.push(MemoryId(uuid));
}
}
}
Ok(ids)
}
fn search_by_importance(&self, min: f32, max: f32) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let min_bucket = (min * 10.0) as u32;
let max_bucket = (max * 10.0) as u32;
for bucket in min_bucket..=max_bucket {
let prefix = format!("importance:{bucket}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
ids.push(MemoryId(uuid));
}
}
}
}
Ok(ids)
}
fn search_by_entity(&self, entity: &str) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let normalized_entity = entity.to_lowercase();
let prefix = format!("entity:{normalized_entity}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
ids.push(MemoryId(uuid));
}
}
}
Ok(ids)
}
fn search_by_tags(&self, tags: &[String]) -> Result<Vec<MemoryId>> {
use std::collections::HashSet;
let mut all_ids = HashSet::new();
for tag in tags {
let normalized_tag = tag.to_lowercase();
let prefix = format!("tag:{normalized_tag}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
all_ids.insert(MemoryId(uuid));
}
}
}
}
Ok(all_ids.into_iter().collect())
}
fn search_by_episode(&self, episode_id: &str) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let prefix = format!("episode:{episode_id}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
ids.push(MemoryId(uuid));
}
}
}
Ok(ids)
}
fn search_by_episode_sequence(
&self,
episode_id: &str,
min_sequence: Option<u32>,
max_sequence: Option<u32>,
) -> Result<Vec<MemoryId>> {
let mut results: Vec<(u32, MemoryId)> = Vec::new();
let prefix = format!("episode_seq:{episode_id}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(rest) = key_str.strip_prefix(&prefix) {
let parts: Vec<&str> = rest.splitn(2, ':').collect();
if parts.len() == 2 {
if let (Ok(seq), Ok(uuid)) =
(parts[0].parse::<u32>(), uuid::Uuid::parse_str(parts[1]))
{
let passes_min = min_sequence.map_or(true, |min| seq >= min);
let passes_max = max_sequence.map_or(true, |max| seq <= max);
if passes_min && passes_max {
results.push((seq, MemoryId(uuid)));
}
}
}
}
}
results.sort_by_key(|(seq, _)| *seq);
Ok(results.into_iter().map(|(_, id)| id).collect())
}
fn search_by_robot(&self, robot_id: &str) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let prefix = format!("robot:{robot_id}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
ids.push(MemoryId(uuid));
}
}
}
Ok(ids)
}
fn search_by_mission(&self, mission_id: &str) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let prefix = format!("mission:{mission_id}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
ids.push(MemoryId(uuid));
}
}
}
Ok(ids)
}
fn search_by_location(
&self,
center_lat: f64,
center_lon: f64,
radius_meters: f64,
) -> Result<Vec<MemoryId>> {
use super::types::{geohash_decode, geohash_search_prefixes, GeoFilter};
let geo_filter = GeoFilter::new(center_lat, center_lon, radius_meters);
let mut ids = Vec::new();
let prefixes = geohash_search_prefixes(center_lat, center_lon, radius_meters);
for geohash_prefix in prefixes {
let prefix = format!("geo:{}", geohash_prefix);
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _value) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
let parts: Vec<&str> = key_str.split(':').collect();
if parts.len() >= 3 {
let geohash = parts[1];
let (min_lat, min_lon, max_lat, max_lon) = geohash_decode(geohash);
let approx_lat = (min_lat + max_lat) / 2.0;
let approx_lon = (min_lon + max_lon) / 2.0;
if geo_filter.contains(approx_lat, approx_lon) {
if let Ok(uuid) = uuid::Uuid::parse_str(parts[2]) {
ids.push(MemoryId(uuid));
}
}
}
}
}
Ok(ids)
}
fn search_by_action_type(&self, action_type: &str) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let prefix = format!("action:{action_type}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
ids.push(MemoryId(uuid));
}
}
}
Ok(ids)
}
fn search_by_reward(&self, min: f32, max: f32) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let clamped_min = min.clamp(-1.0, 1.0);
let clamped_max = max.clamp(-1.0, 1.0);
let min_bucket = ((clamped_min + 1.0) * 10.0) as i32; let max_bucket = ((clamped_max + 1.0) * 10.0) as i32;
for bucket in min_bucket..=max_bucket {
let prefix = format!("reward:{bucket}:");
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
ids.push(MemoryId(uuid));
}
}
}
}
Ok(ids)
}
fn search_by_parent(&self, parent_id: &MemoryId) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let prefix = format!("parent:{}:", parent_id.0);
let iter = self.db.iterator_cf(
self.index_cf(),
IteratorMode::From(prefix.as_bytes(), rocksdb::Direction::Forward),
);
for (key, _) in iter.log_errors() {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with(&prefix) {
break;
}
if let Some(child_id_str) = key_str.strip_prefix(&prefix) {
if let Ok(uuid) = uuid::Uuid::parse_str(child_id_str) {
ids.push(MemoryId(uuid));
}
}
}
Ok(ids)
}
fn search_roots(&self) -> Result<Vec<MemoryId>> {
let mut roots = Vec::new();
let iter = self.db.iterator(IteratorMode::Start);
for item in iter {
if let Ok((key, value)) = item {
if key.len() != 16 {
continue;
}
if let Ok((memory, _)) = deserialize_memory(&value) {
if memory.parent_id.is_none() {
roots.push(memory.id);
}
}
}
}
Ok(roots)
}
pub fn get_children(&self, parent_id: &MemoryId) -> Result<Vec<Memory>> {
let child_ids = self.search_by_parent(parent_id)?;
let mut children = Vec::new();
for id in child_ids {
if let Ok(memory) = self.get(&id) {
children.push(memory);
}
}
Ok(children)
}
pub fn get_ancestors(&self, memory_id: &MemoryId) -> Result<Vec<Memory>> {
let mut ancestors = Vec::new();
let mut current_id = memory_id.clone();
for _ in 0..100 {
let memory = self.get(¤t_id)?;
if let Some(parent_id) = &memory.parent_id {
let parent = self.get(parent_id)?;
ancestors.push(parent.clone());
current_id = parent_id.clone();
} else {
break; }
}
Ok(ancestors)
}
pub fn get_hierarchy_context(
&self,
memory_id: &MemoryId,
) -> Result<(Vec<Memory>, Memory, Vec<Memory>)> {
let memory = self.get(memory_id)?;
let ancestors = self.get_ancestors(memory_id)?;
let children = self.get_children(memory_id)?;
Ok((ancestors, memory, children))
}
pub fn get_subtree(&self, root_id: &MemoryId, max_depth: usize) -> Result<Vec<Memory>> {
let mut result = Vec::new();
let mut queue = vec![(root_id.clone(), 0usize)];
while let Some((id, depth)) = queue.pop() {
if depth > max_depth {
continue;
}
if let Ok(memory) = self.get(&id) {
result.push(memory);
if depth < max_depth {
let child_ids = self.search_by_parent(&id)?;
for child_id in child_ids {
queue.push((child_id, depth + 1));
}
}
}
}
Ok(result)
}
pub fn get_all_ids(&self) -> Result<Vec<MemoryId>> {
let mut ids = Vec::new();
let mut read_opts = rocksdb::ReadOptions::default();
read_opts.fill_cache(false);
let iter = self.db.iterator_opt(IteratorMode::Start, read_opts);
for item in iter {
if let Ok((key, _)) = item {
if key.len() == 16 {
let uuid_bytes: [u8; 16] = key[..16].try_into().unwrap();
ids.push(MemoryId(uuid::Uuid::from_bytes(uuid_bytes)));
}
}
}
Ok(ids)
}
pub fn get_all(&self) -> Result<Vec<Memory>> {
let mut memories = Vec::new();
let mut read_opts = rocksdb::ReadOptions::default();
read_opts.fill_cache(false);
let iter = self.db.iterator_opt(IteratorMode::Start, read_opts);
for item in iter {
if let Ok((key, value)) = item {
if key.len() != 16 {
continue;
}
if let Ok((memory, _)) = deserialize_memory(&value) {
if !memory.is_forgotten() {
memories.push(memory);
}
}
}
}
Ok(memories)
}
pub fn get_uncompressed_older_than(&self, cutoff: DateTime<Utc>) -> Result<Vec<Memory>> {
let mut memories = Vec::new();
let iter = self.db.iterator(IteratorMode::Start);
for item in iter {
if let Ok((key, value)) = item {
if key.len() != 16 {
continue;
}
if let Ok((memory, _)) = deserialize_memory(&value) {
if !memory.compressed && !memory.is_forgotten() && memory.created_at < cutoff {
memories.push(memory);
}
}
}
}
Ok(memories)
}
pub fn mark_forgotten_by_age(&self, cutoff: DateTime<Utc>) -> Result<Vec<MemoryId>> {
let mut batch = rocksdb::WriteBatch::default();
let mut flagged_ids = Vec::new();
let now = Utc::now().to_rfc3339();
let iter = self.db.iterator(IteratorMode::Start);
for item in iter {
if let Ok((key, value)) = item {
if key.len() != 16 {
continue;
}
if let Ok((mut memory, _)) = deserialize_memory(&value) {
if memory.is_forgotten() {
continue;
}
if memory.created_at < cutoff {
flagged_ids.push(memory.id.clone());
memory
.experience
.metadata
.insert("forgotten".to_string(), "true".to_string());
memory
.experience
.metadata
.insert("forgotten_at".to_string(), now.clone());
let updated_value = crate::serialization::encode_sho(&memory)?;
batch.put(&key, updated_value);
}
}
}
}
if !flagged_ids.is_empty() {
let mut write_opts = WriteOptions::default();
write_opts.set_sync(true);
self.db.write_opt(batch, &write_opts)?;
}
Ok(flagged_ids)
}
pub fn mark_forgotten_by_importance(&self, threshold: f32) -> Result<Vec<MemoryId>> {
let mut batch = rocksdb::WriteBatch::default();
let mut flagged_ids = Vec::new();
let now = Utc::now().to_rfc3339();
let iter = self.db.iterator(IteratorMode::Start);
for item in iter {
if let Ok((key, value)) = item {
if key.len() != 16 {
continue;
}
if let Ok((mut memory, _)) = deserialize_memory(&value) {
if memory.is_forgotten() {
continue;
}
if memory.importance() < threshold {
flagged_ids.push(memory.id.clone());
memory
.experience
.metadata
.insert("forgotten".to_string(), "true".to_string());
memory
.experience
.metadata
.insert("forgotten_at".to_string(), now.clone());
let updated_value = crate::serialization::encode_sho(&memory)?;
batch.put(&key, updated_value);
}
}
}
}
if !flagged_ids.is_empty() {
let mut write_opts = WriteOptions::default();
write_opts.set_sync(true);
self.db.write_opt(batch, &write_opts)?;
}
Ok(flagged_ids)
}
pub fn remove_matching(&self, regex: ®ex::Regex) -> Result<usize> {
let mut count = 0;
let mut to_delete: Vec<MemoryId> = Vec::new();
let iter = self.db.iterator(IteratorMode::Start);
for item in iter {
if let Ok((key, value)) = item {
if key.len() != 16 {
continue;
}
if let Ok((memory, _)) = deserialize_memory(&value) {
if regex.is_match(&memory.experience.content) {
to_delete.push(memory.id);
count += 1;
}
}
}
}
for memory_id in to_delete {
if let Err(e) = self.delete(&memory_id) {
tracing::warn!("Failed to delete matching memory {}: {}", memory_id.0, e);
}
}
Ok(count)
}
pub fn update_access(&self, id: &MemoryId) -> Result<()> {
if let Ok(memory) = self.get(id) {
memory.update_access();
self.update(&memory)?;
}
Ok(())
}
pub fn get_stats(&self) -> Result<StorageStats> {
let mut stats = StorageStats::default();
let mut raw_count = 0;
let mut skipped_non_memory = 0;
let mut deserialize_errors = 0;
let stats_prefix = b"stats:";
let iter = self.db.iterator(IteratorMode::Start);
for item in iter {
match item {
Ok((key, value)) => {
raw_count += 1;
if key.starts_with(stats_prefix) {
skipped_non_memory += 1;
continue;
}
if key.len() != 16 {
skipped_non_memory += 1;
continue;
}
match deserialize_memory(&value) {
Ok((memory, _)) => {
if memory.is_forgotten() {
continue;
}
stats.total_count += 1;
stats.total_size_bytes += value.len();
if memory.compressed {
stats.compressed_count += 1;
}
stats.importance_sum += memory.importance();
}
Err(e) => {
deserialize_errors += 1;
tracing::warn!(
"Corrupted memory entry (key len: {}, value len: {}): {}",
key.len(),
value.len(),
e
);
}
}
}
Err(e) => {
tracing::error!("Iterator error: {}", e);
}
}
}
tracing::debug!(
"get_stats: raw_count={}, memories={}, skipped={}, corrupted={}",
raw_count,
stats.total_count,
skipped_non_memory,
deserialize_errors
);
if stats.total_count > 0 {
stats.average_importance = stats.importance_sum / stats.total_count as f32;
}
stats.total_retrievals = self.get_retrieval_count().unwrap_or(0);
Ok(stats)
}
pub fn get_retrieval_count(&self) -> Result<usize> {
const RETRIEVAL_KEY: &[u8] = b"stats:total_retrievals";
match self.db.get(RETRIEVAL_KEY)? {
Some(data) => {
if data.len() >= 8 {
Ok(usize::from_le_bytes(data[..8].try_into().unwrap_or([0; 8])))
} else {
Ok(0)
}
}
None => Ok(0),
}
}
pub fn increment_retrieval_count(&self) -> Result<usize> {
const RETRIEVAL_KEY: &[u8] = b"stats:total_retrievals";
let current = self.get_retrieval_count().unwrap_or(0);
let new_count = current + 1;
self.db.put(RETRIEVAL_KEY, new_count.to_le_bytes())?;
Ok(new_count)
}
pub fn cleanup_corrupted(&self) -> Result<usize> {
let mut to_delete = Vec::new();
let skip_prefixes: &[&[u8]] = &[
b"stats:",
b"vmapping:",
b"interference:",
b"interference_meta:",
b"_watermark:",
b"facts:",
b"facts_by_entity:",
b"facts_by_type:",
b"facts_embedding:",
b"temporal_facts:",
b"temporal_by_time:",
b"temporal_by_entity:",
b"lineage:",
b"learning:",
b"geo:",
];
let iter = self.db.iterator(IteratorMode::Start);
for item in iter {
if let Ok((key, value)) = item {
if skip_prefixes.iter().any(|p| key.starts_with(p)) {
continue;
}
let is_valid_memory_key = key.len() == 16;
if !is_valid_memory_key {
tracing::debug!(
"Marking for deletion: invalid key length {} (expected 16)",
key.len()
);
to_delete.push(key.to_vec());
} else if deserialize_memory(&value).is_err() {
tracing::debug!(
"Marking for deletion: valid key but corrupted value ({} bytes)",
value.len()
);
to_delete.push(key.to_vec());
}
}
}
let count = to_delete.len();
if count > 0 {
tracing::info!("Cleaning up {} corrupted memory entries", count);
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
for key in to_delete {
if let Err(e) = self.db.delete_opt(&key, &write_opts) {
tracing::warn!("Failed to delete corrupted entry: {}", e);
}
}
self.flush()?;
}
Ok(count)
}
pub fn migrate_legacy(&self) -> Result<(usize, usize, usize)> {
let mut migrated = 0;
let mut already_current = 0;
let mut failed = 0;
let stats_prefix = b"stats:";
let iter = self.db.iterator(IteratorMode::Start);
let mut to_migrate = Vec::new();
for item in iter {
if let Ok((key, value)) = item {
if key.starts_with(stats_prefix) {
continue;
}
if key.len() != 16 {
continue;
}
let is_current = deserialize_memory(&value).is_ok();
if is_current {
already_current += 1;
continue;
}
match deserialize_memory(&value) {
Ok((memory, _)) => {
to_migrate.push((key.to_vec(), memory));
}
Err(_) => {
failed += 1;
}
}
}
}
if !to_migrate.is_empty() {
tracing::info!(
"Migrating {} legacy memories to current format",
to_migrate.len()
);
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
for (key, memory) in to_migrate {
match crate::serialization::encode_sho(&memory) {
Ok(serialized) => {
if let Err(e) = self.db.put_opt(&key, &serialized, &write_opts) {
tracing::warn!("Failed to migrate memory: {e}");
failed += 1;
} else {
migrated += 1;
}
}
Err(e) => {
tracing::warn!("Failed to serialize migrated memory: {e}");
failed += 1;
}
}
}
self.flush()?;
}
tracing::info!(
"Migration complete: {} migrated, {} already current, {} failed",
migrated,
already_current,
failed
);
Ok((migrated, already_current, failed))
}
pub fn flush(&self) -> Result<()> {
use rocksdb::FlushOptions;
let mut flush_opts = FlushOptions::default();
flush_opts.set_wait(true);
self.db
.flush_opt(&flush_opts)
.map_err(|e| anyhow::anyhow!("Failed to flush memory storage: {e}"))?;
self.db
.flush_cf_opt(self.index_cf(), &flush_opts)
.map_err(|e| anyhow::anyhow!("Failed to flush index CF: {e}"))?;
Ok(())
}
pub fn db(&self) -> Arc<DB> {
self.db.clone()
}
}
#[derive(Debug, Clone)]
pub enum SearchCriteria {
ByDate {
start: DateTime<Utc>,
end: DateTime<Utc>,
},
ByType(ExperienceType),
ByImportance {
min: f32,
max: f32,
},
ByEntity(String),
ByTags(Vec<String>),
ByEpisode(String),
ByEpisodeSequence {
episode_id: String,
min_sequence: Option<u32>,
max_sequence: Option<u32>,
},
ByRobot(String),
ByMission(String),
ByLocation {
lat: f64,
lon: f64,
radius_meters: f64,
},
ByActionType(String),
ByReward {
min: f32,
max: f32,
},
Combined(Vec<SearchCriteria>),
ByParent(MemoryId),
RootsOnly,
}
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct StorageStats {
pub total_count: usize,
pub compressed_count: usize,
pub total_size_bytes: usize,
pub average_importance: f32,
pub importance_sum: f32,
#[serde(default)]
pub total_retrievals: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Modality {
Text,
Image,
Audio,
Video,
Unified,
}
impl Modality {
pub fn dimension(&self) -> usize {
match self {
Modality::Text => 384, Modality::Image => 1024,
Modality::Audio => 1024,
Modality::Video => 1024,
Modality::Unified => 1024,
}
}
pub fn as_str(&self) -> &'static str {
match self {
Modality::Text => "text",
Modality::Image => "image",
Modality::Audio => "audio",
Modality::Video => "video",
Modality::Unified => "unified",
}
}
}
impl std::fmt::Display for Modality {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ModalityVectors {
pub vector_ids: Vec<u32>,
pub dimension: usize,
pub chunk_ranges: Option<Vec<(usize, usize)>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorMappingEntry {
pub modalities: HashMap<Modality, ModalityVectors>,
pub created_at: i64,
pub version: u8,
}
impl Default for VectorMappingEntry {
fn default() -> Self {
Self {
modalities: HashMap::new(),
created_at: chrono::Utc::now().timestamp_millis(),
version: 1,
}
}
}
impl VectorMappingEntry {
pub fn with_text(vector_ids: Vec<u32>) -> Self {
let mut modalities = HashMap::new();
modalities.insert(
Modality::Text,
ModalityVectors {
vector_ids,
dimension: 384,
chunk_ranges: None,
},
);
Self {
modalities,
created_at: chrono::Utc::now().timestamp_millis(),
version: 1,
}
}
pub fn text_vectors(&self) -> Option<&Vec<u32>> {
self.modalities.get(&Modality::Text).map(|m| &m.vector_ids)
}
pub fn all_vector_ids(&self) -> Vec<(Modality, u32)> {
self.modalities
.iter()
.flat_map(|(modality, mv)| mv.vector_ids.iter().map(|id| (*modality, *id)))
.collect()
}
pub fn is_empty(&self) -> bool {
self.modalities.values().all(|mv| mv.vector_ids.is_empty())
}
pub fn add_modality(&mut self, modality: Modality, vector_ids: Vec<u32>) {
self.modalities.insert(
modality,
ModalityVectors {
dimension: modality.dimension(),
vector_ids,
chunk_ranges: None,
},
);
}
#[allow(dead_code)]
pub fn with_image(mut self, vector_ids: Vec<u32>) -> Self {
self.add_modality(Modality::Image, vector_ids);
self
}
#[allow(dead_code)]
pub fn with_audio(mut self, vector_ids: Vec<u32>) -> Self {
self.add_modality(Modality::Audio, vector_ids);
self
}
#[allow(dead_code)]
pub fn with_video(mut self, vector_ids: Vec<u32>) -> Self {
self.add_modality(Modality::Video, vector_ids);
self
}
}
impl MemoryStorage {
pub fn store_with_vectors(&self, memory: &Memory, vector_ids: Vec<u32>) -> Result<()> {
self.store_with_multimodal_vectors(memory, Modality::Text, vector_ids)
}
pub fn store_with_multimodal_vectors(
&self,
memory: &Memory,
modality: Modality,
vector_ids: Vec<u32>,
) -> Result<()> {
let mut batch = WriteBatch::default();
let memory_key = memory.id.0.as_bytes();
let memory_value = crate::serialization::encode_sho(memory)
.context(format!("Failed to serialize memory {}", memory.id.0))?;
batch.put(memory_key, &memory_value);
let mapping_key = format!("vmapping:{}", memory.id.0);
let mut mapping_entry = self.get_vector_mapping(&memory.id)?.unwrap_or_default();
mapping_entry.add_modality(modality, vector_ids);
let mapping_value = crate::serialization::encode(&mapping_entry)
.context("Failed to serialize vector mapping")?;
batch.put(mapping_key.as_bytes(), &mapping_value);
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db
.write_opt(batch, &write_opts)
.context("Atomic write of memory + vector mapping failed")?;
if let Err(e) = self.update_indices(memory) {
tracing::warn!("Secondary index update failed (non-fatal): {}", e);
}
Ok(())
}
pub fn get_vector_mapping(&self, memory_id: &MemoryId) -> Result<Option<VectorMappingEntry>> {
let mapping_key = format!("vmapping:{}", memory_id.0);
match self.db.get(mapping_key.as_bytes())? {
Some(data) => {
let (entry, _) = crate::serialization::try_decode::<VectorMappingEntry>(&data)
.context("Failed to deserialize vector mapping")?;
Ok(Some(entry))
}
None => Ok(None),
}
}
pub fn get_all_vector_mappings(&self) -> Result<Vec<(MemoryId, VectorMappingEntry)>> {
let mut mappings = Vec::new();
let prefix = b"vmapping:";
let iter = self
.db
.iterator(IteratorMode::From(prefix, rocksdb::Direction::Forward));
for item in iter {
match item {
Ok((key, value)) => {
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with("vmapping:") {
break;
}
if let Some(id_str) = key_str.strip_prefix("vmapping:") {
if let Ok(uuid) = uuid::Uuid::parse_str(id_str) {
if let Ok((entry, _)) =
crate::serialization::try_decode::<VectorMappingEntry>(&value)
{
mappings.push((MemoryId(uuid), entry));
}
}
}
}
Err(e) => {
tracing::warn!("Error reading vector mapping: {}", e);
}
}
}
Ok(mappings)
}
pub fn delete_vector_mapping(&self, memory_id: &MemoryId) -> Result<()> {
let mapping_key = format!("vmapping:{}", memory_id.0);
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db.delete_opt(mapping_key.as_bytes(), &write_opts)?;
Ok(())
}
pub fn update_vector_mapping(&self, memory_id: &MemoryId, vector_ids: Vec<u32>) -> Result<()> {
self.update_modality_vectors(memory_id, Modality::Text, vector_ids)
}
pub fn update_modality_vectors(
&self,
memory_id: &MemoryId,
modality: Modality,
vector_ids: Vec<u32>,
) -> Result<()> {
let mapping_key = format!("vmapping:{}", memory_id.0);
let mut mapping_entry = self.get_vector_mapping(memory_id)?.unwrap_or_default();
mapping_entry.add_modality(modality, vector_ids);
let mapping_value = crate::serialization::encode(&mapping_entry)?;
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db
.put_opt(mapping_key.as_bytes(), &mapping_value, &write_opts)?;
Ok(())
}
pub fn delete_with_vectors(&self, id: &MemoryId) -> Result<()> {
let mut batch = WriteBatch::default();
batch.delete(id.0.as_bytes());
let mapping_key = format!("vmapping:{}", id.0);
batch.delete(mapping_key.as_bytes());
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db.write_opt(batch, &write_opts)?;
if let Err(e) = self.remove_from_indices(id) {
tracing::warn!("Index cleanup failed (non-fatal): {}", e);
}
Ok(())
}
pub fn count_vector_mappings(&self) -> usize {
let prefix = b"vmapping:";
let iter = self
.db
.iterator(IteratorMode::From(prefix, rocksdb::Direction::Forward));
let mut count = 0;
for item in iter {
if let Ok((key, _)) = item {
if key.starts_with(prefix) {
count += 1;
} else {
break;
}
}
}
count
}
pub fn find_memories_without_mappings(&self) -> Result<Vec<MemoryId>> {
let mut orphans = Vec::new();
let iter = self.db.iterator(IteratorMode::Start);
for item in iter {
if let Ok((key, value)) = item {
if key.len() != 16 {
continue;
}
if let Ok((memory, _)) = deserialize_memory(&value) {
let has_mapping = match self.get_vector_mapping(&memory.id) {
Ok(Some(entry)) => entry.text_vectors().is_some_and(|v| !v.is_empty()),
_ => false,
};
if !has_mapping && memory.experience.embeddings.is_some() {
orphans.push(memory.id);
}
}
}
}
Ok(orphans)
}
pub fn get_all_text_vector_ids(&self) -> Result<Vec<u32>> {
let mut all_ids = Vec::new();
let mappings = self.get_all_vector_mappings()?;
for (_, entry) in mappings {
if let Some(text_vecs) = entry.text_vectors() {
all_ids.extend(text_vecs.iter().copied());
}
}
Ok(all_ids)
}
pub fn get_modality_stats(&self) -> Result<HashMap<Modality, usize>> {
let mut stats: HashMap<Modality, usize> = HashMap::new();
let mappings = self.get_all_vector_mappings()?;
for (_, entry) in mappings {
for (modality, mv) in entry.modalities {
*stats.entry(modality).or_insert(0) += mv.vector_ids.len();
}
}
Ok(stats)
}
pub fn save_interference_records(
&self,
memory_id: &str,
records: &[super::replay::InterferenceRecord],
) -> Result<()> {
let key = format!("interference:{memory_id}");
let value =
serde_json::to_vec(records).context("Failed to serialize interference records")?;
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db
.put_opt(key.as_bytes(), &value, &write_opts)
.context("Failed to persist interference records")?;
Ok(())
}
pub fn load_all_interference_records(
&self,
) -> Result<(
HashMap<String, Vec<super::replay::InterferenceRecord>>,
usize,
)> {
let prefix = b"interference:";
let mut history: HashMap<String, Vec<super::replay::InterferenceRecord>> = HashMap::new();
let mut total_events: usize = 0;
let iter = self
.db
.iterator(IteratorMode::From(prefix, rocksdb::Direction::Forward));
for item in iter.log_errors() {
let (key, value) = item;
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with("interference:") {
break;
}
if let Some(memory_id) = key_str.strip_prefix("interference:") {
match serde_json::from_slice::<Vec<super::replay::InterferenceRecord>>(&value) {
Ok(records) => {
total_events += records.len();
history.insert(memory_id.to_string(), records);
}
Err(e) => {
tracing::warn!(
key = %key_str,
error = %e,
"Failed to deserialize interference records, skipping"
);
}
}
}
}
let persisted_total = self
.db
.get(b"interference_meta:total")
.ok()
.flatten()
.and_then(|v| {
if v.len() == 8 {
Some(u64::from_le_bytes(v[..8].try_into().unwrap()) as usize)
} else {
None
}
})
.unwrap_or(total_events);
Ok((history, persisted_total.max(total_events)))
}
pub fn delete_interference_records(&self, memory_id: &str) -> Result<()> {
let key = format!("interference:{memory_id}");
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db
.delete_opt(key.as_bytes(), &write_opts)
.context("Failed to delete interference records")?;
Ok(())
}
pub fn save_interference_event_count(&self, count: usize) -> Result<()> {
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db
.put_opt(
b"interference_meta:total",
&(count as u64).to_le_bytes(),
&write_opts,
)
.context("Failed to persist interference event count")?;
Ok(())
}
pub fn get_fact_watermark(&self, user_id: &str) -> Option<i64> {
let key = format!("_watermark:fact_extraction:{user_id}");
match self.db.get(key.as_bytes()) {
Ok(Some(bytes)) if bytes.len() == 8 => {
Some(i64::from_le_bytes(bytes[..8].try_into().unwrap()))
}
_ => None,
}
}
pub fn set_fact_watermark(&self, user_id: &str, timestamp_millis: i64) {
let key = format!("_watermark:fact_extraction:{user_id}");
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
if let Err(e) =
self.db
.put_opt(key.as_bytes(), ×tamp_millis.to_le_bytes(), &write_opts)
{
tracing::warn!("Failed to persist fact extraction watermark: {e}");
}
}
pub fn clear_all_interference_records(&self) -> Result<usize> {
let prefix = b"interference";
let mut batch = WriteBatch::default();
let mut count = 0;
let iter = self
.db
.iterator(IteratorMode::From(prefix, rocksdb::Direction::Forward));
for item in iter.log_errors() {
let (key, _) = item;
let key_str = String::from_utf8_lossy(&key);
if !key_str.starts_with("interference") {
break;
}
batch.delete(&key);
count += 1;
}
if count > 0 {
let mut write_opts = WriteOptions::default();
write_opts.set_sync(self.write_mode == WriteMode::Sync);
self.db.write_opt(batch, &write_opts)?;
}
Ok(count)
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde::Serialize;
#[derive(Serialize)]
struct LegacyMinimalFixture {
id: MemoryId,
content: String,
}
fn sample_memory(id: MemoryId, content: &str) -> Memory {
let now = Utc::now();
let experience = Experience {
experience_type: ExperienceType::Observation,
content: content.to_string(),
..Default::default()
};
Memory::from_legacy(
id,
experience,
0.5,
0,
now,
now,
false,
MemoryTier::LongTerm,
Vec::new(),
1.0,
None,
None,
None,
None,
0.0,
None,
None,
1,
Vec::new(),
Vec::new(),
)
}
#[test]
fn test_deserialize_with_fallback_records_current_bincode2_branch() {
let id = MemoryId(uuid::Uuid::new_v4());
let memory = sample_memory(id.clone(), "current format memory");
let bytes = bincode::serde::encode_to_vec(&memory, bincode::config::standard()).unwrap();
let counter =
crate::metrics::LEGACY_FALLBACK_BRANCH_TOTAL.with_label_values(&["bincode2_memory"]);
let before = counter.get();
let (decoded, is_legacy) = deserialize_with_fallback(&bytes).unwrap();
let after = counter.get();
assert_eq!(decoded.id, id);
assert!(!is_legacy);
assert_eq!(after, before);
}
#[test]
fn test_deserialize_with_fallback_bincode1_minimal_fixture() {
let id = MemoryId(uuid::Uuid::new_v4());
let fixture = LegacyMinimalFixture {
id: id.clone(),
content: "legacy bincode1 minimal".to_string(),
};
let bytes = bincode1::serialize(&fixture).unwrap();
let counter =
crate::metrics::LEGACY_FALLBACK_BRANCH_TOTAL.with_label_values(&["bincode1_minimal"]);
let before = counter.get();
let (decoded, is_legacy) = deserialize_with_fallback(&bytes).unwrap();
let after = counter.get();
assert_eq!(decoded.id, id);
assert!(is_legacy);
assert_eq!(after, before + 1);
}
#[test]
fn test_deserialize_with_fallback_msgpack_minimal_fixture() {
let id = MemoryId(uuid::Uuid::new_v4());
let fixture = LegacyMinimalFixture {
id: id.clone(),
content: "legacy msgpack minimal".to_string(),
};
let bytes = rmp_serde::to_vec(&fixture).unwrap();
let counter =
crate::metrics::LEGACY_FALLBACK_BRANCH_TOTAL.with_label_values(&["msgpack_minimal"]);
let before = counter.get();
let (decoded, is_legacy) = deserialize_with_fallback(&bytes).unwrap();
let after = counter.get();
assert_eq!(decoded.id, id);
assert!(is_legacy);
assert_eq!(after, before + 1);
}
#[test]
fn test_write_mode_default_async() {
std::env::remove_var("SHODH_WRITE_MODE");
let mode = WriteMode::default();
assert_eq!(mode, WriteMode::Async);
}
#[test]
fn test_crc32_simple() {
let data = b"test data for CRC32";
let crc1 = crc32_simple(data);
let crc2 = crc32_simple(data);
assert_eq!(crc1, crc2);
assert_ne!(crc1, 0);
let crc3 = crc32_simple(b"different data");
assert_ne!(crc1, crc3);
}
#[test]
fn test_crc32_empty() {
let crc = crc32_simple(b"");
assert_eq!(
crc, 0,
"IEEE CRC32 of empty input is 0 (init 0xFFFFFFFF XOR final 0xFFFFFFFF)"
);
}
#[test]
fn test_modality_dimension() {
assert_eq!(Modality::Text.dimension(), 384);
assert_eq!(Modality::Image.dimension(), 1024);
assert_eq!(Modality::Audio.dimension(), 1024);
assert_eq!(Modality::Video.dimension(), 1024);
assert_eq!(Modality::Unified.dimension(), 1024);
}
#[test]
fn test_modality_as_str() {
assert_eq!(Modality::Text.as_str(), "text");
assert_eq!(Modality::Image.as_str(), "image");
assert_eq!(Modality::Audio.as_str(), "audio");
assert_eq!(Modality::Video.as_str(), "video");
}
#[test]
fn test_vector_mapping_entry_with_text() {
let entry = VectorMappingEntry::with_text(vec![1, 2, 3]);
assert_eq!(entry.text_vectors(), Some(&vec![1, 2, 3]));
assert!(!entry.is_empty());
}
#[test]
fn test_vector_mapping_entry_multimodal() {
let entry = VectorMappingEntry::with_text(vec![1])
.with_image(vec![2])
.with_audio(vec![3])
.with_video(vec![4]);
let all = entry.all_vector_ids();
assert_eq!(all.len(), 4);
assert!(all.contains(&(Modality::Text, 1)));
assert!(all.contains(&(Modality::Image, 2)));
assert!(all.contains(&(Modality::Audio, 3)));
assert!(all.contains(&(Modality::Video, 4)));
}
#[test]
fn test_vector_mapping_entry_empty() {
let entry = VectorMappingEntry::default();
assert!(entry.is_empty());
assert!(entry.text_vectors().is_none());
assert!(entry.all_vector_ids().is_empty());
}
#[test]
fn test_vector_mapping_entry_add_modality() {
let mut entry = VectorMappingEntry::default();
entry.add_modality(Modality::Text, vec![1, 2]);
assert_eq!(entry.text_vectors(), Some(&vec![1, 2]));
}
#[test]
fn test_storage_stats_default() {
let stats = StorageStats::default();
assert_eq!(stats.total_count, 0);
assert_eq!(stats.compressed_count, 0);
assert_eq!(stats.total_size_bytes, 0);
assert_eq!(stats.total_retrievals, 0);
}
#[test]
fn test_search_criteria_variants() {
let criteria1 = SearchCriteria::ByEntity("test".to_string());
let criteria2 = SearchCriteria::ByImportance { min: 0.5, max: 1.0 };
let criteria3 = SearchCriteria::ByType(ExperienceType::Observation);
assert!(matches!(criteria1, SearchCriteria::ByEntity(_)));
assert!(matches!(criteria2, SearchCriteria::ByImportance { .. }));
assert!(matches!(criteria3, SearchCriteria::ByType(_)));
}
#[test]
fn test_search_criteria_by_date() {
let now = Utc::now();
let start = now - chrono::Duration::days(7);
let criteria = SearchCriteria::ByDate { start, end: now };
if let SearchCriteria::ByDate { start: s, end: e } = criteria {
assert!(s < e);
} else {
panic!("Expected ByDate");
}
}
#[test]
fn test_search_criteria_combined() {
let criteria = SearchCriteria::Combined(vec![
SearchCriteria::ByEntity("test".to_string()),
SearchCriteria::ByImportance { min: 0.5, max: 1.0 },
]);
if let SearchCriteria::Combined(inner) = criteria {
assert_eq!(inner.len(), 2);
} else {
panic!("Expected Combined");
}
}
#[test]
fn test_modality_vectors_struct() {
let mv = ModalityVectors {
vector_ids: vec![1, 2, 3],
dimension: 384,
chunk_ranges: None,
};
assert_eq!(mv.vector_ids.len(), 3);
assert_eq!(mv.dimension, 384);
assert!(mv.chunk_ranges.is_none());
}
}