use crate::models::field_names;
use anyhow::{Context, Result};
use rusqlite::Connection;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::models::Memory;
pub const REFLECTION_ORIGIN_KEY: &str = "reflection_origin";
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ReflectionOrigin {
pub memory_id: String,
pub peer_origin: Option<String>,
pub signing_agent: Option<String>,
pub original_depth: i32,
pub local_depth_at_arrival: Option<u32>,
pub is_reflection: bool,
}
#[must_use]
pub fn stamp_reflection_origin(mem: &Memory, sender_agent_id: &str, local_cap: u32) -> Memory {
if mem.reflection_depth <= 0 {
return mem.clone();
}
let mut out = mem.clone();
let mut meta_map: Map<String, Value> = match out.metadata.take() {
Value::Object(m) => m,
_ => Map::new(),
};
if !meta_map.contains_key(REFLECTION_ORIGIN_KEY) {
let stamp = serde_json::json!({
(field_names::PEER_ORIGIN): sender_agent_id,
(field_names::ORIGINAL_DEPTH): mem.reflection_depth,
(field_names::LOCAL_DEPTH_AT_ARRIVAL): local_cap,
});
meta_map.insert(REFLECTION_ORIGIN_KEY.to_string(), stamp);
}
out.metadata = Value::Object(meta_map);
out
}
pub fn reflection_origin(conn: &Connection, id: &str) -> Result<Option<ReflectionOrigin>> {
let mem = match crate::storage::get(conn, id).context("storage::get for reflection_origin")? {
Some(m) => m,
None => return Ok(None),
};
Ok(Some(reflection_origin_from_memory(&mem)))
}
#[must_use]
pub fn reflection_origin_from_memory(mem: &Memory) -> ReflectionOrigin {
let is_reflection = mem.reflection_depth > 0;
let signing_agent = mem
.metadata
.get("agent_id")
.and_then(Value::as_str)
.map(str::to_string);
let origin_obj = mem.metadata.get(REFLECTION_ORIGIN_KEY);
let peer_origin = origin_obj
.and_then(|v| v.get(field_names::PEER_ORIGIN))
.and_then(Value::as_str)
.map(str::to_string);
let local_depth_at_arrival = origin_obj
.and_then(|v| v.get(field_names::LOCAL_DEPTH_AT_ARRIVAL))
.and_then(Value::as_u64)
.and_then(|n| u32::try_from(n).ok());
ReflectionOrigin {
memory_id: mem.id.clone(),
peer_origin,
signing_agent,
original_depth: mem.reflection_depth,
local_depth_at_arrival,
is_reflection,
}
}
pub fn enforce_local_cap_on_derived(
new_depth: u32,
local_cap: u32,
sources: &[Memory],
) -> std::result::Result<(), LocalCapRefusal> {
if new_depth <= local_cap {
return Ok(());
}
let imported_peer = sources.iter().find_map(|m| {
m.metadata
.get(REFLECTION_ORIGIN_KEY)
.and_then(|v| v.get(field_names::PEER_ORIGIN))
.and_then(Value::as_str)
.map(str::to_string)
});
let max_source_depth = sources
.iter()
.map(|m| m.reflection_depth)
.max()
.unwrap_or(0)
.max(0);
Err(LocalCapRefusal {
attempted: new_depth,
local_cap,
max_source_depth: u32::try_from(max_source_depth).unwrap_or(u32::MAX),
imported_peer,
})
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LocalCapRefusal {
pub attempted: u32,
pub local_cap: u32,
pub max_source_depth: u32,
pub imported_peer: Option<String>,
}
impl std::fmt::Display for LocalCapRefusal {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self.imported_peer.as_deref() {
Some(peer) => write!(
f,
"remote reflection at depth {} (from peer {}), local depth limit {}",
self.max_source_depth, peer, self.local_cap,
),
None => write!(
f,
"reflection depth {} would exceed local cap {}",
self.attempted, self.local_cap,
),
}
}
}
impl std::error::Error for LocalCapRefusal {}
#[cfg(test)]
mod tests {
use super::*;
use crate::models::Tier;
use chrono::Utc;
fn reflection_memory(id: &str, depth: i32) -> Memory {
let now = Utc::now().to_rfc3339();
Memory {
id: id.to_string(),
tier: Tier::Mid,
namespace: "test".to_string(),
title: format!("reflection-{id}"),
content: "body".to_string(),
tags: vec![],
priority: 5,
confidence: 1.0,
source: "test".to_string(),
access_count: 0,
created_at: now.clone(),
updated_at: now,
last_accessed_at: None,
expires_at: None,
metadata: serde_json::json!({"agent_id": "ai:test"}),
reflection_depth: depth,
memory_kind: crate::models::MemoryKind::Observation,
entity_id: None,
persona_version: None,
citations: Vec::new(),
source_uri: None,
source_span: None,
confidence_source: crate::models::ConfidenceSource::CallerProvided,
confidence_signals: None,
confidence_decayed_at: None,
version: 1,
}
}
#[test]
fn reflection_origin_from_memory_derives_all_fields() {
let mut mem = reflection_memory("m-derive", 2);
let mut meta = serde_json::Map::new();
meta.insert("agent_id".to_string(), serde_json::json!("ai:signer@host"));
meta.insert(
REFLECTION_ORIGIN_KEY.to_string(),
serde_json::json!({ "peer_origin": "peer-x", "local_depth_at_arrival": 5 }),
);
mem.metadata = serde_json::Value::Object(meta);
let origin = reflection_origin_from_memory(&mem);
assert_eq!(origin.memory_id, "m-derive");
assert!(origin.is_reflection);
assert_eq!(origin.original_depth, 2);
assert_eq!(origin.signing_agent.as_deref(), Some("ai:signer@host"));
assert_eq!(origin.peer_origin.as_deref(), Some("peer-x"));
assert_eq!(origin.local_depth_at_arrival, Some(5));
}
#[test]
fn reflection_origin_from_memory_non_reflection_is_flagged_false() {
let origin = reflection_origin_from_memory(&reflection_memory("m-base", 0));
assert!(!origin.is_reflection);
assert_eq!(origin.original_depth, 0);
assert!(origin.peer_origin.is_none());
assert!(origin.local_depth_at_arrival.is_none());
}
#[test]
fn stamp_skips_non_reflection() {
let mut mem = reflection_memory("m1", 0);
let before = mem.metadata.clone();
let stamped = stamp_reflection_origin(&mem, "peer-A", 3);
assert_eq!(stamped.metadata, before);
mem.reflection_depth = 0;
assert!(stamped.metadata.get(REFLECTION_ORIGIN_KEY).is_none());
}
#[test]
fn stamp_records_peer_and_local_cap() {
let mem = reflection_memory("m1", 2);
let stamped = stamp_reflection_origin(&mem, "peer-A", 3);
let origin = stamped
.metadata
.get(REFLECTION_ORIGIN_KEY)
.expect("origin stamped");
assert_eq!(origin["peer_origin"].as_str(), Some("peer-A"));
assert_eq!(origin["original_depth"].as_i64(), Some(2));
assert_eq!(origin["local_depth_at_arrival"].as_u64(), Some(3));
}
#[test]
fn stamp_is_idempotent_first_writer_wins() {
let mem = reflection_memory("m1", 2);
let first = stamp_reflection_origin(&mem, "peer-A", 3);
let second = stamp_reflection_origin(&first, "peer-B", 5);
let origin = second
.metadata
.get(REFLECTION_ORIGIN_KEY)
.expect("origin preserved");
assert_eq!(origin["peer_origin"].as_str(), Some("peer-A"));
assert_eq!(origin["local_depth_at_arrival"].as_u64(), Some(3));
}
#[test]
fn enforce_local_cap_allows_when_under_limit() {
let sources = vec![reflection_memory("s1", 1)];
assert!(enforce_local_cap_on_derived(2, 3, &sources).is_ok());
}
#[test]
fn enforce_local_cap_refuses_with_imported_peer_named() {
let mut imported = reflection_memory("s1", 2);
imported.metadata = serde_json::json!({
"agent_id": "ai:test",
REFLECTION_ORIGIN_KEY: {
"peer_origin": "peer-A",
"original_depth": 2,
"local_depth_at_arrival": 3,
},
});
let refusal = enforce_local_cap_on_derived(3, 2, &[imported]).unwrap_err();
let msg = refusal.to_string();
assert!(
msg.contains("peer-A") && msg.contains("local depth limit 2"),
"refusal msg should name peer + local cap: {msg}"
);
}
}