lean_ctx_server/tools/

ctx_knowledge_relations.rs

// SPDX-License-Identifier: Apache-2.0
use lean_ctx::core::knowledge::ProjectKnowledge;
use lean_ctx::core::knowledge_relations::{
    KnowledgeEdge, KnowledgeEdgeKind, KnowledgeNodeRef, KnowledgeRelationGraph, format_mermaid,
    parse_node_ref,
};

fn load_policy_or_error() -> Result<lean_ctx::core::memory_policy::MemoryPolicy, String> {
    super::knowledge_shared::load_policy_or_error()
}

fn ensure_current_fact_exists(knowledge: &ProjectKnowledge, node: &KnowledgeNodeRef) -> bool {
    knowledge
        .facts
        .iter()
        .any(|f| f.is_current() && f.category == node.category && f.key == node.key)
}

fn parse_kind_or_default(value: Option<&str>) -> Result<KnowledgeEdgeKind, String> {
    let kind_str = value.unwrap_or("related_to");
    KnowledgeEdgeKind::parse(kind_str).ok_or_else(|| {
        "Error: relation kind must be one of depends_on|related_to|supports|contradicts|supersedes"
            .to_string()
    })
}

fn parse_target_or_error(query: Option<&str>) -> Result<KnowledgeNodeRef, String> {
    let Some(q) = query else {
        return Err("Error: query is required and must be 'category/key'".to_string());
    };
    parse_node_ref(q).ok_or_else(|| "Error: query must be 'category/key'".to_string())
}

fn parse_direction(query: Option<&str>) -> &'static str {
    match query.unwrap_or("all").trim().to_lowercase().as_str() {
        "in" | "incoming" => "in",
        "out" | "outgoing" => "out",
        _ => "all",
    }
}

fn derived_supersedes_edges(
    knowledge: &ProjectKnowledge,
    focus: &KnowledgeNodeRef,
) -> Vec<KnowledgeEdge> {
    let mut out = Vec::new();
    let focus_id = focus.id();

    for f in knowledge.facts.iter().filter(|f| f.is_current()) {
        if f.category == focus.category && f.key == focus.key {
            if let Some(s) = &f.supersedes
                && let Some(to) = parse_node_ref(s)
            {
                if to == *focus {
                    continue;
                }
                out.push(KnowledgeEdge {
                    from: focus.clone(),
                    to,
                    kind: KnowledgeEdgeKind::Supersedes,
                    created_at: f.created_at,
                    last_seen: None,
                    count: 0,
                    source_session: f.source_session.clone(),
                    strength: 0.5,
                    decay_rate: 0.02,
                });
            }
        } else if f.supersedes.as_deref() == Some(&focus_id) {
            out.push(KnowledgeEdge {
                from: KnowledgeNodeRef::new(&f.category, &f.key),
                to: focus.clone(),
                kind: KnowledgeEdgeKind::Supersedes,
                created_at: f.created_at,
                last_seen: None,
                count: 0,
                source_session: f.source_session.clone(),
                strength: 0.5,
                decay_rate: 0.02,
            });
        }
    }

    out
}

pub fn handle_relate(
    project_root: &str,
    category: Option<&str>,
    key: Option<&str>,
    value: Option<&str>,
    query: Option<&str>,
    session_id: &str,
) -> String {
    let Some(cat) = category else {
        return "Error: category is required for relate".to_string();
    };
    let Some(k) = key else {
        return "Error: key is required for relate".to_string();
    };

    let from = KnowledgeNodeRef::new(cat, k);
    let to = match parse_target_or_error(query) {
        Ok(n) => n,
        Err(e) => return e,
    };
    let kind = match parse_kind_or_default(value) {
        Ok(k) => k,
        Err(e) => return e,
    };

    // An invalid memory policy still refuses the write.
    if let Err(e) = load_policy_or_error() {
        return e;
    }

    let knowledge = ProjectKnowledge::load_or_create(project_root);
    if !ensure_current_fact_exists(&knowledge, &from) {
        return format!(
            "Error: no current fact exists for [{}] {}. Use ctx_knowledge remember first.",
            from.category, from.key
        );
    }
    if !ensure_current_fact_exists(&knowledge, &to) {
        return format!(
            "Error: no current fact exists for [{}] {}. Use ctx_knowledge remember first.",
            to.category, to.key
        );
    }

    let max_edges = lean_ctx::core::knowledge_relations::edge_bound(
        knowledge.facts.iter().filter(|f| f.is_current()).count(),
    );
    let saved =
        KnowledgeRelationGraph::mutate_locked(project_root, &knowledge.project_hash, |graph| {
            let created = graph.upsert_edge(from.clone(), to.clone(), kind, session_id);
            Some((created, graph.enforce_cap(max_edges)))
        });

    match saved.map(Option::unwrap_or_default) {
        Ok((created, capped)) => {
            let verb = if created { "added" } else { "reinforced" };
            let mut out = format!(
                "Relation {verb}: {} -({})-> {}",
                from.id(),
                kind.as_str(),
                to.id()
            );
            if capped {
                out.push_str(&format!(" (note: capped to {max_edges} edges)"));
            }
            out
        }
        Err(e) => format!(
            "Relation not saved: {e} ({} -({})-> {}); retry once the knowledge store is free",
            from.id(),
            kind.as_str(),
            to.id()
        ),
    }
}

