ai_memory/cli/commands/
kg_query.rs

1// Copyright 2026 AlphaOne LLC
2// SPDX-License-Identifier: Apache-2.0
3
4//! v0.7.0 ARCH-3 / FX-12 — `ai-memory kg-query` CLI subcommand.
5//!
6//! Closes the three-surface-parity gap on `memory_kg_query`. The MCP
7//! tool ([`crate::mcp::handle_kg_query`]) and the HTTP route landed
8//! previously; this module wires the CLI surface so operators can
9//! traverse the knowledge graph from a terminal without driving MCP
10//! stdio JSON-RPC.
11//!
12//! ## DRY contract
13//!
14//! No business logic lives here — this module is a clap arg-parser
15//! plus an output formatter. The actual traversal semantics live in
16//! [`crate::mcp::handle_kg_query`]. The MCP, HTTP, and CLI surfaces
17//! all share that one implementation.
18
19use anyhow::Result;
20use clap::Args;
21use serde_json::{Value, json};
22
23use crate::cli::CliOutput;
24use crate::storage as db;
25
26/// CLI args for `ai-memory kg-query`. Mirrors the MCP `memory_kg_query`
27/// `input_schema` shape.
28#[derive(Args, Debug, Clone)]
29pub struct KgQueryArgs {
30    /// Source memory id (full UUID or unique prefix).
31    #[arg(long = "source-id", value_name = "ID")]
32    pub source_id: Option<String>,
33
34    /// #889 — list every memory rooted at the given source_uri instead
35    /// of traversing from a specific source memory.
36    #[arg(long = "by-source-uri", value_name = "URI")]
37    pub by_source_uri: Option<String>,
38
39    /// Max hops, 1..=5. Defaults to 1 server-side.
40    #[arg(long = "max-depth", value_name = "N")]
41    pub max_depth: Option<u32>,
42
43    /// Restrict to a specific namespace.
44    #[arg(long, value_name = "NS")]
45    pub namespace: Option<String>,
46
47    /// RFC3339 timestamp — keep only links valid at this instant.
48    #[arg(long = "valid-at", value_name = "RFC3339")]
49    pub valid_at: Option<String>,
50
51    /// Comma-separated allowlist of observed_by agent ids.
52    #[arg(long = "allowed-agents", value_name = "CSV")]
53    pub allowed_agents: Option<String>,
54
55    /// Hard cap across all depths (1..=1000).
56    #[arg(long, value_name = "N")]
57    pub limit: Option<u32>,
58
59    /// When set, traverse historically-invalidated edges as well.
60    #[arg(long = "include-invalidated")]
61    pub include_invalidated: bool,
62
63    /// Emit the raw JSON envelope (the same shape MCP / HTTP return)
64    /// instead of a human-readable table.
65    #[arg(long)]
66    pub json: bool,
67}
68
69/// `ai-memory kg-query` dispatch entry. Opens the DB at `db_path`,
70/// builds the MCP-shaped JSON params bag, and routes through the
71/// shared substrate primitive — guaranteeing the wire envelope is
72/// byte-equal across MCP / HTTP / CLI.
73///
74/// # Errors
75///
76/// - The DB at `db_path` cannot be opened.
77/// - The substrate validation rejects the supplied params.
78/// - `serde_json::to_string` cannot serialise the envelope (in
79///   practice never happens with the shapes used here).
80pub fn cmd_kg_query(
81    db_path: &std::path::Path,
82    args: &KgQueryArgs,
83    out: &mut CliOutput<'_>,
84) -> Result<()> {
85    if args.source_id.is_none() && args.by_source_uri.is_none() {
86        anyhow::bail!("kg-query: either --source-id or --by-source-uri is required");
87    }
88    let conn = db::open(db_path)?;
89
90    let mut params = json!({});
91    if let Some(sid) = &args.source_id {
92        params["source_id"] = json!(sid);
93    }
94    if let Some(uri) = &args.by_source_uri {
95        params[crate::models::field_names::BY_SOURCE_URI] = json!(uri);
96    }
97    if let Some(d) = args.max_depth {
98        params["max_depth"] = json!(d);
99    }
100    if let Some(ns) = &args.namespace {
101        params["namespace"] = json!(ns);
102    }
103    if let Some(t) = &args.valid_at {
104        params["valid_at"] = json!(t);
105    }
106    if let Some(csv) = &args.allowed_agents {
107        let agents: Vec<&str> = csv
108            .split(',')
109            .map(str::trim)
110            .filter(|s| !s.is_empty())
111            .collect();
112        params["allowed_agents"] = json!(agents);
113    }
114    if let Some(l) = args.limit {
115        params["limit"] = json!(l);
116    }
117    if args.include_invalidated {
118        params[crate::models::field_names::INCLUDE_INVALIDATED] = json!(true);
119    }
120
121    let envelope = crate::mcp::handle_kg_query(&conn, &params)
122        .map_err(|e| anyhow::anyhow!("kg-query: {e}"))?;
123
124    if args.json {
125        writeln!(out.stdout, "{}", serde_json::to_string(&envelope)?)?;
126        return Ok(());
127    }
128
129    // Human-readable summary.
130    let count = envelope.get("count").and_then(Value::as_u64).unwrap_or(0);
131    writeln!(out.stdout, "kg-query: {count} row(s)")?;
132    if let Some(arr) = envelope.get("memories").and_then(Value::as_array) {
133        for m in arr {
134            let target = m.get("target_id").and_then(Value::as_str).unwrap_or("?");
135            let title = m.get("title").and_then(Value::as_str).unwrap_or("");
136            let depth = m.get("depth").and_then(Value::as_u64).unwrap_or(0);
137            let relation = m.get("relation").and_then(Value::as_str).unwrap_or("");
138            writeln!(out.stdout, "  [d={depth}] {target}  {relation}  {title}",)?;
139        }
140    }
141    Ok(())
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use crate::cli::test_utils::{TestEnv, seed_memory};
148
149    #[test]
150    fn kg_query_cli_requires_source_or_uri() {
151        let mut env = TestEnv::fresh();
152        let db = env.db_path.clone();
153        let args = KgQueryArgs {
154            source_id: None,
155            by_source_uri: None,
156            max_depth: None,
157            namespace: None,
158            valid_at: None,
159            allowed_agents: None,
160            limit: None,
161            include_invalidated: false,
162            json: true,
163        };
164        let mut out = env.output();
165        let err = cmd_kg_query(&db, &args, &mut out).expect_err("must fail");
166        assert!(err.to_string().contains("source-id"), "got: {err}");
167    }
168
169    #[test]
170    fn kg_query_cli_empty_db_returns_zero_rows() {
171        let mut env = TestEnv::fresh();
172        let db = env.db_path.clone();
173        let src_id = seed_memory(&db, "ns", "kg-source", "content");
174        let args = KgQueryArgs {
175            source_id: Some(src_id),
176            by_source_uri: None,
177            max_depth: None,
178            namespace: None,
179            valid_at: None,
180            allowed_agents: None,
181            limit: None,
182            include_invalidated: false,
183            json: true,
184        };
185        {
186            let mut out = env.output();
187            cmd_kg_query(&db, &args, &mut out).expect("kg-query ok");
188        }
189        let stdout = env.stdout_str();
190        let envelope: Value = serde_json::from_str(stdout.trim()).expect("parse envelope");
191        assert_eq!(envelope["count"].as_u64(), Some(0));
192    }
193
194    /// Seed source + target memories with a `related_to` edge, then
195    /// drive the text-output path (the memory-row loop) plus every
196    /// optional param arm.
197    #[test]
198    fn kg_query_cli_text_output_with_rows_and_all_params() {
199        let mut env = TestEnv::fresh();
200        let db = env.db_path.clone();
201        let src = seed_memory(&db, "ns", "kg-src", "content");
202        let tgt = seed_memory(&db, "ns", "kg-tgt", "target content");
203        {
204            let conn = db::open(&db).unwrap();
205            let now = chrono::Utc::now().to_rfc3339();
206            conn.execute(
207                "INSERT INTO memory_links (source_id, target_id, relation, created_at, valid_from)
208                 VALUES (?1, ?2, 'related_to', ?3, ?3)",
209                rusqlite::params![src, tgt, now],
210            )
211            .expect("insert link");
212        }
213        let args = KgQueryArgs {
214            source_id: Some(src),
215            by_source_uri: None,
216            max_depth: Some(2),
217            namespace: Some("ns".into()),
218            valid_at: None,
219            allowed_agents: None,
220            limit: Some(100),
221            include_invalidated: true,
222            json: false,
223        };
224        {
225            let mut out = env.output();
226            cmd_kg_query(&db, &args, &mut out).expect("kg-query ok");
227        }
228        let stdout = env.stdout_str();
229        assert!(stdout.contains("row(s)"), "got: {stdout}");
230        assert!(stdout.contains("related_to"), "got: {stdout}");
231        assert!(stdout.contains("kg-tgt"), "got: {stdout}");
232    }
233
234    #[test]
235    fn kg_query_cli_by_source_uri_path() {
236        let mut env = TestEnv::fresh();
237        let db = env.db_path.clone();
238        seed_memory(&db, "ns", "kg-uri", "content");
239        let args = KgQueryArgs {
240            source_id: None,
241            by_source_uri: Some("doc://nonexistent".into()),
242            max_depth: None,
243            namespace: None,
244            // Exercise the valid_at + allowed_agents param-build arms
245            // (they thread into the params bag regardless of row count).
246            valid_at: Some(chrono::Utc::now().to_rfc3339()),
247            allowed_agents: Some("test-agent, , ai:other".into()),
248            limit: None,
249            include_invalidated: false,
250            json: true,
251        };
252        {
253            let mut out = env.output();
254            cmd_kg_query(&db, &args, &mut out).expect("kg-query ok");
255        }
256        let envelope: Value = serde_json::from_str(env.stdout_str().trim()).expect("json");
257        assert_eq!(envelope["count"].as_u64(), Some(0));
258    }
259}
ai_memory/cli/commands/kg_query.rs

ai_memory/cli/commands/
kg_query.rs
