use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use crate::profile::{
BalancedRecallSnapshot, BalancedRecallState, ProfileBinding, ProfileLifecycle, ProfileRecord,
};
use crate::section_state::{SectionPosteriorSnapshot, SectionPosteriorState};
pub fn sort_fallback_candidates(candidates: &mut [&ProfileRecord]) {
candidates.sort_by(|a, b| a.created_at.cmp(&b.created_at).then(a.id.cmp(&b.id)));
}
pub struct BrainState {
pub profiles: HashMap<String, ProfileRecord>,
pub balanced_recall: BalancedRecallState,
pub profile_states: HashMap<String, BalancedRecallState>,
pub bindings: Vec<ProfileBinding>,
pub section_states: HashMap<String, SectionPosteriorState>,
pub router_state: HashMap<String, RouterStateBlob>,
pub adapter_set: HashMap<String, Vec<AdapterRecord>>,
}
impl BrainState {
pub fn new(entity_capacity: usize) -> Self {
let mut profiles = HashMap::new();
let record = ProfileRecord::new_balanced_recall(entity_capacity);
let profile_id = record.id.clone();
profiles.insert(profile_id, record);
Self {
profiles,
balanced_recall: BalancedRecallState::new(entity_capacity),
profile_states: HashMap::new(),
bindings: Vec::new(),
section_states: HashMap::new(),
router_state: HashMap::new(),
adapter_set: HashMap::new(),
}
}
pub fn to_snapshot(&self) -> BrainStateSnapshot {
let extra: HashMap<String, BalancedRecallSnapshot> = self
.profile_states
.iter()
.map(|(id, s)| (id.clone(), s.to_snapshot()))
.collect();
let section_states: HashMap<String, SectionPosteriorSnapshot> = self
.section_states
.iter()
.map(|(id, s)| (id.clone(), s.to_snapshot()))
.collect();
BrainStateSnapshot {
profiles: self.profiles.clone(),
balanced_recall: self.balanced_recall.to_snapshot(),
profile_states: extra,
bindings: self.bindings.clone(),
section_states,
router_state: self.router_state.clone(),
adapter_set: self.adapter_set.clone(),
}
}
pub fn from_snapshot(snapshot: BrainStateSnapshot, entity_capacity: usize) -> Self {
let extra: HashMap<String, BalancedRecallState> = snapshot
.profile_states
.into_iter()
.map(|(id, s)| (id, BalancedRecallState::from_snapshot(s, entity_capacity)))
.collect();
let section_states: HashMap<String, SectionPosteriorState> = snapshot
.section_states
.into_iter()
.map(|(id, s)| (id, SectionPosteriorState::from_snapshot(s)))
.collect();
Self {
profiles: snapshot.profiles,
balanced_recall: BalancedRecallState::from_snapshot(
snapshot.balanced_recall,
entity_capacity,
),
profile_states: extra,
bindings: snapshot.bindings,
section_states,
router_state: snapshot.router_state,
adapter_set: snapshot.adapter_set,
}
}
pub fn reset_posteriors(&mut self) {
self.balanced_recall.reset_posteriors();
if let Some(record) = self.profiles.get_mut("balanced-recall-v1") {
record.exploration_epoch = self.balanced_recall.exploration_epoch;
record.state_snapshot = serde_json::to_value(self.balanced_recall.to_snapshot()).ok();
}
if let Some(ss) = self.section_states.get_mut("balanced-recall-v1") {
ss.reset_posteriors();
}
}
pub fn reset_profile_posteriors(&mut self, profile_id: &str) {
if let Some(ps) = self.profile_states.get_mut(profile_id) {
ps.reset_posteriors();
let snap = serde_json::to_value(ps.to_snapshot()).ok();
let epoch = ps.exploration_epoch;
if let Some(record) = self.profiles.get_mut(profile_id) {
record.exploration_epoch = epoch;
record.state_snapshot = snap;
}
}
if let Some(ss) = self.section_states.get_mut(profile_id) {
ss.reset_posteriors();
}
}
pub fn resolve(
&self,
actor: Option<&str>,
namespace: Option<&str>,
consumer_kind: &str,
) -> Option<&ProfileRecord> {
self.resolve_with_match(actor, namespace, consumer_kind)
.map(|(record, _, _)| record)
}
pub fn resolve_with_match(
&self,
actor: Option<&str>,
namespace: Option<&str>,
consumer_kind: &str,
) -> Option<(&ProfileRecord, String, bool)> {
let actor_val = actor.unwrap_or("*");
let namespace_val = namespace.unwrap_or("*");
let best = self
.bindings
.iter()
.filter(|b| {
(b.actor == "*" || b.actor == actor_val)
&& (b.namespace == "*" || b.namespace == namespace_val)
&& (b.consumer_kind == "*" || b.consumer_kind == consumer_kind)
&& self
.profiles
.get(&b.profile_id)
.is_some_and(|p| p.lifecycle != ProfileLifecycle::Archived)
})
.max_by_key(|b| {
let actor_score = if b.actor != "*" { 4 } else { 0 };
let ns_score = if b.namespace != "*" { 2 } else { 0 };
let kind_score = if b.consumer_kind != "*" { 1 } else { 0 };
(
actor_score + ns_score + kind_score,
b.priority,
-(b.created_at.timestamp()),
)
});
if let Some(binding) = best {
if let Some(record) = self.profiles.get(&binding.profile_id) {
return Some((record, binding.consumer_kind.clone(), true));
}
}
if let Some(default) = self.profiles.get("balanced-recall-v1") {
if default.lifecycle == ProfileLifecycle::Active
&& (default.consumer_kind == consumer_kind
|| consumer_kind == "*"
|| default.consumer_kind == "*")
{
return Some((default, default.consumer_kind.clone(), false));
}
}
let mut candidates: Vec<&ProfileRecord> = self
.profiles
.values()
.filter(|p| p.consumer_kind == consumer_kind && p.lifecycle == ProfileLifecycle::Active)
.collect();
sort_fallback_candidates(&mut candidates);
candidates
.into_iter()
.next()
.map(|p| (p, p.consumer_kind.clone(), false))
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct RouterStateBlob {
pub schema_version: u32,
pub gate_bytes: Vec<u8>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct AdapterRecord {
pub adapter_id: String,
pub slot: u32,
pub content_hash: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BrainStateSnapshot {
pub profiles: HashMap<String, ProfileRecord>,
pub balanced_recall: BalancedRecallSnapshot,
#[serde(default)]
pub profile_states: HashMap<String, BalancedRecallSnapshot>,
pub bindings: Vec<ProfileBinding>,
#[serde(default)]
pub section_states: HashMap<String, SectionPosteriorSnapshot>,
#[serde(default)]
pub router_state: HashMap<String, RouterStateBlob>,
#[serde(default)]
pub adapter_set: HashMap<String, Vec<AdapterRecord>>,
}
pub fn validate_brain_state_snapshot(snapshot: &BrainStateSnapshot) -> Result<(), String> {
let br = &snapshot.balanced_recall;
br.relevance
.validate()
.map_err(|e| format!("balanced_recall.relevance: {e}"))?;
br.salience
.validate()
.map_err(|e| format!("balanced_recall.salience: {e}"))?;
br.temporal
.validate()
.map_err(|e| format!("balanced_recall.temporal: {e}"))?;
for (id, p) in &br.entity_posteriors {
p.validate()
.map_err(|e| format!("balanced_recall.entity_posteriors[{id}]: {e}"))?;
}
for (pid, ps) in &snapshot.profile_states {
ps.relevance
.validate()
.map_err(|e| format!("profile_states[{pid}].relevance: {e}"))?;
ps.salience
.validate()
.map_err(|e| format!("profile_states[{pid}].salience: {e}"))?;
ps.temporal
.validate()
.map_err(|e| format!("profile_states[{pid}].temporal: {e}"))?;
for (id, p) in &ps.entity_posteriors {
p.validate()
.map_err(|e| format!("profile_states[{pid}].entity_posteriors[{id}]: {e}"))?;
}
}
for (pid, ss) in &snapshot.section_states {
for (st, p) in &ss.posteriors {
p.validate()
.map_err(|e| format!("section_states[{pid}].posteriors[{st:?}]: {e}"))?;
}
for (st, p) in &ss.priors {
p.validate()
.map_err(|e| format!("section_states[{pid}].priors[{st:?}]: {e}"))?;
}
}
Ok(())
}
pub fn validate_brain_state_snapshot_with_capacity(
snapshot: &BrainStateSnapshot,
entity_capacity: usize,
) -> Result<(), String> {
validate_brain_state_snapshot(snapshot)?;
fn check_recall(
label: &str,
br: &BalancedRecallSnapshot,
entity_capacity: usize,
) -> Result<(), String> {
if br.entity_posteriors.len() > entity_capacity {
return Err(format!(
"{label}.entity_posteriors: {} entries exceeds capacity {entity_capacity}",
br.entity_posteriors.len()
));
}
if br.entity_posteriors_version > 1 {
return Err(format!(
"{label}.entity_posteriors_version: unknown version {}",
br.entity_posteriors_version
));
}
if !br.entity_posterior_order.is_empty() {
let mut seen =
std::collections::HashSet::with_capacity(br.entity_posterior_order.len());
for id in &br.entity_posterior_order {
if !seen.insert(*id) {
return Err(format!("{label}.entity_posterior_order: duplicate id {id}"));
}
if !br.entity_posteriors.contains_key(id) {
return Err(format!(
"{label}.entity_posterior_order: id {id} not present in entity_posteriors"
));
}
}
}
if br.entity_posteriors_version == 0 && !br.entity_posterior_order.is_empty() {
return Err(format!(
"{label}.entity_posterior_order: non-empty order requires entity_posteriors_version >= 1 (got version 0, the legacy empty-order compatibility version)",
));
}
if br.entity_posteriors_version >= 1
&& br.entity_posterior_order.len() != br.entity_posteriors.len()
{
return Err(format!(
"{label}.entity_posterior_order: length {} does not cover all {} entity_posteriors entries (version {} requires full order coverage)",
br.entity_posterior_order.len(),
br.entity_posteriors.len(),
br.entity_posteriors_version,
));
}
Ok(())
}
check_recall(
"balanced_recall",
&snapshot.balanced_recall,
entity_capacity,
)?;
for (pid, ps) in &snapshot.profile_states {
check_recall(&format!("profile_states[{pid}]"), ps, entity_capacity)?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use chrono::{TimeZone, Utc};
use super::*;
use crate::profile::ProfileLifecycle;
fn make_profile(id: &str, consumer_kind: &str, ts_secs: i64) -> ProfileRecord {
ProfileRecord {
id: id.to_owned(),
description: String::new(),
consumer_kind: consumer_kind.to_owned(),
state_class: "Bayesian".into(),
lifecycle: ProfileLifecycle::Active,
created_at: Utc.timestamp_opt(ts_secs, 0).unwrap(),
state_snapshot: None,
total_events: 0,
exploration_epoch: 0,
}
}
#[test]
fn sort_fallback_candidates_produces_created_at_then_id_order() {
let p1 = make_profile("aardvark", "recall", 1_000); let p2 = make_profile("mango", "recall", 2_000);
let p3 = make_profile("zebra", "recall", 3_000);
let mut candidates = vec![&p3, &p2, &p1];
sort_fallback_candidates(&mut candidates);
assert_eq!(
candidates[0].id, "aardvark",
"earliest profile must come first"
);
assert_eq!(candidates[1].id, "mango");
assert_eq!(candidates[2].id, "zebra", "latest profile must come last");
}
#[test]
fn sort_fallback_candidates_breaks_ties_by_id() {
let ts = 5_000i64;
let p_z = make_profile("z-profile", "recall", ts);
let p_a = make_profile("a-profile", "recall", ts);
let p_m = make_profile("m-profile", "recall", ts);
let mut candidates = vec![&p_z, &p_m, &p_a];
sort_fallback_candidates(&mut candidates);
assert_eq!(
candidates[0].id, "a-profile",
"lowest id must come first on equal timestamps"
);
assert_eq!(candidates[1].id, "m-profile");
assert_eq!(candidates[2].id, "z-profile");
}
#[test]
fn resolve_fallback_is_deterministic() {
let p_early = make_profile("alpha", "recall", 1_000);
let p_later = make_profile("zeta", "recall", 2_000);
let mut state_a = BrainState {
profiles: HashMap::new(),
balanced_recall: BalancedRecallState::new(8),
profile_states: HashMap::new(),
bindings: Vec::new(),
section_states: HashMap::new(),
router_state: HashMap::new(),
adapter_set: HashMap::new(),
};
state_a.profiles.insert(p_early.id.clone(), p_early.clone());
state_a.profiles.insert(p_later.id.clone(), p_later.clone());
let mut state_b = BrainState {
profiles: HashMap::new(),
balanced_recall: BalancedRecallState::new(8),
profile_states: HashMap::new(),
bindings: Vec::new(),
section_states: HashMap::new(),
router_state: HashMap::new(),
adapter_set: HashMap::new(),
};
state_b.profiles.insert(p_later.id.clone(), p_later.clone());
state_b.profiles.insert(p_early.id.clone(), p_early.clone());
let result_a = state_a.resolve(None, None, "recall");
let result_b = state_b.resolve(None, None, "recall");
let id_a = result_a.map(|p| p.id.clone());
let id_b = result_b.map(|p| p.id.clone());
assert_eq!(id_a, Some("alpha".to_owned()));
assert_eq!(id_a, id_b, "fallback resolution must be deterministic");
}
#[test]
fn router_state_and_adapter_set_round_trip() {
let mut state = BrainState::new(8);
let blob = RouterStateBlob {
schema_version: 1,
gate_bytes: vec![1, 2, 3],
};
state
.router_state
.insert("profile-x".to_owned(), blob.clone());
let record = AdapterRecord {
adapter_id: "lora-42".to_owned(),
slot: 0,
content_hash: "abc123".to_owned(),
};
state
.adapter_set
.insert("profile-x".to_owned(), vec![record.clone()]);
let snapshot = state.to_snapshot();
let restored = BrainState::from_snapshot(snapshot, 8);
assert_eq!(
restored.router_state.get("profile-x"),
Some(&blob),
"router_state must survive round-trip"
);
assert_eq!(
restored.adapter_set.get("profile-x"),
Some(&vec![record]),
"adapter_set must survive round-trip"
);
}
#[test]
fn snapshot_missing_router_and_adapter_fields_defaults_to_empty() {
let state = BrainState::new(8);
let snapshot = state.to_snapshot();
let mut value: serde_json::Value =
serde_json::to_value(&snapshot).expect("serialize snapshot");
let obj = value.as_object_mut().expect("snapshot is a JSON object");
obj.remove("router_state");
obj.remove("adapter_set");
let restored: BrainStateSnapshot =
serde_json::from_value(value).expect("deserialize old-format snapshot");
assert!(
restored.router_state.is_empty(),
"router_state must default to empty map when absent from JSON"
);
assert!(
restored.adapter_set.is_empty(),
"adapter_set must default to empty map when absent from JSON"
);
}
#[test]
fn validate_brain_state_snapshot_with_capacity_rejects_entity_posteriors_over_capacity() {
use uuid::Uuid;
let capacity = 2;
let mut state = BrainState::new(capacity);
for _ in 0..(capacity + 1) {
state
.balanced_recall
.entity_posteriors
.get_or_insert(Uuid::new_v4(), crate::posterior::BetaPosterior::default);
}
let mut snapshot = state.to_snapshot();
snapshot
.balanced_recall
.entity_posteriors
.insert(Uuid::new_v4(), crate::posterior::BetaPosterior::default());
let result = validate_brain_state_snapshot_with_capacity(&snapshot, capacity);
assert!(
result.is_err(),
"snapshot with entity_posteriors.len() > capacity must be rejected"
);
}
#[test]
fn validate_brain_state_snapshot_with_capacity_rejects_profile_state_over_capacity() {
use uuid::Uuid;
let capacity = 1;
let mut state = BrainState::new(capacity);
let mut extra = BalancedRecallState::new(capacity);
extra
.entity_posteriors
.get_or_insert(Uuid::new_v4(), crate::posterior::BetaPosterior::default);
state
.profile_states
.insert("extra-profile".to_owned(), extra);
let mut snapshot = state.to_snapshot();
snapshot
.profile_states
.get_mut("extra-profile")
.unwrap()
.entity_posteriors
.insert(Uuid::new_v4(), crate::posterior::BetaPosterior::default());
let result = validate_brain_state_snapshot_with_capacity(&snapshot, capacity);
assert!(
result.is_err(),
"profile_states entry with entity_posteriors.len() > capacity must be rejected"
);
}
fn make_versioned_recall_snapshot(capacity: usize, n: usize) -> BalancedRecallSnapshot {
let mut state = BalancedRecallState::new(capacity);
for _ in 0..n {
state.entity_posteriors.get_or_insert(
uuid::Uuid::new_v4(),
crate::posterior::BetaPosterior::default,
);
}
state.to_snapshot()
}
#[test]
fn validate_brain_state_snapshot_with_capacity_rejects_missing_order_ids() {
let capacity = 4;
let mut snapshot = make_versioned_recall_snapshot(capacity, 2);
assert_eq!(snapshot.entity_posteriors_version, 1);
assert_eq!(snapshot.entity_posterior_order.len(), 2);
snapshot.entity_posterior_order.truncate(1);
let mut full_snapshot = BrainState::new(capacity).to_snapshot();
full_snapshot.balanced_recall = snapshot;
let result = validate_brain_state_snapshot_with_capacity(&full_snapshot, capacity);
assert!(
result.is_err(),
"version-1 snapshot with order shorter than entity_posteriors must be rejected"
);
assert!(
result.unwrap_err().contains("does not cover all"),
"error must name the missing-coverage reason"
);
}
#[test]
fn validate_brain_state_snapshot_with_capacity_rejects_duplicate_order_ids() {
let capacity = 4;
let mut snapshot = make_versioned_recall_snapshot(capacity, 2);
let dup = snapshot.entity_posterior_order[0];
snapshot.entity_posterior_order[1] = dup;
let mut full_snapshot = BrainState::new(capacity).to_snapshot();
full_snapshot.balanced_recall = snapshot;
let result = validate_brain_state_snapshot_with_capacity(&full_snapshot, capacity);
assert!(
result.is_err(),
"duplicate entity_posterior_order ids must be rejected"
);
assert!(result.unwrap_err().contains("duplicate id"));
}
#[test]
fn validate_brain_state_snapshot_with_capacity_rejects_unknown_order_ids() {
let capacity = 4;
let mut snapshot = make_versioned_recall_snapshot(capacity, 2);
snapshot.entity_posterior_order.push(uuid::Uuid::new_v4());
let mut full_snapshot = BrainState::new(capacity).to_snapshot();
full_snapshot.balanced_recall = snapshot;
let result = validate_brain_state_snapshot_with_capacity(&full_snapshot, capacity);
assert!(
result.is_err(),
"entity_posterior_order id absent from entity_posteriors must be rejected"
);
assert!(result
.unwrap_err()
.contains("not present in entity_posteriors"));
}
#[test]
fn validate_brain_state_snapshot_with_capacity_rejects_unknown_version() {
let capacity = 4;
let mut snapshot = make_versioned_recall_snapshot(capacity, 1);
snapshot.entity_posteriors_version = 2;
let mut full_snapshot = BrainState::new(capacity).to_snapshot();
full_snapshot.balanced_recall = snapshot;
let result = validate_brain_state_snapshot_with_capacity(&full_snapshot, capacity);
assert!(
result.is_err(),
"unknown entity_posteriors_version must be rejected"
);
assert!(result.unwrap_err().contains("unknown version"));
}
#[test]
fn validate_brain_state_snapshot_with_capacity_accepts_legacy_empty_order() {
let capacity = 4;
let mut snapshot = make_versioned_recall_snapshot(capacity, 2);
snapshot.entity_posteriors_version = 0;
snapshot.entity_posterior_order.clear();
let mut full_snapshot = BrainState::new(capacity).to_snapshot();
full_snapshot.balanced_recall = snapshot;
let result = validate_brain_state_snapshot_with_capacity(&full_snapshot, capacity);
assert!(
result.is_ok(),
"legacy version-0 snapshot with empty order must be accepted: {result:?}"
);
}
#[test]
fn validate_brain_state_snapshot_with_capacity_rejects_version_zero_with_partial_order() {
let capacity = 4;
let mut snapshot = make_versioned_recall_snapshot(capacity, 2);
snapshot.entity_posteriors_version = 0;
snapshot.entity_posterior_order.truncate(1);
let mut full_snapshot = BrainState::new(capacity).to_snapshot();
full_snapshot.balanced_recall = snapshot;
let result = validate_brain_state_snapshot_with_capacity(&full_snapshot, capacity);
assert!(
result.is_err(),
"version-0 snapshot with non-empty partial order must be rejected"
);
assert!(
result
.unwrap_err()
.contains("requires entity_posteriors_version >= 1"),
"error must name the version-0-with-non-empty-order reason"
);
}
#[test]
fn validate_brain_state_snapshot_with_capacity_accepts_full_order_current_version() {
let capacity = 4;
let mut full_snapshot = BrainState::new(capacity).to_snapshot();
full_snapshot.balanced_recall = make_versioned_recall_snapshot(capacity, 3);
let result = validate_brain_state_snapshot_with_capacity(&full_snapshot, capacity);
assert!(
result.is_ok(),
"version-1 snapshot with a fully-covering order must be accepted: {result:?}"
);
}
#[test]
fn restore_to_snapshot_restore_idempotency() {
let capacity = 4;
let mut state = BrainState::new(capacity);
for _ in 0..3 {
state.balanced_recall.entity_posteriors.get_or_insert(
uuid::Uuid::new_v4(),
crate::posterior::BetaPosterior::default,
);
}
let snapshot_1 = state.to_snapshot();
validate_brain_state_snapshot_with_capacity(&snapshot_1, capacity)
.expect("first snapshot must validate");
let restored_1 = BrainState::from_snapshot(snapshot_1, capacity);
let snapshot_2 = restored_1.to_snapshot();
validate_brain_state_snapshot_with_capacity(&snapshot_2, capacity)
.expect("round-tripped snapshot must still validate");
let restored_2 = BrainState::from_snapshot(snapshot_2.clone(), capacity);
let snapshot_3 = restored_2.to_snapshot();
assert_eq!(
snapshot_2.balanced_recall.entity_posteriors,
snapshot_3.balanced_recall.entity_posteriors,
"entity_posteriors must be stable across a second restore/snapshot cycle"
);
assert_eq!(
snapshot_2.balanced_recall.entity_posterior_order,
snapshot_3.balanced_recall.entity_posterior_order,
"entity_posterior_order must be stable across a second restore/snapshot cycle"
);
}
}