pub fn handle_unrelate(
    project_root: &str,
    category: Option<&str>,
    key: Option<&str>,
    value: Option<&str>,
    query: Option<&str>,
) -> String {
    let Some(cat) = category else {
        return "Error: category is required for unrelate".to_string();
    };
    let Some(k) = key else {
        return "Error: key is required for unrelate".to_string();
    };

    let from = KnowledgeNodeRef::new(cat, k);
    let to = match parse_target_or_error(query) {
        Ok(n) => n,
        Err(e) => return e,
    };
    let kind = if let Some(v) = value {
        match KnowledgeEdgeKind::parse(v) {
            Some(k) => Some(k),
            None => {
                return "Error: relation kind must be one of depends_on|related_to|supports|contradicts|supersedes".to_string();
            }
        }
    } else {
        None
    };

    let knowledge = ProjectKnowledge::load_or_create(project_root);
    let removed =
        KnowledgeRelationGraph::mutate_locked(project_root, &knowledge.project_hash, |graph| {
            let removed = graph.remove_edge(&from, &to, kind);
            (removed > 0).then_some(removed)
        });

    match removed {
        Ok(Some(removed)) => format!("Relation removed ({removed}): {} -> {}", from.id(), to.id()),
        Ok(None) => format!("No matching relation found: {} -> {}", from.id(), to.id()),
        Err(e) => format!(
            "Relation not removed: {e} ({} -> {}); retry once the knowledge store is free",
            from.id(),
            to.id()
        ),
    }
}

pub fn handle_relations(
    project_root: &str,
    category: Option<&str>,
    key: Option<&str>,
    value: Option<&str>,
    query: Option<&str>,
) -> String {
    let Some(cat) = category else {
        return "Error: category is required for relations".to_string();
    };
    let Some(k) = key else {
        return "Error: key is required for relations".to_string();
    };

    let focus = KnowledgeNodeRef::new(cat, k);
    let dir = parse_direction(query);
    let kind_filter = match value {
        Some(v) => match KnowledgeEdgeKind::parse(v) {
            Some(k) => Some(k),
            None => {
                return "Error: relation kind must be one of depends_on|related_to|supports|contradicts|supersedes".to_string();
            }
        },
        None => None,
    };

    let policy = match load_policy_or_error() {
        Ok(p) => p,
        Err(e) => return e,
    };
    let limit = policy.knowledge.relations_limit;

    let knowledge = ProjectKnowledge::load_or_create(project_root);
    let graph = KnowledgeRelationGraph::load_or_create(&knowledge.project_hash);

    let mut edges: Vec<&KnowledgeEdge> = graph
        .incident_edges(&focus)
        .into_iter()
        .filter(|e| match dir {
            "in" => e.to == focus,
            "out" => e.from == focus,
            _ => e.from == focus || e.to == focus,
        })
        .filter(|e| kind_filter.is_none_or(|k| e.kind == k))
        .collect();

    edges.sort_by(|a, b| {
        a.kind
            .as_str()
            .cmp(b.kind.as_str())
            .then_with(|| a.from.category.cmp(&b.from.category))
            .then_with(|| a.from.key.cmp(&b.from.key))
            .then_with(|| a.to.category.cmp(&b.to.category))
            .then_with(|| a.to.key.cmp(&b.to.key))
            .then_with(|| b.count.cmp(&a.count))
            .then_with(|| b.last_seen.cmp(&a.last_seen))
            .then_with(|| b.created_at.cmp(&a.created_at))
    });

    let derived = derived_supersedes_edges(&knowledge, &focus);
    let mut derived_filtered: Vec<KnowledgeEdge> = derived
        .into_iter()
        .filter(|e| match dir {
            "in" => e.to == focus,
            "out" => e.from == focus,
            _ => e.from == focus || e.to == focus,
        })
        .filter(|e| kind_filter.is_none_or(|k| e.kind == k))
        .collect();
    derived_filtered.sort_by(|a, b| {
        a.kind
            .as_str()
            .cmp(b.kind.as_str())
            .then_with(|| a.from.category.cmp(&b.from.category))
            .then_with(|| a.from.key.cmp(&b.from.key))
            .then_with(|| a.to.category.cmp(&b.to.category))
            .then_with(|| a.to.key.cmp(&b.to.key))
    });

    let mut seen = std::collections::HashSet::<(String, String, KnowledgeEdgeKind)>::new();
    for e in &edges {
        let _ = seen.insert((e.from.id(), e.to.id(), e.kind));
    }
    let derived_filtered: Vec<_> = derived_filtered
        .into_iter()
        .filter(|e| seen.insert((e.from.id(), e.to.id(), e.kind)))
        .collect();

    if edges.is_empty() && derived_filtered.is_empty() {
        return format!("No relations for {}.", focus.id());
    }

    let mut out = Vec::new();
    let total = edges.len() + derived_filtered.len();
    let mut shown = 0usize;
    let mut remaining = limit;

    for e in edges.iter().take(remaining) {
        let arrow = if e.from == focus { "->" } else { "<-" };
        let other = if e.from == focus { &e.to } else { &e.from };
        out.push(format!(
            "  {arrow} {} {} (count={}, last_seen={})",
            e.kind.as_str(),
            other.id(),
            e.count.max(1),
            e.last_seen
                .map_or_else(|| "n/a".to_string(), |t| t.format("%Y-%m-%d").to_string(),)
        ));
        shown += 1;
        remaining = remaining.saturating_sub(1);
    }

    for e in derived_filtered.into_iter().take(remaining) {
        let arrow = if e.from == focus { "->" } else { "<-" };
        let other = if e.from == focus { &e.to } else { &e.from };
        out.push(format!(
            "  {arrow} {} {} (derived)",
            e.kind.as_str(),
            other.id()
        ));
        shown += 1;
        remaining = remaining.saturating_sub(1);
    }

    out.insert(
        0,
        format!(
            "Relations for {} (dir={dir}, showing {shown}/{total}):",
            focus.id()
        ),
    );
    if total > shown {
        out.push(format!("  … +{} more", total - shown));
    }
    out.join("\n")
}

