use std::sync::Arc;
use rmcp::ErrorData as McpError;
use rmcp::model::CallToolResult;
use super::OutlineEntry;
use super::ServerState;
use super::helpers::{HashMode, json_result, symbol_fingerprint};
use super::types_governance::{AuditResult, AuditVerdict, MemoryAuditParams, MemoryAuditResponse};
use super::types_memory::{MemoryRecord, VerifyState};
const STALE_DECAY: f32 = 0.5;
const ARCHIVE_AFTER_MICROS: i64 = 90 * 24 * 60 * 60 * 1_000_000;
pub(super) const DEFAULT_AUDIT_LIMIT: u32 = 100;
const MAX_AUDIT_LIMIT: u32 = 1000;
pub(super) fn write_live(
idx: &crate::index::IndexDb,
scope: &str,
vis_byte: u8,
owner: &str,
key: &str,
record: &MemoryRecord,
) -> Result<(), McpError> {
let raw_key = crate::index::keys::memory_by_key(scope, vis_byte, owner, key);
let bytes = rmp_serde::to_vec_named(record)
.map_err(|e| McpError::internal_error(format!("serialize memory record: {e}"), None))?;
idx.memory_by_key
.insert(raw_key, bytes)
.map_err(|e| McpError::internal_error(format!("fjall insert: {e}"), None))
}
fn write_archive(
idx: &crate::index::IndexDb,
scope: &str,
vis_byte: u8,
owner: &str,
key: &str,
record: &MemoryRecord,
) -> Result<(), McpError> {
let raw_key = crate::index::keys::memory_by_key(scope, vis_byte, owner, key);
let bytes = rmp_serde::to_vec_named(record)
.map_err(|e| McpError::internal_error(format!("serialize archive record: {e}"), None))?;
idx.memory_archive
.insert(raw_key, bytes)
.map_err(|e| McpError::internal_error(format!("fjall archive insert: {e}"), None))
}
fn delete_live(
idx: &crate::index::IndexDb,
scope: &str,
vis_byte: u8,
owner: &str,
key: &str,
) -> Result<(), McpError> {
let raw_key = crate::index::keys::memory_by_key(scope, vis_byte, owner, key);
idx.memory_by_key
.remove(raw_key)
.map_err(|e| McpError::internal_error(format!("fjall remove: {e}"), None))
}
pub(super) fn audit_one_record(
cache: &super::MapCache,
store: &crate::store::Store,
root: &std::path::Path,
record: &MemoryRecord,
) -> AuditVerdict {
let prov = &record.provenance;
if prov.files.is_empty() && prov.symbols.is_empty() && prov.commands.is_empty() {
return AuditVerdict {
state: VerifyState::Unverified,
reasons: vec!["no provenance".to_string()],
};
}
let mut reasons: Vec<String> = Vec::new();
let mut stale = false;
for rel in &prov.files {
if !cache.by_path.contains_key(rel) {
reasons.push(format!("file deleted: {}", rel.to_str_lossy()));
stale = true;
}
}
for sym_ref in &prov.symbols {
let Some(l1) = cache.by_path.get(&sym_ref.path) else {
reasons.push(format!(
"symbol not found: {} (file gone: {})",
sym_ref.name,
sym_ref.path.to_str_lossy()
));
stale = true;
continue;
};
let sym = l1.symbols.iter().find(|s| {
s.name == sym_ref.name
&& sym_ref
.kind
.as_deref()
.is_none_or(|k| crate::mcp::helpers::kind_to_str(s.kind) == k)
});
let Some(sym) = sym else {
reasons.push(format!("symbol not found: {}", sym_ref.name));
stale = true;
continue;
};
if let Some(stored_hash) = sym_ref.structural_hash {
if let Some(lang) = crate::lang::intern(&l1.language) {
let abs_path = root.join(sym_ref.path.to_path_buf());
if let Ok(source) = std::fs::read(&abs_path) {
let entry = OutlineEntry {
map: Arc::new(l1.clone()),
source: Arc::new(source),
};
let kind_opt = sym_ref.kind.as_deref().and_then(parse_kind_opt);
if let Some(current_hash) = symbol_fingerprint(
&entry,
&sym_ref.name,
kind_opt,
lang,
HashMode::Structural,
) && current_hash != stored_hash
{
reasons.push(format!("symbol body changed: {}", sym_ref.name));
stale = true;
}
}
}
}
let _ = sym; }
for cmd in &prov.commands {
let first_token = cmd.split_whitespace().next().unwrap_or(cmd.as_str());
let cmd_rel = crate::path::RelPath::from(first_token);
if cache.by_path.contains_key(&cmd_rel) {
if store.lookup(first_token).is_none() {
reasons.push(format!("command may be stale: {cmd}"));
}
}
}
if stale {
return AuditVerdict {
state: VerifyState::Stale,
reasons,
};
}
if !prov.files.is_empty() || !prov.symbols.is_empty() {
return AuditVerdict {
state: VerifyState::Verified,
reasons,
};
}
AuditVerdict {
state: VerifyState::Unverified,
reasons,
}
}
fn parse_kind_opt(k: &str) -> Option<crate::extract::SymbolKind> {
use crate::extract::SymbolKind;
Some(match k {
"function" => SymbolKind::Function,
"method" => SymbolKind::Method,
"struct" => SymbolKind::Struct,
"enum" => SymbolKind::Enum,
"class" => SymbolKind::Class,
"interface" => SymbolKind::Interface,
"trait" => SymbolKind::Trait,
"type" => SymbolKind::Type,
"const" => SymbolKind::Const,
"module" => SymbolKind::Module,
"macro" => SymbolKind::Macro,
"impl" => SymbolKind::Impl,
"namespace" => SymbolKind::Namespace,
"getter" => SymbolKind::Getter,
"setter" => SymbolKind::Setter,
"field" => SymbolKind::Field,
"variable" => SymbolKind::Variable,
"enum_variant" => SymbolKind::EnumVariant,
"constructor" => SymbolKind::Constructor,
"decorator" => SymbolKind::Decorator,
_ => return None,
})
}
struct EntryOutcome {
record: MemoryRecord,
audit_result: AuditResult,
}
pub(super) struct AuditCtx<'a> {
cache: &'a super::MapCache,
store: &'a crate::store::Store,
root: &'a std::path::Path,
dry_run: bool,
now: i64,
}
fn evaluate_one(
ctx: &AuditCtx<'_>,
key: &str,
raw_val: &[u8],
from_archive: bool,
) -> Option<EntryOutcome> {
let mut record: MemoryRecord = rmp_serde::from_slice(raw_val).ok()?;
let verdict = audit_one_record(ctx.cache, ctx.store, ctx.root, &record);
let state_str = verdict.state_str().to_string();
let mut archived = false;
if !ctx.dry_run && !from_archive {
let prev_last_verified = record.last_verified;
record.verified = verdict.state;
record.last_verified = ctx.now;
if verdict.state == VerifyState::Stale {
record.importance *= STALE_DECAY;
let stale_since = if prev_last_verified > 0 {
prev_last_verified
} else {
record.updated_at
};
if ctx.now.saturating_sub(stale_since) > ARCHIVE_AFTER_MICROS {
archived = true;
}
}
}
Some(EntryOutcome {
record,
audit_result: AuditResult {
key: key.to_string(),
state: state_str,
reasons: verdict.reasons,
archived,
},
})
}
pub(super) async fn run_memory_audit(
state: &ServerState,
params: MemoryAuditParams,
) -> Result<CallToolResult, McpError> {
let limit = params
.limit
.unwrap_or(DEFAULT_AUDIT_LIMIT)
.min(MAX_AUDIT_LIMIT) as usize;
let scan_cap = limit.saturating_mul(8).max(1_000);
let vis_byte = params.visibility.vis_byte();
let owner: &str = match params.visibility {
super::types_memory::Visibility::Individual => &state.agent_id,
super::types_memory::Visibility::Group => "",
};
let cache = state.cache.load_full();
let root = state.root.clone();
let store_guard = state.store.read().await;
let idx = store_guard
.index_db
.as_ref()
.ok_or_else(|| McpError::internal_error("memory_by_key index not available", None))?;
let now = crate::lance::now_micros();
let ctx = AuditCtx {
cache: &cache,
store: &store_guard,
root: &root,
dry_run: params.dry_run,
now,
};
let mut results: Vec<AuditResult> = Vec::new();
if let Some(ref single_key) = params.key {
let raw_key = crate::index::keys::memory_by_key(&state.scope, vis_byte, owner, single_key);
let keyspace = if params.include_archived {
&idx.memory_archive
} else {
&idx.memory_by_key
};
let raw_val_opt = keyspace
.get(&raw_key)
.map_err(|e| McpError::internal_error(format!("fjall get: {e}"), None))?;
if let Some(raw_val) = raw_val_opt
&& let Some(outcome) = evaluate_one(&ctx, single_key, &raw_val, params.include_archived)
{
if !ctx.dry_run {
if outcome.audit_result.archived {
write_archive(
idx,
&state.scope,
vis_byte,
owner,
single_key,
&outcome.record,
)?;
delete_live(idx, &state.scope, vis_byte, owner, single_key)?;
} else {
write_live(
idx,
&state.scope,
vis_byte,
owner,
single_key,
&outcome.record,
)?;
}
}
results.push(outcome.audit_result);
}
} else {
let ns_prefix = crate::index::keys::memory_by_key_ns_prefix(&state.scope, vis_byte, owner);
for (scanned, guard) in idx.memory_by_key.prefix(&ns_prefix).enumerate() {
if results.len() >= limit || scanned >= scan_cap {
break;
}
let (raw_key_bytes, raw_val) = guard
.into_inner()
.map_err(|e| McpError::internal_error(format!("index iter: {e}"), None))?;
let Some(key) = crate::index::keys::parse_memory_key_only(&raw_key_bytes) else {
continue;
};
let key_str = key.to_string();
if let Some(outcome) = evaluate_one(&ctx, &key_str, &raw_val, false) {
if !ctx.dry_run {
if outcome.audit_result.archived {
write_archive(
idx,
&state.scope,
vis_byte,
owner,
&key_str,
&outcome.record,
)?;
delete_live(idx, &state.scope, vis_byte, owner, &key_str)?;
} else {
write_live(
idx,
&state.scope,
vis_byte,
owner,
&key_str,
&outcome.record,
)?;
}
}
results.push(outcome.audit_result);
}
}
if params.include_archived {
for (arch_scanned, guard) in idx.memory_archive.prefix(&ns_prefix).enumerate() {
if results.len() >= limit || arch_scanned >= scan_cap {
break;
}
let (raw_key_bytes, raw_val) = guard
.into_inner()
.map_err(|e| McpError::internal_error(format!("archive iter: {e}"), None))?;
let Some(key) = crate::index::keys::parse_memory_key_only(&raw_key_bytes) else {
continue;
};
let key_str = key.to_string();
if let Some(outcome) = evaluate_one(&ctx, &key_str, &raw_val, true) {
results.push(outcome.audit_result);
}
}
}
}
let audited = results.len();
json_result(&MemoryAuditResponse { audited, results })
}
#[cfg(all(test, feature = "memory"))]
mod tests;
pub(super) async fn audit_scope_on_rescan(state: &Arc<ServerState>) {
let cache = state.cache.load_full();
let root = state.root.clone();
let store_guard = state.store.read().await;
let idx = match store_guard.index_db.as_ref() {
Some(idx) => idx,
None => return,
};
let ctx = AuditCtx {
cache: &cache,
store: &store_guard,
root: &root,
dry_run: false,
now: crate::lance::now_micros(),
};
audit_scope_persist(idx, &ctx, &state.scope, DEFAULT_AUDIT_LIMIT as usize);
}
pub(super) fn audit_scope_persist(
idx: &crate::index::IndexDb,
ctx: &AuditCtx<'_>,
scope: &str,
limit: usize,
) {
let scope_prefix = crate::index::keys::memory_scope_prefix(scope);
let scan_cap = limit.saturating_mul(8).max(1_000);
let mut evaluated = 0usize;
for (scanned, guard) in idx.memory_by_key.prefix(&scope_prefix).enumerate() {
if evaluated >= limit || scanned >= scan_cap {
break;
}
let (raw_key_bytes, raw_val) = match guard.into_inner() {
Ok(pair) => pair,
Err(e) => {
tracing::warn!(error = %e, "audit_scope_on_rescan: iter error");
continue;
}
};
let Some((_scope, vis_byte, owner, key_str)) =
crate::index::keys::parse_memory_by_key(&raw_key_bytes)
else {
continue;
};
let Some(outcome) = evaluate_one(ctx, &key_str, &raw_val, false) else {
continue;
};
evaluated += 1;
if outcome.record.verified != VerifyState::Stale {
continue;
}
let write = if outcome.audit_result.archived {
write_archive(idx, scope, vis_byte, &owner, &key_str, &outcome.record)
.and_then(|()| delete_live(idx, scope, vis_byte, &owner, &key_str))
} else {
write_live(idx, scope, vis_byte, &owner, &key_str, &outcome.record)
};
if let Err(e) = write {
tracing::warn!(key = key_str, error = ?e, "audit_scope_on_rescan: persist failed");
}
}
}