cortex_mcp/tools/

suggest.rs

//! `cortex_suggest` — server-initiated memory suggestions.
//!
//! Surfaces the memories most worth paying attention to right now:
//!
//! - With `query`: run FTS5 full-text search and return the top-N hits,
//!   filtered to `status = 'active'`.
//! - Without `query`: return the top-N active memories by confidence score
//!   (highest first), as a salience-proxy ranking.
//!
//! ## Schema
//!
//! ```text
//! cortex_suggest(
//!   query?:  string,   // optional search query
//!   limit?:  u32,      // default 3, cap 20
//! ) -> {
//!   suggestions: [{ memory_id: string, claim: string, salience: number, score: number }],
//!   query_used: string | null,
//! }
//! ```
//!
//! ## Salience and score
//!
//! - `salience` is the memory's `confidence` value, the best uniformly
//!   available proxy for how salient a memory is without a secondary lookup.
//! - `score` is the FTS5 normalized BM25 score (`exp(rank)` mapped to
//!   `(0, 1]`) when a query was provided, or equals `salience` when no
//!   query was provided.
//!
//! ## Gate
//!
//! [`GateId::FtsRead`] — session tier, no confirmation token required.

use std::sync::{Arc, Mutex};

use cortex_store::repo::MemoryRepo;
use serde_json::{json, Value};

use crate::tool_handler::{GateId, ToolError, ToolHandler};

/// Default number of suggestions returned when `limit` is absent.
const DEFAULT_LIMIT: u32 = 3;
/// Server-side cap on `limit`.
const MAX_LIMIT: u32 = 20;

/// MCP tool: `cortex_suggest`.
///
/// Returns the top-N memories most relevant to `query` (FTS5 ranked), or the
/// top-N highest-confidence active memories when no query is provided.
#[derive(Debug)]
pub struct CortexSuggestTool {
    pool: Arc<Mutex<cortex_store::Pool>>,
}

impl CortexSuggestTool {
    /// Construct the tool over a shared store connection.
    #[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> {
        // Parse optional query string.
        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}"
                )));
            }
        };

        // Parse optional limit (default 3, cap 20).
        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 {
            // FTS5 path: search for relevant memories, filter to active status.
            let limit_usize = usize::try_from(limit).unwrap_or(usize::MAX);
            // Ask for more hits than needed to allow filtering by status.
            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}"))
                })?;

                // Skip non-active memories (candidates, quarantined, etc.).
                if let Some(m) = record {
                    if m.status == "active" {
                        // Normalize BM25 rank: exp(rank) maps (-inf, 0] to (0, 1]
                        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 {
            // No-query path: highest-confidence active memories.
            let mut memories = repo.list_by_status("active").map_err(|err| {
                ToolError::Internal(format!("failed to list active memories: {err}"))
            })?;

            // Sort by confidence descending (best salience proxy available
            // without a secondary table lookup).
            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,
        }))
    }
}

/// Backward-compat type alias for callers that referenced the old stub name.
///
/// # Deprecated
///
/// Use [`CortexSuggestTool`] instead.
#[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()
    }

    /// Insert a source event (idempotent guard on the same pool).
    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");
    }

    /// Insert a memory as active with the given confidence.
    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: empty store returns empty suggestions --------------------------

    #[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: no-query returns highest-salience memories --------------------

    #[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");

        // First suggestion must be the highest-confidence memory.
        assert_eq!(
            suggestions[0]["claim"].as_str().unwrap(),
            "high confidence memory",
            "first suggestion must be highest-confidence memory"
        );
        // Second suggestion must be the medium-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: query returns relevant memories --------------------------------

    #[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");

        // Must have found at least the relevant memory.
        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"
        );

        // Each suggestion must have the required fields.
        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() {
        // Insert a memory as candidate (not active) -- the suggest tool must not
        // return it even if it matches the query.
        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");
            // Do NOT accept -- stays as 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: tool metadata -------------------------------------------------

    #[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);
        // An empty/whitespace query is treated as no query -- must not error.
        let result = tool
            .call(json!({"query": "   "}))
            .expect("whitespace-only query must not error");
        assert!(result["query_used"].is_null());
    }
}