pub fn handle_relations_diagram(
    project_root: &str,
    category: Option<&str>,
    key: Option<&str>,
    value: Option<&str>,
    query: Option<&str>,
) -> String {
    let Some(cat) = category else {
        return "Error: category is required for relations_diagram".to_string();
    };
    let Some(k) = key else {
        return "Error: key is required for relations_diagram".to_string();
    };

    let focus = KnowledgeNodeRef::new(cat, k);
    let dir = parse_direction(query);
    let kind_filter = match value {
        Some(v) => match KnowledgeEdgeKind::parse(v) {
            Some(k) => Some(k),
            None => {
                return "Error: relation kind must be one of depends_on|related_to|supports|contradicts|supersedes".to_string();
            }
        },
        None => None,
    };

    let policy = match load_policy_or_error() {
        Ok(p) => p,
        Err(e) => return e,
    };
    let limit = policy.knowledge.relations_limit;

    let knowledge = ProjectKnowledge::load_or_create(project_root);
    let graph = KnowledgeRelationGraph::load_or_create(&knowledge.project_hash);

    let mut edges: Vec<KnowledgeEdge> = graph
        .incident_edges(&focus)
        .into_iter()
        .filter(|e| match dir {
            "in" => e.to == focus,
            "out" => e.from == focus,
            _ => e.from == focus || e.to == focus,
        })
        .filter(|e| kind_filter.is_none_or(|k| e.kind == k))
        .cloned()
        .collect();

    let derived = derived_supersedes_edges(&knowledge, &focus);
    let derived_filtered = derived
        .into_iter()
        .filter(|e| match dir {
            "in" => e.to == focus,
            "out" => e.from == focus,
            _ => e.from == focus || e.to == focus,
        })
        .filter(|e| kind_filter.is_none_or(|k| e.kind == k))
        .collect::<Vec<_>>();

    let mut seen = std::collections::HashSet::<(String, String, KnowledgeEdgeKind)>::new();
    edges.retain(|e| seen.insert((e.from.id(), e.to.id(), e.kind)));
    for e in derived_filtered {
        if seen.insert((e.from.id(), e.to.id(), e.kind)) {
            edges.push(e);
        }
    }

    edges.sort_by(|a, b| {
        a.kind
            .as_str()
            .cmp(b.kind.as_str())
            .then_with(|| a.from.category.cmp(&b.from.category))
            .then_with(|| a.from.key.cmp(&b.from.key))
            .then_with(|| a.to.category.cmp(&b.to.category))
            .then_with(|| a.to.key.cmp(&b.to.key))
    });

    let truncated = edges.len() > limit;
    if truncated {
        edges.truncate(limit);
    }

    let mut out = format_mermaid(&edges);
    if truncated {
        out = format!("%% truncated to {limit} edges\n{out}");
    }
    out
}