use std::collections::HashSet;
use std::path::PathBuf;
use rmcp::handler::server::tool::ToolRouter;
use rmcp::handler::server::wrapper::Parameters;
use rmcp::tool_router;
use serde_json::json;
use crate::graph::edges::EdgeKind;
use crate::graph::Graph;
use crate::store::record::{
Category, ContextPacket, FileRecord, GotchaRecord, Priority, QualityTier, Record,
RecordLifecycle, StaleReviewPayload, StalenessTier,
};
use super::protocol::{
self as proto, Command, DecisionUpsertInput, DevNoteUpsertInput, GotchaConfirmInput,
GotchaDraftInput, GotchaTombstoneInput,
};
use super::server::{proxy_daemon_result, proxy_daemon_v2, ProxyDaemonResult};
use super::types::{MemBootstrapParams, MemGetParams, MemQueryParams, MemSetParams};
pub(crate) const VECTOR_B: &str =
"\n\n[mati] Before reading any file: call mem_get(\"file:<path>\").\n\
confidence>=0.6 + confirmed=true \u{2192} use record, skip file read.\n\
confidence<0.3 \u{2192} read file, consider mem_set to improve.\n\
\"add gotcha\" \u{2192} mem_set(Gotcha) then mati gotcha confirm <key>.";
const TOKEN_BUDGET: usize = 2_000;
const VECTOR_B_TOKENS: usize = 77;
fn estimate_tokens(text: &str) -> usize {
text.len() / 4
}
fn priority_weight(priority: &Priority) -> f32 {
match priority {
Priority::Low => 0.25,
Priority::Normal => 0.50,
Priority::High => 0.75,
Priority::Critical => 1.00,
}
}
pub(crate) fn record_to_agent_json(record: &Record) -> serde_json::Value {
let mut obj = serde_json::Map::new();
obj.insert("key".into(), serde_json::json!(record.key));
obj.insert("value".into(), serde_json::json!(record.value));
obj.insert("category".into(), serde_json::json!(record.category));
obj.insert("priority".into(), serde_json::json!(record.priority));
if !record.tags.is_empty() {
obj.insert("tags".into(), serde_json::json!(record.tags));
}
obj.insert(
"confidence".into(),
serde_json::json!(record.confidence.value),
);
obj.insert(
"confirmation_count".into(),
serde_json::json!(record.confidence.confirmation_count),
);
obj.insert("quality".into(), serde_json::json!(record.quality.value));
obj.insert(
"quality_tier".into(),
serde_json::json!(record.quality.tier),
);
if !record.quality.signals.is_empty() {
obj.insert(
"quality_signals".into(),
serde_json::json!(record.quality.signals),
);
}
obj.insert("source".into(), serde_json::json!(record.source));
obj.insert(
"staleness_tier".into(),
serde_json::json!(record.staleness.tier),
);
if let Some(ref url) = record.ref_url {
obj.insert("ref_url".into(), serde_json::json!(url));
}
if let Some(ref payload) = record.payload {
obj.insert("payload".into(), strip_payload(payload, &record.category));
}
serde_json::Value::Object(obj)
}
fn strip_payload(payload: &serde_json::Value, category: &Category) -> serde_json::Value {
let Some(obj) = payload.as_object() else {
return payload.clone();
};
let internal_fields: &[&str] = match category {
Category::File => &[
"token_cost_estimate",
"last_modified_session",
"content_hash",
],
Category::Gotcha => &["discovered_session"],
_ => &[],
};
if internal_fields.is_empty() {
return payload.clone();
}
let mut stripped = obj.clone();
for field in internal_fields {
stripped.remove(*field);
}
if matches!(category, Category::File) {
stripped.retain(|_, v| !matches!(v, serde_json::Value::Array(a) if a.is_empty()));
}
serde_json::Value::Object(stripped)
}
#[derive(Clone)]
pub struct MatiServer {
root: PathBuf,
pub(crate) tool_router: ToolRouter<Self>,
}
impl MatiServer {
pub fn with_socket_root(root: PathBuf) -> Self {
Self {
root,
tool_router: Self::tool_router(),
}
}
fn socket_error(op: &str, result: ProxyDaemonResult) -> String {
let message = match result {
ProxyDaemonResult::NotRunning => format!("{op}: daemon not running"),
ProxyDaemonResult::StaleSocket => format!("{op}: daemon socket stale"),
ProxyDaemonResult::Unresponsive => format!("{op}: daemon unresponsive"),
ProxyDaemonResult::Ok(v) => format!("{op}: malformed daemon response: {v}"),
};
json!({ "error": message }).to_string()
}
async fn socket_call(&self, op: &str, args: serde_json::Value) -> String {
match proxy_daemon_result(&self.root, op, args).await {
ProxyDaemonResult::Ok(v) => Self::format_envelope(op, v),
other => Self::socket_error(op, other),
}
}
async fn socket_call_typed(&self, cmd: Command) -> String {
let op = cmd.kind();
match proxy_daemon_v2(&self.root, cmd).await {
ProxyDaemonResult::Ok(v) => Self::format_envelope(op, v),
other => Self::socket_error(op, other),
}
}
fn format_envelope(op: &str, v: serde_json::Value) -> String {
if v.get("ok") != Some(&serde_json::Value::Bool(true)) {
let err = v
.get("error")
.and_then(|e| e.as_str())
.unwrap_or("daemon request failed");
let code = v.get("code").and_then(|c| c.as_str()).unwrap_or("");
if code.is_empty() {
return json!({ "error": err, "op": op }).to_string();
}
return json!({ "error": err, "op": op, "code": code }).to_string();
}
match v.get("data") {
Some(serde_json::Value::String(s)) => s.clone(),
Some(other) => other.to_string(),
None => json!({ "error": "daemon response missing data" }).to_string(),
}
}
}
#[tool_router]
impl MatiServer {
#[rmcp::tool(
name = "mem_get",
description = "Look up one mati knowledge record by key. Before reading a file directly, call this with \"file:<path>\" and use the record instead when it is confirmed and high-confidence.",
annotations(read_only_hint = true)
)]
pub(crate) async fn mem_get(&self, Parameters(params): Parameters<MemGetParams>) -> String {
self.socket_call("mem_get", json!({ "key": params.key }))
.await
}
#[rmcp::tool(
name = "mem_query",
description = "Search the mati knowledge store. Use mode \"text\" for BM25 full-text search, mode \"tag\" to filter by tag, or mode \"graph\" for a 1-hop traversal from a seed key.",
annotations(read_only_hint = true)
)]
pub(crate) async fn mem_query(&self, Parameters(params): Parameters<MemQueryParams>) -> String {
self.socket_call(
"mem_query",
json!({ "query": params.query, "mode": params.mode, "limit": params.limit }),
)
.await
}
#[rmcp::tool(
name = "mem_bootstrap",
description = "Assemble a compact context packet for the current coding session from relevant gotchas, file records, and decisions. Call this at session start.",
annotations(read_only_hint = true)
)]
pub(crate) async fn mem_bootstrap(
&self,
Parameters(params): Parameters<MemBootstrapParams>,
) -> String {
self.socket_call(
"mem_bootstrap",
json!({ "context_files": params.context_files }),
)
.await
}
#[rmcp::tool(
name = "mem_set",
description = "Write, confirm, or delete a knowledge record. Actions: \"write\" (default) creates/updates a record, \"confirm\" activates a gotcha for hook enforcement, \"delete\" tombstones a gotcha.",
annotations(
read_only_hint = false,
destructive_hint = true,
idempotent_hint = false
)
)]
pub(crate) async fn mem_set(&self, Parameters(params): Parameters<MemSetParams>) -> String {
match build_mem_set_command(¶ms) {
Ok(cmd) => self.socket_call_typed(cmd).await,
Err(error) => json!({ "error": error }).to_string(),
}
}
}
fn build_mem_set_command(params: &MemSetParams) -> Result<Command, String> {
match params.action.as_str() {
"confirm" => {
if !params.key.starts_with("gotcha:") {
return Err("confirm action only applies to gotcha: keys".into());
}
Ok(Command::GotchaConfirm(GotchaConfirmInput {
key: params.key.clone(),
}))
}
"delete" => {
if !params.key.starts_with("gotcha:") {
return Err("delete action only applies to gotcha: keys".into());
}
Ok(Command::GotchaTombstone(GotchaTombstoneInput {
key: params.key.clone(),
}))
}
"write" | "" => build_mem_set_write_command(params),
other => Err(format!(
"unknown action: {other}. Valid: write, confirm, delete"
)),
}
}
fn build_mem_set_write_command(params: &MemSetParams) -> Result<Command, String> {
let payload = match ¶ms.payload {
serde_json::Value::String(s) => {
serde_json::from_str::<serde_json::Value>(s).unwrap_or_else(|_| params.payload.clone())
}
other => other.clone(),
};
let priority = parse_protocol_priority(¶ms.priority);
if let Some(stripped) = params.key.strip_prefix("gotcha:") {
if stripped.is_empty() {
return Err("gotcha key must not be just the prefix".into());
}
let rule = field_string(&payload, "rule")
.ok_or_else(|| "gotcha payload requires non-empty 'rule'".to_string())?;
let reason = field_string(&payload, "reason")
.ok_or_else(|| "gotcha payload requires non-empty 'reason'".to_string())?;
let severity = parse_protocol_severity(payload.get("severity").and_then(|v| v.as_str()));
let affected_files = field_string_list(&payload, "affected_files");
let ref_url = payload
.get("ref_url")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
return Ok(Command::GotchaUpsert(GotchaDraftInput {
key: params.key.clone(),
rule,
reason,
severity,
affected_files,
ref_url,
tags: params.tags.clone(),
priority,
source: None,
}));
}
if let Some(slug) = params.key.strip_prefix("decision:") {
if slug.is_empty() {
return Err("decision key must not be just the prefix".into());
}
let summary = field_string(&payload, "summary")
.ok_or_else(|| "decision payload requires non-empty 'summary'".to_string())?;
let rationale = field_string(&payload, "rationale")
.ok_or_else(|| "decision payload requires non-empty 'rationale'".to_string())?;
return Ok(Command::DecisionUpsert(DecisionUpsertInput {
slug: slug.to_string(),
value: params.value.clone(),
summary,
rationale,
tags: params.tags.clone(),
priority,
}));
}
if let Some(stripped) = params.key.strip_prefix("dev_note:") {
if stripped.is_empty() {
return Err("dev_note key must not be just the prefix".into());
}
if params.value.is_empty() {
return Err("dev_note requires non-empty value".into());
}
return Ok(Command::DevNoteUpsert(DevNoteUpsertInput {
key: Some(params.key.clone()),
text: params.value.clone(),
tags: params.tags.clone(),
priority,
}));
}
Err("mem_set write requires key with gotcha:/decision:/dev_note: prefix".into())
}
fn field_string(payload: &serde_json::Value, field: &str) -> Option<String> {
payload
.get(field)
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.filter(|s| !s.is_empty())
}
fn field_string_list(payload: &serde_json::Value, field: &str) -> Vec<String> {
payload
.get(field)
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(|s| s.to_string()))
.collect()
})
.unwrap_or_default()
}
fn parse_protocol_priority(s: &str) -> proto::Priority {
match s {
"Critical" | "critical" => proto::Priority::Critical,
"High" | "high" => proto::Priority::High,
"Low" | "low" => proto::Priority::Low,
_ => proto::Priority::Normal,
}
}
fn parse_protocol_severity(s: Option<&str>) -> proto::Severity {
match s.map(|s| s.to_ascii_lowercase()).as_deref() {
Some("critical") => proto::Severity::Critical,
Some("high") => proto::Severity::High,
Some("low") => proto::Severity::Low,
_ => proto::Severity::Normal,
}
}
fn is_injectable_gotcha(r: &Record) -> bool {
if !matches!(r.lifecycle, RecordLifecycle::Active) {
return false;
}
if r.staleness.tier == StalenessTier::Tombstone {
return false;
}
if r.quality.value < 0.4 {
return false;
}
if let Some(gotcha) = r.payload_as::<GotchaRecord>() {
if gotcha.confirmed {
return true;
}
}
r.key.starts_with("gotcha:cochange:")
|| r.key.starts_with("gotcha:revert:")
|| r.key.starts_with("gotcha:ownership:")
}
pub(crate) fn resolve_existing_for_write(
store_result: anyhow::Result<Option<Record>>,
) -> Result<Option<Record>, String> {
match store_result {
Ok(record) => Ok(record),
Err(e) => Err(serde_json::json!({
"error": format!("store read failed \u{2014} refusing to write: {e}")
})
.to_string()),
}
}
pub async fn assemble_context_packet(
store: &crate::store::Store,
graph: &Graph,
context_files: &[String],
) -> anyhow::Result<ContextPacket> {
let stage = store.get("stage:current").await?;
let mut file_records = Vec::new();
let mut context_gotcha_keys = HashSet::new();
let mut decision_keys = HashSet::new();
let mut unconfirmed_candidates = Vec::new();
let mut stale_warnings: Vec<String> = Vec::new();
let mut seen_stale_keys: HashSet<String> = HashSet::new();
for file_path in context_files {
let file_key = if file_path.starts_with("file:") {
file_path.clone()
} else {
format!("file:{file_path}")
};
if let Ok(Some(record)) = store.get(&file_key).await {
if record.staleness.tier == StalenessTier::Tombstone
|| !matches!(record.lifecycle, RecordLifecycle::Active)
{
continue;
}
match record.staleness.tier {
StalenessTier::Stale => {
let path = file_key.strip_prefix("file:").unwrap_or(&file_key);
if seen_stale_keys.insert(file_key.clone()) {
stale_warnings.push(format!(
"`{path}` record is stale (staleness {:.2}) — verify before trusting",
record.staleness.value
));
}
}
StalenessTier::Liability => {
let path = file_key.strip_prefix("file:").unwrap_or(&file_key);
if seen_stale_keys.insert(file_key.clone()) {
stale_warnings.push(format!(
"`{path}` record is a liability (staleness {:.2}) — do not trust, read the file",
record.staleness.value
));
}
}
_ => {}
}
if let Some(fr) = record.payload_as::<FileRecord>() {
for key in &fr.gotcha_keys {
context_gotcha_keys.insert(key.clone());
}
let is_nudge_candidate = record.access_count >= 3 && fr.gotcha_keys.is_empty();
file_records.push(fr);
if is_nudge_candidate {
unconfirmed_candidates.push(file_key.clone());
}
}
}
for key in graph.neighbors(&file_key, &EdgeKind::HasGotcha) {
context_gotcha_keys.insert(key);
}
for imported in graph.neighbors(&file_key, &EdgeKind::Imports) {
for key in graph.neighbors(&imported, &EdgeKind::HasGotcha) {
context_gotcha_keys.insert(key);
}
}
for key in graph.neighbors(&file_key, &EdgeKind::AffectedBy) {
decision_keys.insert(key);
}
}
let mut confirmed_gotchas: Vec<Record> = if context_files.is_empty() {
let all_gotchas = store.scan_prefix("gotcha:").await?;
all_gotchas
.into_iter()
.filter(is_injectable_gotcha)
.collect()
} else {
let mut gotchas = Vec::with_capacity(context_gotcha_keys.len());
for key in &context_gotcha_keys {
if let Ok(Some(record)) = store.get(key).await {
if is_injectable_gotcha(&record) {
gotchas.push(record);
}
}
}
gotchas
};
{
let now = chrono::Utc::now();
for days_ago in 0..2 {
let date = (now - chrono::Duration::days(days_ago)).format("%Y-%m-%d");
let review_key = format!("analytics:stale_review_{date}");
if let Ok(Some(record)) = store.get(&review_key).await {
if let Some(payload) = record.payload_as::<StaleReviewPayload>() {
for entry in &payload.entries {
if seen_stale_keys.insert(entry.key.clone()) {
let path = entry.key.strip_prefix("file:").unwrap_or(&entry.key);
stale_warnings.push(format!(
"`{path}` staleness {:.2} ({:?}) — review recommended",
entry.staleness_value, entry.tier
));
}
}
}
}
}
}
let mut related_decisions = Vec::new();
for key in &decision_keys {
if let Ok(Some(record)) = store.get(key).await {
related_decisions.push(record);
}
}
if related_decisions.is_empty() {
if let Ok(mut all_decisions) = store.scan_prefix("decision:").await {
all_decisions.retain(|r| matches!(r.lifecycle, RecordLifecycle::Active));
all_decisions.sort_by(|a, b| {
b.confidence
.value
.partial_cmp(&a.confidence.value)
.unwrap_or(std::cmp::Ordering::Equal)
});
const DECISION_FALLBACK_LIMIT: usize = 5;
related_decisions = all_decisions
.into_iter()
.take(DECISION_FALLBACK_LIMIT)
.collect();
}
}
confirmed_gotchas.sort_by(|a, b| {
let score_a = a.confidence.value * priority_weight(&a.priority);
let score_b = b.confidence.value * priority_weight(&b.priority);
score_b
.partial_cmp(&score_a)
.unwrap_or(std::cmp::Ordering::Equal)
});
let critical_gotchas: Vec<Record> = confirmed_gotchas
.into_iter()
.filter(|r| r.quality.tier != QualityTier::Suppressed)
.collect();
let available_tokens = TOKEN_BUDGET - VECTOR_B_TOKENS;
let mut sections = Vec::new();
let mut used_tokens = 0;
if let Some(ref stage_record) = stage {
let section = format!("## Current Stage\n{}\n", stage_record.value);
let tokens = estimate_tokens(§ion);
if used_tokens + tokens <= available_tokens {
sections.push(section);
used_tokens += tokens;
}
}
if !critical_gotchas.is_empty() {
let mut gotcha_section = String::from("## Gotchas\n");
for record in &critical_gotchas {
if record.key.starts_with("gotcha:cochange:") {
continue;
}
let caveat = if record.staleness.tier == StalenessTier::Liability {
" [STALE — verify]"
} else if record.quality.tier == QualityTier::Poor {
" [LOW QUALITY — verify]"
} else {
""
};
let line = format!("- **{}**{}: {}\n", record.key, caveat, record.value);
let tokens = estimate_tokens(&line);
if used_tokens + tokens > available_tokens {
break;
}
gotcha_section.push_str(&line);
used_tokens += tokens;
}
let mut cochange_map: std::collections::BTreeMap<String, Vec<(String, String)>> =
std::collections::BTreeMap::new();
for record in &critical_gotchas {
if !record.key.starts_with("gotcha:cochange:") {
continue;
}
if let Some(pair) = record.key.strip_prefix("gotcha:cochange:") {
if let Some((src, tgt)) = pair.split_once('|') {
let pct = record
.value
.rfind('(')
.and_then(|i| {
record.value[i + 1..]
.find(')')
.map(|j| &record.value[i + 1..i + 1 + j])
})
.unwrap_or("?");
cochange_map
.entry(src.to_string())
.or_default()
.push((tgt.to_string(), pct.to_string()));
}
}
}
if !cochange_map.is_empty() {
let all_pairs: Vec<String> = cochange_map
.iter()
.flat_map(|(src, targets)| {
targets
.iter()
.map(move |(tgt, pct)| format!("{src}\u{2194}{tgt} ({pct})"))
})
.collect();
let total = all_pairs.len();
let display: Vec<&str> = all_pairs.iter().take(10).map(|s| s.as_str()).collect();
let suffix = if total > 10 {
format!(", +{} more", total - 10)
} else {
String::new()
};
let line = format!("- **Co-change partners**: {}{suffix}\n", display.join(", "));
let tokens = estimate_tokens(&line);
if used_tokens + tokens <= available_tokens {
gotcha_section.push_str(&line);
used_tokens += tokens;
}
}
if gotcha_section.len() > "## Gotchas\n".len() {
sections.push(gotcha_section);
}
}
if !file_records.is_empty() {
let mut file_section = String::from("## Context Files\n");
for fr in &file_records {
if fr.purpose.is_empty() {
continue;
}
let line = format!("- **{}**: {}\n", fr.path, fr.purpose);
let tokens = estimate_tokens(&line);
if used_tokens + tokens > available_tokens {
break;
}
file_section.push_str(&line);
used_tokens += tokens;
}
if file_section.len() > "## Context Files\n".len() {
sections.push(file_section);
}
}
{
use crate::analysis::blast_radius::BlastTier;
let mut impact_files: Vec<(&FileRecord, f32)> = file_records
.iter()
.filter_map(|fr| {
fr.blast_radius.as_ref().and_then(|br| {
if br.tier == BlastTier::Isolated {
None
} else {
Some((fr, br.score))
}
})
})
.collect();
impact_files.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
if !impact_files.is_empty() {
let mut impact_section = String::from("## Highest Impact Files\n");
for (fr, _score) in impact_files.iter().take(3) {
let br = fr
.blast_radius
.as_ref()
.expect("filter_map above kept only files with Some(blast_radius)");
let line = format!(
"- `{}`: {} direct importers ({})\n",
fr.path,
br.direct,
br.tier.label(),
);
let tokens = estimate_tokens(&line);
if used_tokens + tokens > available_tokens {
break;
}
impact_section.push_str(&line);
used_tokens += tokens;
}
if impact_section.len() > "## Highest Impact Files\n".len() {
sections.push(impact_section);
}
}
}
if !stale_warnings.is_empty() {
let mut stale_section = String::from("## Stale Warnings\n");
for warning in &stale_warnings {
let line = format!("- {warning}\n");
let tokens = estimate_tokens(&line);
if used_tokens + tokens > available_tokens {
break;
}
stale_section.push_str(&line);
used_tokens += tokens;
}
if stale_section.len() > "## Stale Warnings\n".len() {
sections.push(stale_section);
}
}
if !related_decisions.is_empty() {
let mut dec_section = String::from("## Decisions\n");
for record in &related_decisions {
let line = format!("- **{}**: {}\n", record.key, record.value);
let tokens = estimate_tokens(&line);
if used_tokens + tokens > available_tokens {
break;
}
dec_section.push_str(&line);
used_tokens += tokens;
}
if dec_section.len() > "## Decisions\n".len() {
sections.push(dec_section);
}
}
if !unconfirmed_candidates.is_empty() {
let mut nudge_section = String::from("## Suggested Actions\n");
for key in &unconfirmed_candidates {
let path = key.strip_prefix("file:").unwrap_or(key);
let line = format!(
"- `{path}` is read frequently but has no recorded gotchas. The developer may want to run `mati gotcha add {path}`.\n"
);
let tokens = estimate_tokens(&line);
if used_tokens + tokens > available_tokens {
break;
}
nudge_section.push_str(&line);
used_tokens += tokens;
}
if nudge_section.len() > "## Suggested Actions\n".len() {
sections.push(nudge_section);
}
}
let mut injection_string = sections.join("\n");
injection_string.push_str(VECTOR_B);
let token_estimate = estimate_tokens(&injection_string) as u32;
Ok(ContextPacket {
stage,
critical_gotchas,
file_records,
related_decisions,
recent_session: None,
token_estimate,
stale_warnings,
unconfirmed_candidates,
knowledge_gaps: vec![],
compliance_rate: None,
injection_string,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::store::record::*;
use crate::store::Store;
use tempfile::TempDir;
fn device_id() -> uuid::Uuid {
uuid::Uuid::nil()
}
fn now() -> u64 {
1_700_000_000
}
fn make_record(key: &str, value: &str, category: Category, quality_value: f32) -> Record {
Record {
key: key.to_string(),
value: value.to_string(),
category,
priority: Priority::Normal,
tags: vec![],
created_at: now(),
updated_at: now(),
ref_url: None,
staleness: StalenessScore::fresh(),
lifecycle: RecordLifecycle::Active,
version: RecordVersion {
device_id: device_id(),
logical_clock: 1,
wall_clock: now(),
},
quality: QualityScore {
value: quality_value,
tier: QualityScore::tier_from_value(quality_value),
signals: vec![],
computed_at: now(),
},
access_count: 0,
last_accessed: 0,
source: RecordSource::DeveloperManual,
confidence: ConfidenceScore {
value: 0.8,
confirmation_count: 1,
contributor_count: 1,
last_challenged: None,
challenge_count: 0,
},
gap_analysis_score: 0.0,
payload: Some(serde_json::json!({})),
}
}
fn make_gotcha_record(key: &str, rule: &str, confirmed: bool, quality_value: f32) -> Record {
let gotcha = GotchaRecord {
rule: rule.to_string(),
reason: "test reason".to_string(),
severity: Priority::High,
affected_files: vec![],
ref_url: None,
discovered_session: now(),
confirmed,
};
let mut record = make_record(key, rule, Category::Gotcha, quality_value);
record.payload = serde_json::to_value(&gotcha).ok();
record
}
fn handler_test_ctx() -> crate::mcp::dispatch_v2::RequestContext {
crate::mcp::dispatch_v2::RequestContext {
peer: crate::mcp::metadata::PeerContext {
uid: 501,
pid: Some(99999),
},
daemon_session: uuid::Uuid::nil(),
repo_root: std::path::PathBuf::new(),
}
}
async fn call_mem_get(
graph_arc: &std::sync::Arc<tokio::sync::RwLock<crate::graph::Graph>>,
key: &str,
) -> String {
let ctx = handler_test_ctx();
let input = crate::mcp::protocol::MemGetInput {
key: key.to_string(),
};
let g = graph_arc.read().await;
match crate::mcp::handlers::handle_mem_get(
g.store(),
graph_arc,
&ctx,
uuid::Uuid::new_v4(),
&input,
)
.await
{
Ok(v) => serde_json::to_string_pretty(&v).unwrap_or_else(|_| "{}".into()),
Err((_code, msg)) => format!("{{\"error\": \"{}\"}}", msg.replace('"', "\\\"")),
}
}
async fn call_mem_query(
graph_arc: &std::sync::Arc<tokio::sync::RwLock<crate::graph::Graph>>,
query: &str,
mode: crate::mcp::protocol::QueryMode,
limit: u32,
) -> String {
let input = crate::mcp::protocol::MemQueryInput {
query: query.to_string(),
mode,
limit,
};
let g = graph_arc.read().await;
match crate::mcp::handlers::handle_mem_query(g.store(), &g, &input).await {
Ok(v) => serde_json::to_string_pretty(&v).unwrap_or_else(|_| "{}".into()),
Err((_code, msg)) => format!("{{\"error\": \"{}\"}}", msg.replace('"', "\\\"")),
}
}
async fn call_mem_bootstrap(
graph_arc: &std::sync::Arc<tokio::sync::RwLock<crate::graph::Graph>>,
context_files: Vec<String>,
) -> String {
let ctx = handler_test_ctx();
let input = crate::mcp::protocol::MemBootstrapInput { context_files };
let g = graph_arc.read().await;
match crate::mcp::handlers::handle_mem_bootstrap(
g.store(),
&g,
graph_arc,
&ctx,
uuid::Uuid::new_v4(),
&input,
)
.await
{
Ok(s) => s,
Err((_code, msg)) => format!("[mati] bootstrap error: {msg}"),
}
}
async fn call_mem_set(
graph_arc: &std::sync::Arc<tokio::sync::RwLock<crate::graph::Graph>>,
params: crate::mcp::types::MemSetParams,
) -> String {
let ctx = handler_test_ctx();
crate::mcp::handlers::handle_mem_set(graph_arc, &ctx, uuid::Uuid::new_v4(), ¶ms).await
}
#[tokio::test]
async fn mem_get_returns_null_for_nonexistent_key() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_get(&graph_arc, "file:nonexistent.rs").await;
assert_eq!(result, "null");
}
#[tokio::test]
async fn mem_get_returns_record_for_existing_key() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let record = make_record("gotcha:test", "test value", Category::Gotcha, 0.8);
store.put("gotcha:test", &record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_get(&graph_arc, "gotcha:test").await;
assert!(result.contains("gotcha:test"));
assert!(result.contains("test value"));
}
#[tokio::test]
async fn mem_get_blast_radius_warning_for_critical_file() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/core.rs".to_string(),
purpose: "Core module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 100,
last_modified_session: 0,
content_hash: None,
line_count: 0,
blast_radius: Some(crate::analysis::blast_radius::BlastRadius {
direct: 45,
transitive: 10,
score: 48.0,
tier: crate::analysis::blast_radius::BlastTier::Critical,
}),
propagated_staleness: None,
};
let mut record = make_record("file:src/core.rs", "Core module", Category::File, 0.5);
record.payload = serde_json::to_value(&fr).ok();
store.put("file:src/core.rs", &record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_get(&graph_arc, "file:src/core.rs").await;
assert!(
result.contains("HIGH IMPACT FILE"),
"response must contain blast radius warning for critical file, got: {result}"
);
assert!(result.contains("45"), "warning must include direct count");
}
#[tokio::test]
async fn mem_get_no_blast_warning_for_low_file() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/leaf.rs".to_string(),
purpose: "Leaf module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 100,
last_modified_session: 0,
content_hash: None,
line_count: 0,
blast_radius: Some(crate::analysis::blast_radius::BlastRadius {
direct: 2,
transitive: 0,
score: 2.0,
tier: crate::analysis::blast_radius::BlastTier::Low,
}),
propagated_staleness: None,
};
let mut record = make_record("file:src/leaf.rs", "Leaf module", Category::File, 0.5);
record.payload = serde_json::to_value(&fr).ok();
store.put("file:src/leaf.rs", &record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_get(&graph_arc, "file:src/leaf.rs").await;
assert!(
!result.contains("HIGH IMPACT FILE"),
"low blast radius file should NOT have warning"
);
}
#[tokio::test]
async fn mem_get_includes_depth_hint_for_file_records() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/tiny.rs".to_string(),
purpose: "Tiny leaf module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 100,
last_modified_session: 0,
content_hash: None,
line_count: 50,
blast_radius: Some(crate::analysis::blast_radius::BlastRadius {
direct: 0,
transitive: 0,
score: 0.0,
tier: crate::analysis::blast_radius::BlastTier::Isolated,
}),
propagated_staleness: None,
};
let mut record = make_record("file:src/tiny.rs", "Tiny leaf", Category::File, 0.5);
record.payload = serde_json::to_value(&fr).ok();
store.put("file:src/tiny.rs", &record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_get(&graph_arc, "file:src/tiny.rs").await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(
parsed.get("enrichment_depth_hint").and_then(|v| v.as_str()),
Some("fast"),
"tiny isolated file should hint Fast tier; got: {result}"
);
}
#[tokio::test]
async fn mem_get_depth_hint_for_hotspot_is_deep() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/core.rs".to_string(),
purpose: "Core hotspot".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: true,
token_cost_estimate: 5000,
last_modified_session: 0,
content_hash: None,
line_count: 500,
blast_radius: Some(crate::analysis::blast_radius::BlastRadius {
direct: 20,
transitive: 30,
score: 35.0,
tier: crate::analysis::blast_radius::BlastTier::High,
}),
propagated_staleness: None,
};
let mut record = make_record("file:src/core.rs", "Core hotspot", Category::File, 0.5);
record.payload = serde_json::to_value(&fr).ok();
store.put("file:src/core.rs", &record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_get(&graph_arc, "file:src/core.rs").await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(
parsed.get("enrichment_depth_hint").and_then(|v| v.as_str()),
Some("deep"),
"large hotspot file should hint Deep tier; got: {result}"
);
}
#[tokio::test]
async fn mem_get_omits_depth_hint_for_non_file_records() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let record = make_record("gotcha:foo", "rule", Category::Gotcha, 0.6);
store.put("gotcha:foo", &record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_get(&graph_arc, "gotcha:foo").await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(
parsed.get("enrichment_depth_hint").is_none(),
"non-file records should not carry enrichment_depth_hint; got: {result}"
);
}
#[tokio::test]
async fn mem_query_text_mode_returns_results() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let record = make_record(
"gotcha:async-race",
"never use inference in async context",
Category::Gotcha,
0.8,
);
store.put("gotcha:async-race", &record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_query(
&graph_arc,
"inference",
crate::mcp::protocol::QueryMode::Text,
10,
)
.await;
assert!(result.contains("gotcha:async-race"));
}
#[tokio::test]
async fn mem_query_semantic_returns_feature_gate_error() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_query(
&graph_arc,
"test",
crate::mcp::protocol::QueryMode::Semantic,
20,
)
.await;
assert!(
result.contains("--features semantic"),
"semantic mode must surface feature-gate error, got: {result}"
);
}
#[tokio::test]
async fn mem_bootstrap_empty_store_returns_vector_b() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_bootstrap(&graph_arc, vec![]).await;
assert!(result.contains("[mati] Before reading any file"));
assert!(result.contains("mem_get"));
}
#[tokio::test]
async fn mem_bootstrap_token_budget_never_exceeds_2000() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
for i in 0..100 {
let record = make_gotcha_record(
&format!("gotcha:test-{i:03}"),
&format!("This is a very long gotcha rule number {i} with lots of text to fill up the token budget and ensure we test the truncation logic properly"),
true,
0.8,
);
store.put(&record.key, &record).await.unwrap();
}
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &[])
.await
.unwrap();
let tokens = estimate_tokens(&packet.injection_string);
assert!(
tokens <= TOKEN_BUDGET,
"token estimate {tokens} exceeds budget {TOKEN_BUDGET}"
);
}
#[tokio::test]
async fn quality_filter_suppressed_excluded() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let suppressed = make_gotcha_record("gotcha:suppressed", "bad rule", true, 0.10);
store.put("gotcha:suppressed", &suppressed).await.unwrap();
let good = make_gotcha_record("gotcha:good", "good rule", true, 0.80);
store.put("gotcha:good", &good).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &[])
.await
.unwrap();
assert!(
!packet.injection_string.contains("gotcha:suppressed"),
"suppressed gotcha must not appear in injection"
);
}
#[tokio::test]
async fn quality_filter_poor_caveated() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let poor = make_gotcha_record("gotcha:poor", "poor rule", true, 0.30);
store.put("gotcha:poor", &poor).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &[])
.await
.unwrap();
if packet.injection_string.contains("gotcha:poor") {
assert!(
packet.injection_string.contains("LOW QUALITY"),
"poor quality gotcha must be caveated"
);
}
}
#[tokio::test]
async fn assemble_context_packet_with_context_files_does_graph_traversal() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let gotcha = make_gotcha_record("gotcha:important", "do not use unwrap", true, 0.80);
store.put("gotcha:important", &gotcha).await.unwrap();
let file_record = make_record("file:src/main.rs", "{}", Category::File, 0.5);
store.put("file:src/main.rs", &file_record).await.unwrap();
let mut graph = Graph::load(store).await.unwrap();
graph
.add_edge("file:src/main.rs", EdgeKind::HasGotcha, "gotcha:important")
.await
.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &["src/main.rs".to_string()])
.await
.unwrap();
assert!(
packet.injection_string.contains("gotcha:important")
|| packet
.critical_gotchas
.iter()
.any(|g| g.key == "gotcha:important"),
"graph-connected gotcha must appear in context packet"
);
}
#[tokio::test]
async fn assemble_context_packet_excludes_unrelated_gotchas_for_context_files() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let relevant = make_gotcha_record("gotcha:relevant", "do not use unwrap", true, 0.80);
let unrelated = make_gotcha_record("gotcha:unrelated", "keep retries bounded", true, 0.80);
store.put("gotcha:relevant", &relevant).await.unwrap();
store.put("gotcha:unrelated", &unrelated).await.unwrap();
let file_record = make_record("file:src/main.rs", "{}", Category::File, 0.5);
store.put("file:src/main.rs", &file_record).await.unwrap();
let mut graph = Graph::load(store).await.unwrap();
graph
.add_edge("file:src/main.rs", EdgeKind::HasGotcha, "gotcha:relevant")
.await
.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &["src/main.rs".to_string()])
.await
.unwrap();
assert!(
packet
.critical_gotchas
.iter()
.any(|g| g.key == "gotcha:relevant"),
"graph-connected gotcha must remain in context packet"
);
assert!(
!packet
.critical_gotchas
.iter()
.any(|g| g.key == "gotcha:unrelated"),
"unrelated gotcha must not be injected for scoped bootstrap"
);
assert!(
!packet.injection_string.contains("gotcha:unrelated"),
"injection string must not mention unrelated gotchas"
);
}
#[tokio::test]
async fn bootstrap_surfaces_confirmed_gotcha_when_graph_edge_missing() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let gotcha = make_gotcha_record(
"gotcha:never-remove-rate-limit",
"Never remove the rate limit check on incoming pipeline events because \
removing it caused a cascade failure in staging",
true,
0.80,
);
store
.put("gotcha:never-remove-rate-limit", &gotcha)
.await
.unwrap();
let file_record = {
let fr = FileRecord {
path: "src/pipeline/prefilter.rs".to_string(),
purpose: String::new(), entry_points: vec![],
imports: vec![],
gotcha_keys: vec!["gotcha:never-remove-rate-limit".to_string()],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 18,
last_author: Some("dev".to_string()),
is_hotspot: true,
token_cost_estimate: 0,
last_modified_session: now(),
content_hash: None,
line_count: 0,
blast_radius: None,
propagated_staleness: None,
};
let mut r = make_record(
"file:src/pipeline/prefilter.rs",
"",
Category::File,
0.10, );
r.payload = serde_json::to_value(&fr).ok();
r
};
store
.put("file:src/pipeline/prefilter.rs", &file_record)
.await
.unwrap();
let graph = Graph::load(store).await.unwrap();
assert_eq!(
graph.neighbors("file:src/pipeline/prefilter.rs", &EdgeKind::HasGotcha),
Vec::<String>::new(),
"test setup: graph must have no HasGotcha edge"
);
let packet = assemble_context_packet(
graph.store(),
&graph,
&["src/pipeline/prefilter.rs".to_string()],
)
.await
.unwrap();
assert!(
packet
.critical_gotchas
.iter()
.any(|g| g.key == "gotcha:never-remove-rate-limit"),
"bootstrap must surface confirmed gotcha even when graph edge is missing"
);
assert!(
packet
.injection_string
.contains("gotcha:never-remove-rate-limit"),
"injection string must include the gotcha"
);
}
#[tokio::test]
async fn bootstrap_low_confidence_file_with_no_gotchas_returns_minimal_packet() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let file_record = {
let fr = FileRecord {
path: "src/empty.rs".to_string(),
purpose: String::new(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 1,
last_author: None,
is_hotspot: false,
token_cost_estimate: 0,
last_modified_session: now(),
content_hash: None,
line_count: 0,
blast_radius: None,
propagated_staleness: None,
};
let mut r = make_record("file:src/empty.rs", "", Category::File, 0.10);
r.payload = serde_json::to_value(&fr).ok();
r
};
store.put("file:src/empty.rs", &file_record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &["src/empty.rs".to_string()])
.await
.unwrap();
assert!(
packet.critical_gotchas.is_empty(),
"no gotchas should be surfaced for a file with no linked gotchas"
);
assert!(
!packet.injection_string.contains("gotcha:"),
"injection string must not mention any gotcha keys"
);
}
#[tokio::test]
async fn nudge_shown_for_hot_file_with_no_gotchas() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/hot.rs".to_string(),
purpose: "Hot module".to_string(),
entry_points: vec!["run".to_string()],
imports: vec![],
gotcha_keys: vec![], decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 10,
last_author: None,
is_hotspot: true,
token_cost_estimate: 100,
last_modified_session: now(),
content_hash: None,
line_count: 0,
blast_radius: None,
propagated_staleness: None,
};
let mut file_record = make_record("file:src/hot.rs", &fr.purpose, Category::File, 0.5);
file_record.payload = serde_json::to_value(&fr).ok();
file_record.access_count = 5; store.put("file:src/hot.rs", &file_record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &["src/hot.rs".to_string()])
.await
.unwrap();
assert!(
packet
.unconfirmed_candidates
.contains(&"file:src/hot.rs".to_string()),
"hot file with no gotchas should be in unconfirmed_candidates"
);
assert!(
packet.injection_string.contains("Suggested Actions"),
"nudge section should appear in injection string"
);
assert!(
packet.injection_string.contains("mati gotcha add"),
"nudge should suggest gotcha add command"
);
}
#[tokio::test]
async fn no_nudge_for_file_with_low_access_count() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/cold.rs".to_string(),
purpose: "Cold module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 50,
last_modified_session: now(),
content_hash: None,
line_count: 0,
blast_radius: None,
propagated_staleness: None,
};
let mut file_record = make_record("file:src/cold.rs", &fr.purpose, Category::File, 0.5);
file_record.payload = serde_json::to_value(&fr).ok();
file_record.access_count = 1; store.put("file:src/cold.rs", &file_record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &["src/cold.rs".to_string()])
.await
.unwrap();
assert!(
packet.unconfirmed_candidates.is_empty(),
"low-access file should not trigger nudge"
);
}
#[tokio::test]
async fn no_nudge_for_file_with_gotchas() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/covered.rs".to_string(),
purpose: "Covered module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec!["gotcha:existing".to_string()],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 10,
last_author: None,
is_hotspot: true,
token_cost_estimate: 100,
last_modified_session: now(),
content_hash: None,
line_count: 0,
blast_radius: None,
propagated_staleness: None,
};
let mut file_record = make_record(
"file:src/covered.rs",
&serde_json::to_string(&fr).unwrap(),
Category::File,
0.5,
);
file_record.access_count = 10;
store
.put("file:src/covered.rs", &file_record)
.await
.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet =
assemble_context_packet(graph.store(), &graph, &["src/covered.rs".to_string()])
.await
.unwrap();
assert!(
packet.unconfirmed_candidates.is_empty(),
"file with gotchas should not trigger nudge"
);
}
#[tokio::test]
async fn tombstone_gotcha_excluded_from_bootstrap() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let mut gotcha = make_gotcha_record("gotcha:tombstone", "tombstone rule", true, 0.80);
gotcha.staleness = StalenessScore {
value: 0.95,
tier: StalenessTier::Tombstone,
signals: vec![],
computed_at: now(),
last_record_sha: String::new(),
};
store.put("gotcha:tombstone", &gotcha).await.unwrap();
let good = make_gotcha_record("gotcha:good", "good rule", true, 0.80);
store.put("gotcha:good", &good).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &[])
.await
.unwrap();
assert!(
!packet.injection_string.contains("gotcha:tombstone"),
"tombstone gotcha must not appear in injection"
);
assert!(
packet.injection_string.contains("gotcha:good"),
"normal gotcha should appear"
);
}
#[tokio::test]
async fn liability_gotcha_gets_stale_caveat() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let mut gotcha = make_gotcha_record("gotcha:liability", "liability rule", true, 0.80);
gotcha.staleness = StalenessScore {
value: 0.75,
tier: StalenessTier::Liability,
signals: vec![],
computed_at: now(),
last_record_sha: String::new(),
};
store.put("gotcha:liability", &gotcha).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &[])
.await
.unwrap();
if packet.injection_string.contains("gotcha:liability") {
assert!(
packet.injection_string.contains("STALE"),
"liability gotcha must have STALE caveat"
);
}
}
#[tokio::test]
async fn stale_file_generates_warning() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/stale.rs".to_string(),
purpose: "Stale module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 50,
last_modified_session: now(),
content_hash: None,
line_count: 0,
blast_radius: None,
propagated_staleness: None,
};
let mut file_record = make_record(
"file:src/stale.rs",
&serde_json::to_string(&fr).unwrap(),
Category::File,
0.5,
);
file_record.staleness = StalenessScore {
value: 0.55,
tier: StalenessTier::Stale,
signals: vec![],
computed_at: now(),
last_record_sha: String::new(),
};
store.put("file:src/stale.rs", &file_record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &["src/stale.rs".to_string()])
.await
.unwrap();
assert!(
!packet.stale_warnings.is_empty(),
"stale file should generate a warning"
);
assert!(
packet.stale_warnings.iter().any(|w| w.contains("stale.rs")),
"warning should mention the stale file"
);
}
#[tokio::test]
async fn tombstone_file_excluded_from_traversal() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/dead.rs".to_string(),
purpose: "Dead module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 50,
last_modified_session: now(),
content_hash: None,
line_count: 0,
blast_radius: None,
propagated_staleness: None,
};
let mut file_record = make_record(
"file:src/dead.rs",
&serde_json::to_string(&fr).unwrap(),
Category::File,
0.5,
);
file_record.staleness = StalenessScore {
value: 0.95,
tier: StalenessTier::Tombstone,
signals: vec![],
computed_at: now(),
last_record_sha: String::new(),
};
store.put("file:src/dead.rs", &file_record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &["src/dead.rs".to_string()])
.await
.unwrap();
assert!(
packet.file_records.is_empty(),
"tombstone file should not appear in file_records"
);
}
#[tokio::test]
async fn stale_warnings_deduplicated() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/dup.rs".to_string(),
purpose: "Dup module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 50,
last_modified_session: now(),
content_hash: None,
line_count: 0,
blast_radius: None,
propagated_staleness: None,
};
let mut file_record = make_record(
"file:src/dup.rs",
&serde_json::to_string(&fr).unwrap(),
Category::File,
0.5,
);
file_record.staleness = StalenessScore {
value: 0.55,
tier: StalenessTier::Stale,
signals: vec![],
computed_at: now(),
last_record_sha: String::new(),
};
store.put("file:src/dup.rs", &file_record).await.unwrap();
let review_payload = StaleReviewPayload {
session_timestamp: now(),
entries: vec![StaleReviewEntry {
key: "file:src/dup.rs".to_string(),
staleness_value: 0.55,
tier: StalenessTier::Stale,
last_updated: now(),
signals: vec!["stale".to_string()],
}],
};
let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
let review_key = format!("analytics:stale_review_{today}");
let review_record = make_record(
&review_key,
&serde_json::to_string(&review_payload).unwrap(),
Category::Analytics,
0.5,
);
store.put(&review_key, &review_record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &["src/dup.rs".to_string()])
.await
.unwrap();
let dup_count = packet
.stale_warnings
.iter()
.filter(|w| w.contains("dup.rs"))
.count();
assert_eq!(
dup_count, 1,
"same key should not produce duplicate warnings"
);
}
#[tokio::test]
async fn stale_warnings_section_before_decisions() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr = FileRecord {
path: "src/stale.rs".to_string(),
purpose: "Stale".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 50,
last_modified_session: now(),
content_hash: None,
line_count: 0,
blast_radius: None,
propagated_staleness: None,
};
let mut file_record = make_record(
"file:src/stale.rs",
&serde_json::to_string(&fr).unwrap(),
Category::File,
0.5,
);
file_record.staleness = StalenessScore {
value: 0.55,
tier: StalenessTier::Stale,
signals: vec![],
computed_at: now(),
last_record_sha: String::new(),
};
store.put("file:src/stale.rs", &file_record).await.unwrap();
let decision = make_record("decision:arch", "Use SurrealKV", Category::Decision, 0.8);
store.put("decision:arch", &decision).await.unwrap();
let mut graph = Graph::load(store).await.unwrap();
graph
.add_edge("file:src/stale.rs", EdgeKind::AffectedBy, "decision:arch")
.await
.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &["src/stale.rs".to_string()])
.await
.unwrap();
let stale_pos = packet.injection_string.find("## Stale Warnings");
let dec_pos = packet.injection_string.find("## Decisions");
if let (Some(s), Some(d)) = (stale_pos, dec_pos) {
assert!(s < d, "Stale Warnings section must appear before Decisions");
}
}
#[tokio::test]
async fn unconfirmed_gotcha_never_injected() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let unconfirmed = make_gotcha_record("gotcha:unconfirmed", "unconfirmed rule", false, 0.80);
store.put("gotcha:unconfirmed", &unconfirmed).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &[])
.await
.unwrap();
assert!(
!packet.injection_string.contains("gotcha:unconfirmed"),
"unconfirmed gotcha must never be injected"
);
}
#[tokio::test]
async fn empty_store_returns_only_vector_b() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(graph.store(), &graph, &[])
.await
.unwrap();
assert!(packet.injection_string.contains("[mati] Before reading"));
assert!(packet.critical_gotchas.is_empty());
assert!(packet.file_records.is_empty());
assert!(packet.stale_warnings.is_empty());
assert!(packet.related_decisions.is_empty());
}
#[tokio::test]
async fn mem_set_rejects_file_writes_on_direct_path() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "file:src/main.rs".to_string(),
value: "Handles CLI dispatch and binary entry point".to_string(),
category: "File".to_string(),
payload: serde_json::json!({
"path": "src/main.rs",
"purpose": "Handles CLI dispatch and binary entry point"
}),
tags: vec!["entry-point".to_string()],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
let error = parsed["error"]
.as_str()
.unwrap_or_else(|| panic!("expected an error envelope, got: {result}"));
assert!(
error.contains("must start with"),
"file: write must be rejected by the prefix gate, got: {error}"
);
assert_eq!(parsed.get("ok"), None, "no record should be written");
let graph = graph_arc.read().await;
assert!(
graph
.store()
.get("file:src/main.rs")
.await
.unwrap()
.is_none(),
"file: write must not create a record"
);
}
#[tokio::test]
async fn mem_set_writes_gotcha_with_quality_score() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams { action: "write".to_string(),
key: "gotcha:always-use-idempotency-keys".to_string(),
value: "Always pass idempotency_key to Stripe charge creation because duplicate charges cause customer refund disputes".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "Always pass idempotency_key to Stripe charge creation",
"reason": "duplicate charges cause customer refund disputes",
"severity": "Critical",
"affected_files": ["src/payments/stripe.go"],
"ref_url": null,
"discovered_session": 0,
"confirmed": false
}),
tags: vec![],
priority: "Critical".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["ok"], true);
assert!(parsed["quality"].as_f64().unwrap() > 0.2);
let graph = graph_arc.read().await;
let record = graph
.store()
.get("gotcha:always-use-idempotency-keys")
.await
.unwrap()
.unwrap();
assert_eq!(record.priority, Priority::Critical);
assert_eq!(record.source, RecordSource::ClaudeEnrich);
}
#[tokio::test]
async fn mem_set_rejects_invalid_key_prefix() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "session:12345".to_string(),
value: "test".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
let error = parsed["error"].as_str().unwrap();
assert!(error.contains("must start with"), "got: {error}");
assert!(!error.contains("file:"), "got: {error}");
}
#[tokio::test]
async fn mem_set_rejects_invalid_category() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:test".to_string(),
value: "test".to_string(),
category: "Unknown".to_string(),
payload: serde_json::json!({}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(parsed["error"]
.as_str()
.unwrap()
.contains("unknown category"));
}
#[tokio::test]
async fn mem_set_rejects_key_category_mismatch() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:should-fail".to_string(),
value: "test".to_string(),
category: "Decision".to_string(),
payload: serde_json::json!({"summary": "x", "rationale": "y"}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(
parsed["error"]
.as_str()
.unwrap()
.contains("requires category"),
"key-category mismatch must be rejected: {result}"
);
}
#[tokio::test]
async fn mem_set_rejects_new_gotcha_without_rule_and_reason() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:missing-fields".to_string(),
value: "test".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({"severity": "Normal"}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(
parsed["error"]
.as_str()
.unwrap()
.contains("'rule' and 'reason'"),
"new gotcha without rule/reason must be rejected: {result}"
);
}
#[tokio::test]
async fn mem_set_rejects_new_decision_without_summary_rationale() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "decision:incomplete".to_string(),
value: "test".to_string(),
category: "Decision".to_string(),
payload: serde_json::json!({}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(
parsed["error"]
.as_str()
.unwrap()
.contains("'summary' and 'rationale'"),
"new decision without summary/rationale must be rejected: {result}"
);
}
#[tokio::test]
async fn mem_set_rejects_new_dev_note_with_empty_value() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "dev_note:empty".to_string(),
value: "".to_string(),
category: "DevNote".to_string(),
payload: serde_json::json!({}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(
parsed["error"]
.as_str()
.unwrap()
.contains("non-empty value"),
"new dev_note with empty value must be rejected: {result}"
);
}
#[tokio::test]
async fn mem_set_allows_partial_payload_on_update() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:partial-update".to_string(),
value: "test rule because test reason".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "test rule",
"reason": "test reason",
"severity": "Normal",
"affected_files": [],
}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["ok"], true, "initial write must succeed");
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:partial-update".to_string(),
value: "updated rule because updated reason".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({"reason": "updated reason"}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(
parsed["ok"], true,
"partial-payload update must succeed: {result}"
);
}
#[tokio::test]
async fn mem_set_preserves_confirmation_state_on_update() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let mut record = make_gotcha_record(
"gotcha:confirmed-edit-test",
"Always test first",
true,
0.70,
);
record.source = RecordSource::DeveloperManual;
record.confidence = ConfidenceScore {
value: 0.80,
confirmation_count: 1,
contributor_count: 1,
last_challenged: None,
challenge_count: 0,
};
record.tags = vec!["important".to_string()];
store
.put("gotcha:confirmed-edit-test", &record)
.await
.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:confirmed-edit-test".to_string(),
value: "Always test first because untested changes cause regressions".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "Always test first",
"reason": "untested changes cause regressions",
"severity": "High",
"affected_files": ["src/main.rs"],
"ref_url": null,
"discovered_session": 0,
"confirmed": true
}),
tags: vec![], priority: "High".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["ok"], true);
let graph = graph_arc.read().await;
let updated = graph
.store()
.get("gotcha:confirmed-edit-test")
.await
.unwrap()
.unwrap();
assert_eq!(updated.source, RecordSource::DeveloperManual);
assert!(
(updated.confidence.value - 0.80).abs() < 0.01,
"confidence should stay 0.80, got {}",
updated.confidence.value
);
assert_eq!(updated.confidence.confirmation_count, 1);
assert_eq!(
updated.tags,
vec!["important".to_string()],
"tags should be preserved when caller sends empty"
);
}
#[tokio::test]
async fn mem_set_moves_gotcha_links_and_edges_on_edit() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let old_file = Record::layer0_file_stub("file:src/old.rs", device_id(), 1, now());
let new_file = Record::layer0_file_stub("file:src/new.rs", device_id(), 1, now());
store.put("file:src/old.rs", &old_file).await.unwrap();
store.put("file:src/new.rs", &new_file).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:test-move".to_string(),
value: "Always update the paired file because drift breaks the feature".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "Always update the paired file",
"reason": "drift breaks the feature",
"severity": "High",
"affected_files": ["src/old.rs"],
"ref_url": null,
"discovered_session": 0,
"confirmed": false
}),
tags: vec![],
priority: "High".to_string(),
},
)
.await;
call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:test-move".to_string(),
value: "Always update the paired file because drift breaks the feature".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "Always update the paired file",
"reason": "drift breaks the feature",
"severity": "High",
"affected_files": ["src/new.rs"],
"ref_url": null,
"discovered_session": 0,
"confirmed": false
}),
tags: vec![],
priority: "High".to_string(),
},
)
.await;
let graph = graph_arc.read().await;
let old_file = graph.store().get("file:src/old.rs").await.unwrap().unwrap();
let new_file = graph.store().get("file:src/new.rs").await.unwrap().unwrap();
let old_payload = old_file.payload.unwrap();
let new_payload = new_file.payload.unwrap();
assert!(
old_payload["gotcha_keys"]
.as_array()
.map(|arr| arr.is_empty())
.unwrap_or(true),
"old file should no longer reference moved gotcha"
);
assert_eq!(new_payload["gotcha_keys"][0], "gotcha:test-move");
assert!(
!graph
.neighbors("file:src/old.rs", &EdgeKind::HasGotcha)
.contains(&"gotcha:test-move".to_string()),
"old file should not keep stale HasGotcha edge"
);
assert!(
graph
.neighbors("file:src/new.rs", &EdgeKind::HasGotcha)
.contains(&"gotcha:test-move".to_string()),
"new file should gain HasGotcha edge"
);
}
#[tokio::test]
async fn test_query_limit_clamped_to_max() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
for i in 0..60 {
let record = make_record(
&format!("gotcha:clamp-test-{i:03}"),
&format!("clamp test rule number {i}"),
Category::Gotcha,
0.8,
);
store
.put(&format!("gotcha:clamp-test-{i:03}"), &record)
.await
.unwrap();
}
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_query(
&graph_arc,
"clamp test rule",
crate::mcp::protocol::QueryMode::Text,
100, )
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(
parsed.get("error").is_none(),
"query with limit > 50 must not error"
);
let results = parsed.as_array().expect("result should be a JSON array");
assert!(
results.len() <= 50,
"result count {} exceeds MAX_QUERY_LIMIT (50)",
results.len()
);
}
#[tokio::test]
async fn test_mem_set_write_new_key_succeeds() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:store-read-ok-none".to_string(),
value: "Regression rule because regression reason".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "Regression rule",
"reason": "regression reason",
"severity": "Normal"
}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(
parsed["ok"], true,
"writing a new key must succeed (Ok(None) arm)"
);
assert_eq!(parsed["key"], "gotcha:store-read-ok-none");
let graph = graph_arc.read().await;
let record = graph
.store()
.get("gotcha:store-read-ok-none")
.await
.unwrap()
.expect("record must exist after write");
assert_eq!(record.value, "Regression rule because regression reason");
}
#[tokio::test]
async fn test_tombstone_removes_in_memory_graph_edges() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let file_record = Record::layer0_file_stub("file:src/target.rs", device_id(), 1, now());
store.put("file:src/target.rs", &file_record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let write_result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:graph-cleanup-test".to_string(),
value: "Never skip validation because it causes silent data corruption".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "Never skip validation",
"reason": "causes silent data corruption",
"severity": "High",
"affected_files": ["src/target.rs"],
"ref_url": null,
"discovered_session": 0,
"confirmed": false
}),
tags: vec![],
priority: "High".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&write_result).unwrap();
assert_eq!(parsed["ok"], true, "gotcha write must succeed");
{
let graph = graph_arc.read().await;
let neighbors = graph.neighbors("file:src/target.rs", &EdgeKind::HasGotcha);
assert!(
neighbors.contains(&"gotcha:graph-cleanup-test".to_string()),
"HasGotcha edge must exist after write; neighbors: {neighbors:?}"
);
}
let delete_result = call_mem_set(
&graph_arc,
MemSetParams {
action: "delete".to_string(),
key: "gotcha:graph-cleanup-test".to_string(),
value: String::new(),
category: String::new(),
payload: serde_json::json!({}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&delete_result).unwrap();
assert_eq!(parsed["ok"], true, "gotcha delete must succeed");
assert_eq!(parsed["tombstoned"], true);
{
let graph = graph_arc.read().await;
let neighbors = graph.neighbors("file:src/target.rs", &EdgeKind::HasGotcha);
assert!(
!neighbors.contains(&"gotcha:graph-cleanup-test".to_string()),
"HasGotcha edge must be removed after delete; neighbors: {neighbors:?}"
);
}
let graph_query_result = call_mem_query(
&graph_arc,
"file:src/target.rs",
crate::mcp::protocol::QueryMode::Graph,
20,
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&graph_query_result).unwrap();
let gotchas = parsed["gotchas"]
.as_array()
.expect("gotchas group must be an array");
assert!(
!gotchas
.iter()
.any(|g| g["key"] == "gotcha:graph-cleanup-test"),
"deleted gotcha must not appear in graph query results"
);
}
#[tokio::test]
async fn test_graph_mode_respects_global_limit() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let file_record =
Record::layer0_file_stub("file:src/graph_limit.rs", device_id(), 1, now());
store
.put("file:src/graph_limit.rs", &file_record)
.await
.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
for i in 0..5 {
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: format!("gotcha:limit-test-{i}"),
value: format!("Limit test gotcha {i}"),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": format!("Limit rule {i}"),
"reason": "testing",
"severity": "Normal",
"affected_files": ["src/graph_limit.rs"],
"ref_url": null,
"discovered_session": 0,
"confirmed": false
}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["ok"], true, "gotcha write {i} must succeed");
}
let result = call_mem_query(
&graph_arc,
"file:src/graph_limit.rs",
crate::mcp::protocol::QueryMode::Graph,
3,
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(parsed.get("error").is_none(), "graph query must not error");
let mut total = 0;
for group in &["gotchas", "co_changes", "imports", "decisions", "notes"] {
if let Some(arr) = parsed[group].as_array() {
total += arr.len();
}
}
assert!(
total <= 3,
"graph mode with limit=3 must return at most 3 total records, got {total}"
);
}
#[tokio::test]
async fn test_graph_mode_limit_zero_returns_empty() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let file_record = Record::layer0_file_stub("file:src/zero.rs", device_id(), 1, now());
store.put("file:src/zero.rs", &file_record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let _ = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:zero-limit-test".to_string(),
value: "Should not appear".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "Zero limit test",
"reason": "testing",
"severity": "Normal",
"affected_files": ["src/zero.rs"],
"ref_url": null,
"discovered_session": 0,
"confirmed": false
}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let result = call_mem_query(
&graph_arc,
"file:src/zero.rs",
crate::mcp::protocol::QueryMode::Graph,
0,
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
let mut total = 0;
for group in &["gotchas", "co_changes", "imports", "decisions", "notes"] {
if let Some(arr) = parsed[group].as_array() {
total += arr.len();
}
}
assert_eq!(total, 0, "limit=0 must return zero records, got {total}");
}
#[tokio::test]
async fn test_mem_set_overwrite_preserves_existing_data() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let mut original =
make_gotcha_record("gotcha:preserve-confirmation", "Original rule", true, 0.7);
original.source = RecordSource::DeveloperManual;
original.confidence.value = 0.85;
original.confidence.confirmation_count = 3;
store
.put("gotcha:preserve-confirmation", &original)
.await
.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:preserve-confirmation".to_string(),
value: "Updated rule from enrichment".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({"reason": "updated reason"}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["ok"], true, "overwrite must succeed");
let graph = graph_arc.read().await;
let record = graph
.store()
.get("gotcha:preserve-confirmation")
.await
.unwrap()
.expect("record must exist");
assert_eq!(
record.value, "Updated rule from enrichment",
"value must be updated"
);
assert!(
record.confidence.value >= 0.80,
"confirmed record confidence must be preserved, got {}",
record.confidence.value
);
}
#[tokio::test]
async fn test_tombstone_multi_file_removes_all_edges() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let f1 = Record::layer0_file_stub("file:src/a.rs", device_id(), 1, now());
let f2 = Record::layer0_file_stub("file:src/b.rs", device_id(), 1, now());
store.put("file:src/a.rs", &f1).await.unwrap();
store.put("file:src/b.rs", &f2).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:multi-file-tombstone".to_string(),
value: "Cross-file gotcha".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "Cross-file rule",
"reason": "testing multi-file cleanup",
"severity": "Normal",
"affected_files": ["src/a.rs", "src/b.rs"],
"ref_url": null,
"discovered_session": 0,
"confirmed": false
}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["ok"], true);
{
let graph = graph_arc.read().await;
assert!(
graph
.neighbors("file:src/a.rs", &EdgeKind::HasGotcha)
.contains(&"gotcha:multi-file-tombstone".to_string()),
"file:src/a.rs must have HasGotcha edge before delete"
);
assert!(
graph
.neighbors("file:src/b.rs", &EdgeKind::HasGotcha)
.contains(&"gotcha:multi-file-tombstone".to_string()),
"file:src/b.rs must have HasGotcha edge before delete"
);
}
let result = call_mem_set(
&graph_arc,
MemSetParams {
action: "delete".to_string(),
key: "gotcha:multi-file-tombstone".to_string(),
value: String::new(),
category: String::new(),
payload: serde_json::json!({}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert_eq!(parsed["ok"], true);
{
let graph = graph_arc.read().await;
assert!(
!graph
.neighbors("file:src/a.rs", &EdgeKind::HasGotcha)
.contains(&"gotcha:multi-file-tombstone".to_string()),
"file:src/a.rs must NOT have HasGotcha edge after delete"
);
assert!(
!graph
.neighbors("file:src/b.rs", &EdgeKind::HasGotcha)
.contains(&"gotcha:multi-file-tombstone".to_string()),
"file:src/b.rs must NOT have HasGotcha edge after delete"
);
}
}
#[tokio::test]
async fn test_confirm_is_non_idempotent() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let file_record = Record::layer0_file_stub("file:src/idem.rs", device_id(), 1, now());
store.put("file:src/idem.rs", &file_record).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let graph_arc = std::sync::Arc::new(tokio::sync::RwLock::new(graph));
let _ = call_mem_set(
&graph_arc,
MemSetParams {
action: "write".to_string(),
key: "gotcha:idem-test".to_string(),
value: "Idempotency test rule".to_string(),
category: "Gotcha".to_string(),
payload: serde_json::json!({
"rule": "Idempotency test",
"reason": "testing",
"severity": "Normal",
"affected_files": ["src/idem.rs"],
"ref_url": null,
"discovered_session": 0,
"confirmed": false
}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let r1 = call_mem_set(
&graph_arc,
MemSetParams {
action: "confirm".to_string(),
key: "gotcha:idem-test".to_string(),
value: String::new(),
category: String::new(),
payload: serde_json::json!({}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let p1: serde_json::Value = serde_json::from_str(&r1).unwrap();
assert_eq!(p1["ok"], true, "first confirm must succeed");
assert_eq!(p1["confirmed"], true);
let count_after_first = {
let graph = graph_arc.read().await;
let record = graph
.store()
.get("gotcha:idem-test")
.await
.unwrap()
.unwrap();
record.confidence.confirmation_count
};
let r2 = call_mem_set(
&graph_arc,
MemSetParams {
action: "confirm".to_string(),
key: "gotcha:idem-test".to_string(),
value: String::new(),
category: String::new(),
payload: serde_json::json!({}),
tags: vec![],
priority: "Normal".to_string(),
},
)
.await;
let p2: serde_json::Value = serde_json::from_str(&r2).unwrap();
assert_eq!(p2["ok"], true, "second confirm must succeed");
let count_after_second = {
let graph = graph_arc.read().await;
let record = graph
.store()
.get("gotcha:idem-test")
.await
.unwrap()
.unwrap();
record.confidence.confirmation_count
};
assert!(
count_after_second > count_after_first,
"confirmation_count must increase on each confirm: first={count_after_first}, second={count_after_second}"
);
}
#[test]
fn test_store_read_error_refuses_write() {
let result = resolve_existing_for_write(Err(anyhow::anyhow!("simulated disk I/O timeout")));
assert!(result.is_err(), "store error must refuse write");
let err = result.unwrap_err();
let parsed: serde_json::Value = serde_json::from_str(&err).unwrap();
assert!(
parsed["error"]
.as_str()
.unwrap()
.contains("store read failed"),
"error must mention store read failure"
);
assert!(
parsed["error"]
.as_str()
.unwrap()
.contains("simulated disk I/O timeout"),
"error must include the underlying cause"
);
}
#[test]
fn test_store_read_ok_none_passes_through() {
let result = resolve_existing_for_write(Ok(None));
assert!(result.is_ok());
assert!(result.unwrap().is_none());
}
#[test]
fn test_store_read_ok_some_passes_through() {
let record = make_record("file:test.rs", "test", Category::File, 0.5);
let result = resolve_existing_for_write(Ok(Some(record.clone())));
assert!(result.is_ok());
let resolved = result.unwrap();
assert!(resolved.is_some());
assert_eq!(resolved.unwrap().key, "file:test.rs");
}
#[tokio::test]
async fn bootstrap_highest_impact_section_appears() {
let dir = TempDir::new().unwrap();
let store = Store::open(dir.path()).await.unwrap();
let fr_critical = FileRecord {
path: "src/core.rs".to_string(),
purpose: "Core module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 100,
last_modified_session: 0,
content_hash: None,
line_count: 0,
blast_radius: Some(crate::analysis::blast_radius::BlastRadius {
direct: 45,
transitive: 10,
score: 48.0,
tier: crate::analysis::blast_radius::BlastTier::Critical,
}),
propagated_staleness: None,
};
let mut rec = make_record("file:src/core.rs", "Core module", Category::File, 0.5);
rec.payload = serde_json::to_value(&fr_critical).ok();
store.put("file:src/core.rs", &rec).await.unwrap();
let fr_low = FileRecord {
path: "src/leaf.rs".to_string(),
purpose: "Leaf module".to_string(),
entry_points: vec![],
imports: vec![],
gotcha_keys: vec![],
decision_keys: vec![],
todos: vec![],
unsafe_count: 0,
unwrap_count: 0,
change_frequency: 0,
last_author: None,
is_hotspot: false,
token_cost_estimate: 100,
last_modified_session: 0,
content_hash: None,
line_count: 0,
blast_radius: Some(crate::analysis::blast_radius::BlastRadius {
direct: 3,
transitive: 0,
score: 3.0,
tier: crate::analysis::blast_radius::BlastTier::Low,
}),
propagated_staleness: None,
};
let mut rec2 = make_record("file:src/leaf.rs", "Leaf module", Category::File, 0.5);
rec2.payload = serde_json::to_value(&fr_low).ok();
store.put("file:src/leaf.rs", &rec2).await.unwrap();
let graph = Graph::load(store).await.unwrap();
let packet = assemble_context_packet(
graph.store(),
&graph,
&["src/core.rs".to_string(), "src/leaf.rs".to_string()],
)
.await
.unwrap();
assert!(
packet.injection_string.contains("Highest Impact"),
"bootstrap must include highest impact section, got: {}",
packet.injection_string
);
assert!(
packet.injection_string.contains("src/core.rs"),
"critical file must appear in impact section"
);
let core_pos = packet.injection_string.find("src/core.rs").unwrap();
let leaf_pos = packet
.injection_string
.find("src/leaf.rs")
.unwrap_or(usize::MAX);
assert!(
core_pos < leaf_pos,
"core.rs should appear before leaf.rs in impact section"
);
}
fn make_params(action: &str, key: &str) -> MemSetParams {
MemSetParams {
action: action.to_string(),
key: key.to_string(),
value: String::new(),
category: String::new(),
payload: serde_json::Value::Object(serde_json::Map::new()),
tags: vec![],
priority: "Normal".to_string(),
}
}
#[test]
fn mem_set_socket_routes_gotcha_confirm() {
let p = make_params("confirm", "gotcha:foo");
let cmd = build_mem_set_command(&p).expect("must build");
assert_eq!(cmd.kind(), "gotcha_confirm");
assert_eq!(cmd.target_key(), "gotcha:foo");
}
#[test]
fn mem_set_socket_rejects_confirm_on_non_gotcha_key() {
let p = make_params("confirm", "decision:not-allowed");
let err = build_mem_set_command(&p).expect_err("must reject");
assert!(
err.contains("gotcha:"),
"error must mention gotcha: prefix, got: {err}"
);
}
#[test]
fn mem_set_socket_routes_gotcha_tombstone() {
let p = make_params("delete", "gotcha:foo");
let cmd = build_mem_set_command(&p).expect("must build");
assert_eq!(cmd.kind(), "gotcha_tombstone");
assert_eq!(cmd.target_key(), "gotcha:foo");
}
#[test]
fn mem_set_socket_routes_gotcha_upsert_by_key_prefix() {
let mut p = make_params("write", "gotcha:stripe-idempotency");
p.payload = serde_json::json!({
"rule": "Always include an idempotency key",
"reason": "Stripe retries cause double charges without it",
"severity": "High",
"affected_files": ["src/payments/stripe.rs"],
});
p.tags = vec!["payments".into()];
p.priority = "High".into();
let cmd = build_mem_set_command(&p).expect("must build");
assert_eq!(cmd.kind(), "gotcha_upsert");
match cmd {
Command::GotchaUpsert(input) => {
assert_eq!(input.key, "gotcha:stripe-idempotency");
assert_eq!(input.rule, "Always include an idempotency key");
assert_eq!(input.severity, proto::Severity::High);
assert_eq!(input.priority, proto::Priority::High);
assert_eq!(input.affected_files, vec!["src/payments/stripe.rs"]);
assert_eq!(input.tags, vec!["payments".to_string()]);
}
_ => panic!("expected GotchaUpsert"),
}
}
#[test]
fn mem_set_socket_routes_decision_upsert_by_key_prefix() {
let mut p = make_params("write", "decision:retry-strategy");
p.value = "We use exponential backoff because linear overloads downstream".into();
p.payload = serde_json::json!({
"summary": "Exponential backoff for all retries",
"rationale": "Linear retry caused cascading failures in prod 2024-01",
});
let cmd = build_mem_set_command(&p).expect("must build");
assert_eq!(cmd.kind(), "decision_upsert");
match cmd {
Command::DecisionUpsert(input) => {
assert_eq!(input.slug, "retry-strategy");
assert_eq!(input.summary, "Exponential backoff for all retries");
assert!(input.rationale.contains("cascading"));
}
_ => panic!("expected DecisionUpsert"),
}
}
#[test]
fn mem_set_socket_routes_dev_note_upsert_by_key_prefix() {
let mut p = make_params("write", "dev_note:remember-changelog");
p.value = "Remember to update the changelog before release".into();
let cmd = build_mem_set_command(&p).expect("must build");
assert_eq!(cmd.kind(), "dev_note_upsert");
match cmd {
Command::DevNoteUpsert(input) => {
assert_eq!(input.key.as_deref(), Some("dev_note:remember-changelog"));
assert!(input.text.contains("changelog"));
}
_ => panic!("expected DevNoteUpsert"),
}
}
#[test]
fn mem_set_socket_rejects_write_with_unknown_prefix() {
let p = make_params("write", "file:src/main.rs");
let err = build_mem_set_command(&p).expect_err("must reject");
assert!(
err.contains("gotcha:") && err.contains("decision:") && err.contains("dev_note:"),
"error must list valid prefixes, got: {err}"
);
}
#[test]
fn mem_set_socket_rejects_unknown_action() {
let p = make_params("smuggle", "gotcha:foo");
let err = build_mem_set_command(&p).expect_err("must reject");
assert!(
err.contains("smuggle"),
"error must echo the bad action, got: {err}"
);
}
#[test]
fn mem_set_socket_rejects_gotcha_write_missing_payload_fields() {
let p = make_params("write", "gotcha:incomplete");
let err = build_mem_set_command(&p).expect_err("must reject");
assert!(
err.contains("rule"),
"error must mention 'rule', got: {err}"
);
}
#[test]
fn mem_set_socket_handles_codex_string_payload() {
let mut p = make_params("write", "gotcha:codex-style");
p.payload = serde_json::Value::String(
r#"{"rule":"do X","reason":"because Y","severity":"low"}"#.to_string(),
);
let cmd = build_mem_set_command(&p).expect("must build from stringified payload");
match cmd {
Command::GotchaUpsert(input) => {
assert_eq!(input.rule, "do X");
assert_eq!(input.severity, proto::Severity::Low);
}
_ => panic!("expected GotchaUpsert"),
}
}
}