#![forbid(unsafe_code)]
#![deny(rust_2018_idioms, unreachable_pub)]
use std::sync::Arc;
use std::time::Duration;
use futures::StreamExt;
use lunaris_consolidate::{
CONSOLIDATE_CONSUMER_GROUP, CONSOLIDATE_TOPIC, ConsolidateEvent, ConsolidationReport,
};
use lunaris_core::keyspace::{episode_key, source_index_key};
use lunaris_core::storage::types::{Filter, Lsn, WriteOp};
use lunaris_core::{Episode, HlcClock, LunarisError, Scope, StorageError, StoragePort};
use lunaris_retrieve::{Hit, Keyword, Query, Vector};
use ulid::Ulid;
use crate::Lunaris;
const DRAIN_CAP: usize = 1024;
const PULL_TIMEOUT_MS: u64 = 50;
const DEFAULT_TOP_K: usize = 8;
const FANOUT: usize = 3;
const RRF_K: u32 = 60;
#[derive(Clone)]
pub struct WorkingMemory {
lunaris: Arc<Lunaris>,
scope: Scope,
scope_prefix: String,
}
impl WorkingMemory {
pub fn new(lunaris: Arc<Lunaris>, scope: Scope, scope_prefix: impl Into<String>) -> Self {
Self { lunaris, scope, scope_prefix: scope_prefix.into() }
}
pub async fn write(&self, k: &str, v: serde_json::Value) -> Result<Lsn, LunarisError> {
self.write_inner(k, v, None).await
}
pub async fn write_dated(
&self,
k: &str,
v: serde_json::Value,
t_ref: chrono::DateTime<chrono::Utc>,
) -> Result<Lsn, LunarisError> {
self.write_inner(k, v, Some(t_ref)).await
}
async fn write_inner(
&self,
k: &str,
v: serde_json::Value,
t_ref: Option<chrono::DateTime<chrono::Utc>>,
) -> Result<Lsn, LunarisError> {
let source = self.scope_key(k);
let content = serde_json::to_string(&v)
.map_err(|e| LunarisError::from(lunaris_core::StorageError::from(e)))?;
let mut episode = Episode::new(
self.scope.clone(),
source.clone(),
content,
self.lunaris.clock().as_ref(),
);
episode.t_ref = t_ref;
let episode_id = episode.id;
let lsn = self.lunaris.ingest(episode).await?;
self.record_source_index(&source, episode_id).await;
Ok(lsn)
}
async fn record_source_index(&self, source: &str, episode_id: Ulid) {
let key = source_index_key(&self.scope, source);
let op = WriteOp::KvPut { key, value: episode_id.to_bytes().to_vec() };
if let Err(err) = self.lunaris.storage().atomic_write(&self.scope, &[op]).await {
tracing::warn!(
error = %err,
source = %source,
"working_memory_source_index_write_failed; read() falls back to the ranked path"
);
}
}
pub async fn read(&self, k: &str) -> Result<Option<serde_json::Value>, LunarisError> {
self.read_at(k, None).await
}
pub(crate) async fn read_at(
&self,
k: &str,
as_of: Option<lunaris_core::Hlc>,
) -> Result<Option<serde_json::Value>, LunarisError> {
let source = self.scope_key(k);
if let Some(id) = self.lookup_source_index(&source, as_of).await? {
return self.recover_value(&id, as_of).await;
}
let filter =
Filter::Eq { field: "source".into(), value: serde_json::Value::String(source) };
match self.find(k, filter).await?.into_iter().max_by(|a, b| a.episode_id.cmp(&b.episode_id))
{
Some(h) => self.recover_value(&h.episode_id, as_of).await,
None => Ok(None),
}
}
async fn lookup_source_index(
&self,
source: &str,
as_of: Option<lunaris_core::Hlc>,
) -> Result<Option<Vec<u8>>, LunarisError> {
let key = source_index_key(&self.scope, source);
let snapshot = as_of.unwrap_or_else(|| HlcClock::new(0).tick());
match self.lunaris.storage().read_as_of(&self.scope, &key, snapshot).await {
Ok(Some(row)) if row.value.len() == 16 => Ok(Some(row.value.to_vec())),
Ok(Some(row)) => {
tracing::warn!(
len = row.value.len(),
source = %source,
"working_memory_source_index_bad_width; falling back to the ranked path"
);
Ok(None)
}
Ok(None) => Ok(None),
Err(err) if lunaris_retrieve::missing_index::is_index_absent(&err) => Ok(None),
Err(err) => Err(LunarisError::from(err)),
}
}
pub async fn grep(
&self,
pattern: &str,
) -> Result<Vec<(String, serde_json::Value)>, LunarisError> {
let filter = Filter::StartsWith { field: "source".into(), prefix: self.scope_key(pattern) };
let hits = self.find(pattern, filter).await?;
let mut seen: std::collections::HashSet<Vec<u8>> = std::collections::HashSet::new();
let mut out = Vec::with_capacity(hits.len());
for h in hits {
if !seen.insert(h.episode_id.clone()) {
continue;
}
if let Some(v) = self.recover_value(&h.episode_id, None).await? {
out.push((h.source, v));
}
}
Ok(out)
}
async fn find(&self, query: &str, filter: Filter) -> Result<Vec<Hit>, LunarisError> {
let fused = Vector::new("chunks", DEFAULT_TOP_K * FANOUT)
.and(Keyword::bm25("chunks", DEFAULT_TOP_K * FANOUT))
.fuse_rrf(RRF_K)
.top(DEFAULT_TOP_K);
let hits = match self
.lunaris
.recall()
.with_scope(self.scope.clone())
.with_root(fused)
.filter(filter.clone())
.execute(Query::text(query))
.await
{
Ok(hits) => hits,
Err(err) if is_ft_index_missing(&err) => return Ok(Vec::new()),
Err(err) if is_keyword_not_supported(&err) || is_ft_query_unusable(&err) => {
match self
.lunaris
.recall()
.with_scope(self.scope.clone())
.with_root(Vector::new("chunks", DEFAULT_TOP_K * FANOUT))
.filter(filter.clone())
.execute(Query::text(query))
.await
{
Ok(hits) => hits,
Err(err) if is_ft_index_missing(&err) => return Ok(Vec::new()),
Err(err) => return Err(err),
}
}
Err(err) => return Err(err),
};
Ok(hits.into_iter().filter(|h| source_filter_matches(&filter, &h.source)).collect())
}
async fn recover_value(
&self,
episode_id: &[u8],
as_of: Option<lunaris_core::Hlc>,
) -> Result<Option<serde_json::Value>, LunarisError> {
let bytes: [u8; 16] = match episode_id.try_into() {
Ok(b) => b,
Err(_) => return Ok(None),
};
let key = episode_key(&self.scope, Ulid::from_bytes(bytes));
let snapshot = as_of.unwrap_or_else(|| HlcClock::new(0).tick());
match self.lunaris.storage().read_as_of(&self.scope, &key, snapshot).await? {
Some(row) => {
let episode: Episode = serde_json::from_slice(&row.value)
.map_err(|e| LunarisError::from(StorageError::from(e)))?;
let value = serde_json::from_str(&episode.content)
.map_err(|e| LunarisError::from(StorageError::from(e)))?;
Ok(Some(value))
}
None => Ok(None),
}
}
pub async fn consolidate(&self) -> Result<ConsolidationReport, LunarisError> {
let storage: Arc<dyn StoragePort> = self.lunaris.storage();
let events = drain_consolidate_events(&storage, &self.scope).await?;
let scope_prefix: &str = &self.scope_prefix;
let mut matching: Vec<ConsolidateEvent> = Vec::with_capacity(events.len());
let mut foreign: Vec<ConsolidateEvent> = Vec::new();
for ev in events {
if ev.source.starts_with(scope_prefix) {
matching.push(ev);
} else {
foreign.push(ev);
}
}
let mut lost: usize = 0;
for ev in &foreign {
match serde_json::to_vec(ev) {
Ok(payload) => {
if let Err(e) =
storage.publish(&self.scope, CONSOLIDATE_TOPIC, 0, payload.into()).await
{
tracing::warn!(
source = %ev.source,
error = %e,
"consolidate: failed to re-queue foreign event; \
it will be lost for this scope's pass"
);
lost += 1;
}
}
Err(e) => {
tracing::warn!(
source = %ev.source,
error = %e,
"consolidate: serde failure serialising foreign event for re-queue"
);
lost += 1;
}
}
}
if lost > 0 {
tracing::warn!(
lost,
scope_prefix,
"consolidate: {} foreign event(s) could not be re-queued and will be lost",
lost
);
}
let pipeline = self.lunaris.consolidator_pipeline();
let consolidator = match pipeline.snapshot_consolidator() {
Some(c) => c,
None => {
return Ok(ConsolidationReport::default());
}
};
let report = consolidator.consolidate_scoped(storage.clone(), &matching, None).await?;
lunaris_consolidate::publish_per_event_audits(&storage, &self.scope, &report).await;
Ok(report)
}
pub async fn consolidate_unfiltered(&self) -> Result<ConsolidationReport, LunarisError> {
let storage: Arc<dyn StoragePort> = self.lunaris.storage();
let events = drain_consolidate_events(&storage, &self.scope).await?;
let pipeline = self.lunaris.consolidator_pipeline();
let consolidator = match pipeline.snapshot_consolidator() {
Some(c) => c,
None => {
return Ok(ConsolidationReport::default());
}
};
let report = consolidator.consolidate_scoped(storage.clone(), &events, None).await?;
lunaris_consolidate::publish_per_event_audits(&storage, &self.scope, &report).await;
Ok(report)
}
fn scope_key(&self, k: &str) -> String {
format!("{}{}", self.scope_prefix, k)
}
}
fn is_keyword_not_supported(err: &LunarisError) -> bool {
matches!(
err,
LunarisError::Storage(StorageError::NotSupported(msg))
if msg.contains("keyword_search") || msg.contains("keyword")
)
}
fn is_ft_query_unusable(err: &LunarisError) -> bool {
matches!(
err,
LunarisError::Storage(StorageError::Backend(msg))
if msg.contains("empty query after analysis")
)
}
fn is_ft_index_missing(err: &LunarisError) -> bool {
matches!(
err,
LunarisError::Storage(e) if lunaris_retrieve::missing_index::is_index_absent(e)
)
}
fn source_filter_matches(filter: &Filter, source: &str) -> bool {
match filter {
Filter::Eq { field, value } if field == "source" => value.as_str() == Some(source),
Filter::StartsWith { field, prefix } if field == "source" => source.starts_with(prefix),
Filter::And(xs) => xs.iter().all(|f| source_filter_matches(f, source)),
Filter::Or(xs) => xs.iter().any(|f| source_filter_matches(f, source)),
_ => true,
}
}
async fn drain_consolidate_events(
storage: &Arc<dyn StoragePort>,
scope: &Scope,
) -> Result<Vec<ConsolidateEvent>, LunarisError> {
let pull_timeout = Duration::from_millis(PULL_TIMEOUT_MS);
let mut stream = storage
.subscribe(scope, CONSOLIDATE_CONSUMER_GROUP, CONSOLIDATE_TOPIC, 0)
.await
.map_err(LunarisError::Storage)?;
let mut events = Vec::with_capacity(64);
while events.len() < DRAIN_CAP {
match tokio::time::timeout(pull_timeout, stream.next()).await {
Ok(Some(Ok(msg))) => {
if let Ok(ev) = serde_json::from_slice::<ConsolidateEvent>(&msg.payload) {
events.push(ev);
}
}
Ok(Some(Err(_))) | Ok(None) | Err(_) => break,
}
}
Ok(events)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn working_memory_public_surface_under_30_loc() {
let src = include_str!("./working_memory.rs");
let production = src.split("#[cfg(test)]").next().unwrap_or(src);
let pub_fns = production.matches(" pub fn ").count()
+ production.matches(" pub async fn ").count();
assert!(
pub_fns <= 7,
"PRIM-04 ≤30-LOC contract: WorkingMemory has {pub_fns} pub fns; cap is 7 \
(write_dated added for Mechanism-B session-date grounding, 2026-07-29)"
);
assert!(
pub_fns >= 3,
"PRIM-04 contract: WorkingMemory needs at least 3 public methods; got {pub_fns}"
);
}
#[test]
fn working_memory_scope_key_prefix_concatenation() {
fn scope_key(prefix: &str, k: &str) -> String {
format!("{prefix}{k}")
}
assert_eq!(scope_key("helios:fs/", "note-1"), "helios:fs/note-1");
assert_eq!(scope_key("chat:user-42/", "draft"), "chat:user-42/draft");
assert_eq!(scope_key("", "raw-key"), "raw-key");
}
#[test]
fn working_memory_grep_uses_starts_with_filter() {
let prefix = "chat:user-42/draft-";
let filter = Filter::StartsWith { field: "source".into(), prefix: prefix.into() };
match filter {
Filter::StartsWith { field, prefix: p } => {
assert_eq!(field, "source");
assert_eq!(p, "chat:user-42/draft-");
}
other => panic!("expected StartsWith variant; got {other:?}"),
}
}
#[test]
fn source_filter_rejects_foreign_sources() {
let eq = Filter::Eq {
field: "source".into(),
value: serde_json::Value::String("scratchpad/sess-b/plan".into()),
};
assert!(source_filter_matches(&eq, "scratchpad/sess-b/plan"));
assert!(!source_filter_matches(&eq, "scratchpad/sess-a/plan"), "Eq must reject leaks");
assert!(!source_filter_matches(&eq, "scratchpad/sess-a/blocker"));
let sw = Filter::StartsWith { field: "source".into(), prefix: "scratchpad/sess-b/".into() };
assert!(source_filter_matches(&sw, "scratchpad/sess-b/anything"));
assert!(!source_filter_matches(&sw, "scratchpad/sess-a/plan"), "prefix must reject leaks");
let other =
Filter::Eq { field: "kind".into(), value: serde_json::Value::String("x".into()) };
assert!(source_filter_matches(&other, "scratchpad/sess-a/plan"));
}
#[test]
fn working_memory_construction_records_scope() {
let s = format!("{}{}", "helios:fs/", "k");
assert!(s.starts_with("helios:fs/"));
assert!(s.ends_with("k"));
assert_eq!(s, "helios:fs/k");
}
}