use std::sync::{Arc, Mutex};
use cortex_store::repo::MemoryRepo;
use serde_json::{json, Value};
use crate::tool_handler::{GateId, ToolError, ToolHandler};
const DEFAULT_LIMIT: u32 = 3;
const MAX_LIMIT: u32 = 20;
#[derive(Debug)]
pub struct CortexSuggestTool {
pool: Arc<Mutex<cortex_store::Pool>>,
}
impl CortexSuggestTool {
#[must_use]
pub fn new(pool: Arc<Mutex<cortex_store::Pool>>) -> Self {
Self { pool }
}
}
impl ToolHandler for CortexSuggestTool {
fn name(&self) -> &'static str {
"cortex_suggest"
}
fn gate_set(&self) -> &'static [GateId] {
&[GateId::FtsRead]
}
fn call(&self, params: Value) -> Result<Value, ToolError> {
let query: Option<String> = match params.get("query") {
None | Some(Value::Null) => None,
Some(Value::String(s)) => {
let trimmed = s.trim();
if trimmed.is_empty() {
None
} else {
Some(trimmed.to_owned())
}
}
Some(other) => {
return Err(ToolError::InvalidParams(format!(
"query must be a string, got {other}"
)));
}
};
let limit: u32 = match params.get("limit") {
None | Some(Value::Null) => DEFAULT_LIMIT,
Some(v) => {
let n = v.as_u64().ok_or_else(|| {
ToolError::InvalidParams("limit must be a non-negative integer".into())
})?;
let n = u32::try_from(n).unwrap_or(MAX_LIMIT);
n.min(MAX_LIMIT)
}
};
if limit == 0 {
return Ok(json!({
"suggestions": [],
"query_used": query,
}));
}
let pool = self
.pool
.lock()
.map_err(|err| ToolError::Internal(format!("failed to acquire store lock: {err}")))?;
let repo = MemoryRepo::new(&pool);
let suggestions: Vec<Value> = if let Some(ref q) = query {
let limit_usize = usize::try_from(limit).unwrap_or(usize::MAX);
let fetch_limit = limit_usize.saturating_mul(4).max(limit_usize + 10);
let hits = repo
.fts5_search(q, fetch_limit)
.map_err(|err| ToolError::Internal(format!("fts5 search failed: {err}")))?;
let mut results = Vec::with_capacity(limit_usize);
for (memory_id, raw_rank) in hits {
if results.len() >= limit_usize {
break;
}
let record = repo.get_by_id(&memory_id).map_err(|err| {
ToolError::Internal(format!("failed to fetch memory {memory_id}: {err}"))
})?;
if let Some(m) = record {
if m.status == "active" {
let score = if raw_rank.is_finite() {
raw_rank.exp().clamp(0.0, 1.0)
} else {
0.0_f32
};
let salience = m.confidence as f32;
results.push(json!({
"memory_id": m.id.to_string(),
"claim": m.claim,
"salience": salience,
"score": score,
}));
}
}
}
results
} else {
let mut memories = repo.list_by_status("active").map_err(|err| {
ToolError::Internal(format!("failed to list active memories: {err}"))
})?;
memories.sort_by(|a, b| {
b.confidence
.partial_cmp(&a.confidence)
.unwrap_or(std::cmp::Ordering::Equal)
});
memories
.into_iter()
.take(limit as usize)
.map(|m| {
let salience = m.confidence as f32;
json!({
"memory_id": m.id.to_string(),
"claim": m.claim,
"salience": salience,
"score": salience,
})
})
.collect()
};
Ok(json!({
"suggestions": suggestions,
"query_used": query,
}))
}
}
#[deprecated(since = "0.1.0", note = "use CortexSuggestTool instead")]
pub type CortexSuggestStub = CortexSuggestTool;
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use chrono::{TimeZone, Utc};
use cortex_core::{AuditRecordId, Event, EventSource, EventType, SCHEMA_VERSION};
use rusqlite::Connection;
use serde_json::json;
use super::*;
use cortex_store::migrate::apply_pending;
use cortex_store::repo::memories::accept_candidate_policy_decision_test_allow;
use cortex_store::repo::{EventRepo, MemoryAcceptanceAudit, MemoryCandidate, MemoryRepo};
fn make_pool() -> Arc<Mutex<Connection>> {
let conn = Connection::open_in_memory().expect("in-memory sqlite");
apply_pending(&conn).expect("apply_pending");
Arc::new(Mutex::new(conn))
}
fn make_tool(pool: Arc<Mutex<Connection>>) -> CortexSuggestTool {
CortexSuggestTool::new(pool)
}
fn ts(second: u32) -> chrono::DateTime<Utc> {
Utc.with_ymd_and_hms(2026, 1, 1, 12, 0, second).unwrap()
}
fn ensure_event(pool: &Connection, event_id: &str, second: u32) {
let parsed_id = event_id.parse().unwrap();
let repo = EventRepo::new(pool);
if repo.get_by_id(&parsed_id).expect("query event").is_some() {
return;
}
repo.append(&Event {
id: parsed_id,
schema_version: SCHEMA_VERSION,
observed_at: ts(second),
recorded_at: ts(second),
source: EventSource::Tool {
name: "suggest-test".into(),
},
event_type: EventType::ToolResult,
trace_id: None,
session_id: Some("suggest-test".into()),
domain_tags: vec!["test".into()],
payload: json!({}),
payload_hash: format!("hash-{second}"),
prev_event_hash: None,
event_hash: format!("ehash-{second}"),
})
.expect("append event");
}
fn insert_active(
pool: &Connection,
memory_id: &str,
claim: &str,
confidence: f64,
event_id: &str,
second: u32,
) {
ensure_event(pool, event_id, second);
let repo = MemoryRepo::new(pool);
let candidate = MemoryCandidate {
id: memory_id.parse().unwrap(),
memory_type: "semantic".into(),
claim: claim.into(),
source_episodes_json: json!([]),
source_events_json: json!([event_id]),
domains_json: json!(["test"]),
salience_json: json!({"score": confidence}),
confidence,
authority: "user".into(),
applies_when_json: json!([]),
does_not_apply_when_json: json!([]),
created_at: ts(second),
updated_at: ts(second),
};
repo.insert_candidate(&candidate).expect("insert candidate");
let audit = MemoryAcceptanceAudit {
id: AuditRecordId::new(),
actor_json: json!({"kind": "test"}),
reason: "suggest test".into(),
source_refs_json: json!([memory_id]),
created_at: ts(second + 1),
};
repo.accept_candidate(
&memory_id.parse().unwrap(),
ts(second + 2),
&audit,
&accept_candidate_policy_decision_test_allow(),
)
.expect("accept candidate");
}
#[test]
fn empty_store_returns_empty_suggestions() {
let pool = make_pool();
let tool = make_tool(pool);
let result = tool.call(json!({})).expect("call must not error");
assert_eq!(
result["suggestions"],
json!([]),
"empty store must produce empty suggestions"
);
assert!(
result["query_used"].is_null(),
"query_used must be null when no query was supplied"
);
}
#[test]
fn empty_store_with_query_returns_empty_suggestions() {
let pool = make_pool();
let tool = make_tool(pool);
let result = tool
.call(json!({"query": "architecture"}))
.expect("call must not error");
assert_eq!(result["suggestions"], json!([]));
assert_eq!(result["query_used"], json!("architecture"));
}
#[test]
fn no_query_returns_highest_confidence_memories() {
let pool = make_pool();
{
let conn = pool.lock().unwrap();
insert_active(
&conn,
"mem_01ARZ3NDEKTSV4RRFFQ69G5FAV",
"low confidence memory",
0.3,
"evt_01ARZ3NDEKTSV4RRFFQ69G5FA1",
0,
);
insert_active(
&conn,
"mem_01ARZ3NDEKTSV4RRFFQ69G5FB1",
"high confidence memory",
0.95,
"evt_01ARZ3NDEKTSV4RRFFQ69G5FA2",
10,
);
insert_active(
&conn,
"mem_01ARZ3NDEKTSV4RRFFQ69G5FC1",
"medium confidence memory",
0.6,
"evt_01ARZ3NDEKTSV4RRFFQ69G5FA3",
20,
);
}
let tool = make_tool(pool);
let result = tool
.call(json!({"limit": 2}))
.expect("call must not error");
let suggestions = result["suggestions"].as_array().expect("suggestions array");
assert_eq!(suggestions.len(), 2, "limit=2 must return 2 suggestions");
assert_eq!(
suggestions[0]["claim"].as_str().unwrap(),
"high confidence memory",
"first suggestion must be highest-confidence memory"
);
assert_eq!(
suggestions[1]["claim"].as_str().unwrap(),
"medium confidence memory",
"second suggestion must be medium-confidence memory"
);
assert!(
result["query_used"].is_null(),
"query_used must be null for no-query path"
);
}
#[test]
fn no_query_default_limit_is_three() {
let pool = make_pool();
{
let conn = pool.lock().unwrap();
for (i, conf) in [0.9_f64, 0.8, 0.7, 0.6, 0.5].iter().enumerate() {
let i = i as u32;
let mem_id = format!("mem_01ARZ3NDEKTSV4RRFFQ69G5F{:02}", i + 10);
let evt_id = format!("evt_01ARZ3NDEKTSV4RRFFQ69G5F{:02}", i + 10);
insert_active(&conn, &mem_id, &format!("memory {i}"), *conf, &evt_id, i * 10);
}
}
let tool = make_tool(pool);
let result = tool.call(json!({})).expect("call must not error");
let suggestions = result["suggestions"].as_array().expect("suggestions array");
assert_eq!(
suggestions.len(),
3,
"default limit must return 3 suggestions"
);
}
#[test]
fn query_returns_relevant_memories() {
let pool = make_pool();
{
let conn = pool.lock().unwrap();
insert_active(
&conn,
"mem_01ARZ3NDEKTSV4RRFFQ69G5FAV",
"trust boundary architecture decisions should be documented",
0.8,
"evt_01ARZ3NDEKTSV4RRFFQ69G5FA1",
0,
);
insert_active(
&conn,
"mem_01ARZ3NDEKTSV4RRFFQ69G5FB1",
"database connection pooling improves throughput",
0.7,
"evt_01ARZ3NDEKTSV4RRFFQ69G5FA2",
10,
);
}
let tool = make_tool(pool);
let result = tool
.call(json!({"query": "trust architecture", "limit": 3}))
.expect("call must not error");
let suggestions = result["suggestions"].as_array().expect("suggestions array");
assert!(
!suggestions.is_empty(),
"FTS5 query must return at least one match"
);
assert!(
suggestions
.iter()
.any(|s| s["claim"].as_str().unwrap_or("").contains("trust boundary")),
"FTS5 query must surface the trust-boundary memory"
);
assert_eq!(
result["query_used"].as_str().unwrap(),
"trust architecture"
);
for s in suggestions {
assert!(s.get("memory_id").is_some(), "suggestion must have memory_id");
assert!(s.get("claim").is_some(), "suggestion must have claim");
assert!(s.get("salience").is_some(), "suggestion must have salience");
assert!(s.get("score").is_some(), "suggestion must have score");
}
}
#[test]
fn query_only_surfaces_active_memories() {
let pool = make_pool();
{
let conn = pool.lock().unwrap();
let repo = MemoryRepo::new(&conn);
ensure_event(&conn, "evt_01ARZ3NDEKTSV4RRFFQ69G5FA1", 0);
let candidate = MemoryCandidate {
id: "mem_01ARZ3NDEKTSV4RRFFQ69G5FAV".parse().unwrap(),
memory_type: "semantic".into(),
claim: "candidate trust architecture memory".into(),
source_episodes_json: json!([]),
source_events_json: json!(["evt_01ARZ3NDEKTSV4RRFFQ69G5FA1"]),
domains_json: json!(["test"]),
salience_json: json!({}),
confidence: 0.9,
authority: "user".into(),
applies_when_json: json!([]),
does_not_apply_when_json: json!([]),
created_at: ts(0),
updated_at: ts(0),
};
repo.insert_candidate(&candidate)
.expect("insert candidate");
}
let tool = make_tool(pool);
let result = tool
.call(json!({"query": "trust architecture"}))
.expect("call must not error");
let suggestions = result["suggestions"].as_array().expect("suggestions array");
assert!(
suggestions.is_empty(),
"candidate memories must not appear in suggestions"
);
}
#[test]
fn name_is_cortex_suggest() {
let pool = make_pool();
let tool = make_tool(pool);
assert_eq!(tool.name(), "cortex_suggest");
}
#[test]
fn gate_set_contains_fts_read() {
let pool = make_pool();
let tool = make_tool(pool);
assert!(
tool.gate_set().contains(&GateId::FtsRead),
"cortex_suggest must declare GateId::FtsRead"
);
}
#[test]
fn zero_limit_returns_empty_immediately() {
let pool = make_pool();
let tool = make_tool(pool);
let result = tool
.call(json!({"limit": 0}))
.expect("zero limit must not error");
assert_eq!(result["suggestions"], json!([]));
}
#[test]
fn invalid_query_type_returns_error() {
let pool = make_pool();
let tool = make_tool(pool);
let err = tool
.call(json!({"query": 42}))
.expect_err("non-string query must return an error");
assert!(matches!(err, ToolError::InvalidParams(_)));
}
#[test]
fn empty_string_query_treated_as_no_query() {
let pool = make_pool();
let tool = make_tool(pool);
let result = tool
.call(json!({"query": " "}))
.expect("whitespace-only query must not error");
assert!(result["query_used"].is_null());
}
}