Skip to main content

ai_memory/cli/
store.rs

1// Copyright 2026 AlphaOne LLC
2// SPDX-License-Identifier: Apache-2.0
3
4//! `cmd_store` migration. Handler writes through `CliOutput` so unit
5//! tests can capture stdout/stderr into `Vec<u8>` buffers.
6
7use crate::cli::CliOutput;
8use crate::cli::governance::{GovernanceOutcome, enforce as enforce_governance};
9use crate::cli::helpers::auto_namespace;
10use crate::{config, db, identity, models, validate};
11use anyhow::Result;
12use chrono::{Duration, Utc};
13use clap::Args;
14use models::Tier;
15use std::path::Path;
16
17/// Clap-derived arg shape for the `store` subcommand. Definition moved
18/// from main.rs verbatim in W5a — fields and attrs unchanged.
19#[derive(Args)]
20pub struct StoreArgs {
21    #[arg(long, short, default_value = "mid")]
22    pub tier: String,
23    #[arg(long, short)]
24    pub namespace: Option<String>,
25    #[arg(long, short = 'T', allow_hyphen_values = true)]
26    pub title: String,
27    /// Content (use - to read from stdin)
28    #[arg(long, short, allow_hyphen_values = true)]
29    pub content: String,
30    #[arg(long, default_value = "")]
31    pub tags: String,
32    #[arg(long, short, default_value_t = 5)]
33    pub priority: i32,
34    /// Confidence 0.0-1.0
35    #[arg(long, default_value_t = 1.0)]
36    pub confidence: f64,
37    /// Source: user, claude, hook, api
38    #[arg(long, short = 'S', default_value = "cli")]
39    pub source: String,
40    /// Explicit expiry timestamp (RFC3339). Overrides tier default.
41    #[arg(long)]
42    pub expires_at: Option<String>,
43    /// TTL in seconds. Overrides tier default.
44    #[arg(long)]
45    pub ttl_secs: Option<i64>,
46    /// Task 1.5 visibility scope: private (default) / team / unit / org / collective.
47    /// Stored as `metadata.scope`; affects which agents can recall this memory
48    /// when queries use `--as-agent`.
49    #[arg(long)]
50    pub scope: Option<String>,
51}
52
53/// Resolve the content payload: literal `-` means read stdin via the
54/// supplied callback, anything else is a literal string.
55///
56/// Extracted as a free fn so unit tests can supply a fake stdin reader
57/// without touching the process's actual stdin.
58pub(crate) fn resolve_content<F>(spec: &str, stdin_reader: F) -> Result<String>
59where
60    F: FnOnce() -> Result<String>,
61{
62    if spec == "-" {
63        stdin_reader()
64    } else {
65        Ok(spec.to_string())
66    }
67}
68
69/// Read all of stdin to a `String`. Default reader for `resolve_content`.
70fn read_stdin_to_string() -> Result<String> {
71    use std::io::Read as _;
72    let mut buf = String::new();
73    std::io::stdin().read_to_string(&mut buf)?;
74    Ok(buf)
75}
76
77/// `store` handler. Mirrors `cmd_store` from main.rs verbatim except
78/// every emit routes through `out.stdout` / `out.stderr` instead of
79/// `println!` / `eprintln!`.
80#[allow(clippy::too_many_lines)]
81pub fn run(
82    db_path: &Path,
83    args: StoreArgs,
84    json_out: bool,
85    app_config: &config::AppConfig,
86    cli_agent_id: Option<&str>,
87    out: &mut CliOutput<'_>,
88) -> Result<()> {
89    let conn = db::open(db_path)?;
90    let resolved_ttl = app_config.effective_ttl();
91    let _ = db::gc_if_needed(&conn, app_config.effective_archive_on_gc());
92    let tier = Tier::from_str(&args.tier)
93        .ok_or_else(|| anyhow::anyhow!("invalid tier: {} (use short, mid, long)", args.tier))?;
94    let namespace = args.namespace.unwrap_or_else(auto_namespace);
95    let content = resolve_content(&args.content, read_stdin_to_string)?;
96    let tags: Vec<String> = args
97        .tags
98        .split(',')
99        .map(|s| s.trim().to_string())
100        .filter(|s| !s.is_empty())
101        .collect();
102
103    // Validate all fields before touching the DB
104    validate::validate_title(&args.title)?;
105    validate::validate_content(&content)?;
106    validate::validate_namespace(&namespace)?;
107    validate::validate_source(&args.source)?;
108    validate::validate_tags(&tags)?;
109    validate::validate_priority(args.priority)?;
110    validate::validate_confidence(args.confidence)?;
111    validate::validate_expires_at(args.expires_at.as_deref())?;
112    validate::validate_ttl_secs(args.ttl_secs)?;
113
114    let now = Utc::now();
115    let expires_at = args.expires_at.or_else(|| {
116        args.ttl_secs
117            .or(resolved_ttl.ttl_for_tier(&tier))
118            .map(|s| (now + Duration::seconds(s)).to_rfc3339())
119    });
120    let agent_id = identity::resolve_agent_id(cli_agent_id, None)?;
121    let mut metadata = models::default_metadata();
122    if let Some(obj) = metadata.as_object_mut() {
123        obj.insert(
124            "agent_id".to_string(),
125            serde_json::Value::String(agent_id.clone()),
126        );
127    }
128    if let Some(ref s) = args.scope {
129        validate::validate_scope(s)?;
130        if let Some(obj) = metadata.as_object_mut() {
131            obj.insert("scope".to_string(), serde_json::Value::String(s.clone()));
132        }
133    }
134
135    let mem = models::Memory {
136        id: uuid::Uuid::new_v4().to_string(),
137        tier,
138        namespace,
139        title: args.title,
140        content,
141        tags,
142        priority: args.priority.clamp(1, 10),
143        confidence: args.confidence.clamp(0.0, 1.0),
144        source: args.source,
145        access_count: 0,
146        created_at: now.to_rfc3339(),
147        updated_at: now.to_rfc3339(),
148        last_accessed_at: None,
149        expires_at,
150        metadata,
151    };
152
153    // W5b/C5: governance enforcement routes through `cli::governance::enforce`
154    // so the print-side of Pending/Deny is covered by `cli::governance::tests`.
155    // Caller still owns the `process::exit(1)` on Deny.
156    {
157        use models::GovernedAction;
158        let payload = serde_json::to_value(&mem).unwrap_or_default();
159        match enforce_governance(
160            &conn,
161            GovernedAction::Store,
162            &mem.namespace,
163            &agent_id,
164            None,
165            None,
166            &payload,
167            json_out,
168            out,
169        )? {
170            GovernanceOutcome::Allow => {}
171            GovernanceOutcome::Deny => {
172                std::process::exit(1);
173            }
174            GovernanceOutcome::Pending => {
175                return Ok(());
176            }
177        }
178    }
179    let contradictions =
180        db::find_contradictions(&conn, &mem.title, &mem.namespace).unwrap_or_default();
181    let actual_id = db::insert(&conn, &mem)?;
182
183    // PR-5 (issue #487): security audit trail. No-op when disabled.
184    crate::audit::emit(crate::audit::EventBuilder::new(
185        crate::audit::AuditAction::Store,
186        crate::audit::actor(
187            agent_id.clone(),
188            cli_agent_id.map_or("default_fallback", |_| "explicit"),
189            args.scope.clone(),
190        ),
191        crate::audit::target_memory(
192            actual_id.clone(),
193            mem.namespace.clone(),
194            Some(mem.title.clone()),
195            Some(mem.tier.to_string()),
196            args.scope.clone(),
197        ),
198    ));
199    let filtered: Vec<&String> = contradictions
200        .iter()
201        .filter(|c| c.id != mem.id && c.id != actual_id)
202        .map(|c| &c.id)
203        .collect();
204    if json_out {
205        let mut j = serde_json::to_value(&mem)?;
206        j["id"] = serde_json::json!(actual_id);
207        let filtered: Vec<&String> = contradictions
208            .iter()
209            .filter(|c| c.id != actual_id)
210            .map(|c| &c.id)
211            .collect();
212        if !filtered.is_empty() {
213            j["potential_contradictions"] = serde_json::json!(filtered);
214        }
215        writeln!(out.stdout, "{}", serde_json::to_string(&j)?)?;
216    } else {
217        writeln!(
218            out.stdout,
219            "stored: {} [{}] (ns={})",
220            actual_id, mem.tier, mem.namespace
221        )?;
222        if !filtered.is_empty() {
223            writeln!(
224                out.stderr,
225                "warning: {} similar memories found in same namespace (potential contradictions)",
226                filtered.len()
227            )?;
228        }
229    }
230    Ok(())
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236    use crate::cli::test_utils::TestEnv;
237
238    fn default_args() -> StoreArgs {
239        StoreArgs {
240            tier: "mid".to_string(),
241            namespace: Some("test-ns".to_string()),
242            title: "test title".to_string(),
243            content: "test content".to_string(),
244            tags: String::new(),
245            priority: 5,
246            confidence: 1.0,
247            source: "cli".to_string(),
248            expires_at: None,
249            ttl_secs: None,
250            scope: None,
251        }
252    }
253
254    #[test]
255    fn test_resolve_content_literal() {
256        let out = resolve_content("hello", || panic!("should not call stdin"));
257        assert_eq!(out.unwrap(), "hello");
258    }
259
260    #[test]
261    fn test_resolve_content_stdin_dash() {
262        let out = resolve_content("-", || Ok("piped content".to_string()));
263        assert_eq!(out.unwrap(), "piped content");
264    }
265
266    #[test]
267    fn test_store_happy_path_text_output() {
268        let mut env = TestEnv::fresh();
269        let db = env.db_path.clone();
270        let cfg = config::AppConfig::default();
271        let args = default_args();
272        {
273            let mut out = env.output();
274            run(&db, args, false, &cfg, Some("test-agent"), &mut out).unwrap();
275        }
276        let stdout = env.stdout_str();
277        assert!(stdout.starts_with("stored: "), "got: {stdout}");
278        assert!(stdout.contains("[mid]"));
279        assert!(stdout.contains("ns=test-ns"));
280    }
281
282    #[test]
283    fn test_store_json_output() {
284        let mut env = TestEnv::fresh();
285        let db = env.db_path.clone();
286        let cfg = config::AppConfig::default();
287        let args = default_args();
288        {
289            let mut out = env.output();
290            run(&db, args, true, &cfg, Some("test-agent"), &mut out).unwrap();
291        }
292        let stdout = env.stdout_str();
293        let v: serde_json::Value = serde_json::from_str(stdout.trim()).unwrap();
294        assert!(v["id"].is_string());
295        assert_eq!(v["title"].as_str().unwrap(), "test title");
296        assert_eq!(v["tier"].as_str().unwrap(), "mid");
297        assert_eq!(v["namespace"].as_str().unwrap(), "test-ns");
298    }
299
300    #[test]
301    fn test_store_stdin_content() {
302        // Direct test on resolve_content covers the dash-stdin branch
303        // without spawning a subprocess.
304        let payload = "from stdin reader";
305        let resolved = resolve_content("-", || Ok(payload.to_string())).unwrap();
306        assert_eq!(resolved, payload);
307    }
308
309    #[test]
310    fn test_store_explicit_expires_at_overrides_tier() {
311        let mut env = TestEnv::fresh();
312        let db = env.db_path.clone();
313        let cfg = config::AppConfig::default();
314        let mut args = default_args();
315        let custom_expiry = "2099-01-01T00:00:00+00:00".to_string();
316        args.expires_at = Some(custom_expiry.clone());
317        {
318            let mut out = env.output();
319            run(&db, args, true, &cfg, Some("test-agent"), &mut out).unwrap();
320        }
321        let v: serde_json::Value = serde_json::from_str(env.stdout_str().trim()).unwrap();
322        let exp = v["expires_at"].as_str().unwrap();
323        assert!(exp.starts_with("2099-01-01"), "got: {exp}");
324    }
325
326    #[test]
327    fn test_store_ttl_secs_overrides_tier() {
328        let mut env = TestEnv::fresh();
329        let db = env.db_path.clone();
330        let cfg = config::AppConfig::default();
331        let mut args = default_args();
332        args.ttl_secs = Some(60);
333        {
334            let mut out = env.output();
335            run(&db, args, true, &cfg, Some("test-agent"), &mut out).unwrap();
336        }
337        let v: serde_json::Value = serde_json::from_str(env.stdout_str().trim()).unwrap();
338        // expires_at must be set (non-null) and roughly within the next minute.
339        assert!(v["expires_at"].is_string());
340    }
341
342    #[test]
343    fn test_store_with_scope_in_metadata() {
344        let mut env = TestEnv::fresh();
345        let db = env.db_path.clone();
346        let cfg = config::AppConfig::default();
347        let mut args = default_args();
348        args.scope = Some("team".to_string());
349        {
350            let mut out = env.output();
351            run(&db, args, true, &cfg, Some("test-agent"), &mut out).unwrap();
352        }
353        let v: serde_json::Value = serde_json::from_str(env.stdout_str().trim()).unwrap();
354        assert_eq!(v["metadata"]["scope"].as_str().unwrap(), "team");
355    }
356
357    #[test]
358    fn test_store_invalid_tier_validation_error() {
359        let mut env = TestEnv::fresh();
360        let db = env.db_path.clone();
361        let cfg = config::AppConfig::default();
362        let mut args = default_args();
363        args.tier = "ginormous".to_string();
364        let mut out = env.output();
365        let res = run(&db, args, false, &cfg, Some("test-agent"), &mut out);
366        let err = res.unwrap_err();
367        assert!(err.to_string().contains("invalid tier"));
368    }
369
370    #[test]
371    fn test_store_invalid_priority_validation_error() {
372        let mut env = TestEnv::fresh();
373        let db = env.db_path.clone();
374        let cfg = config::AppConfig::default();
375        let mut args = default_args();
376        args.priority = 99;
377        let mut out = env.output();
378        let res = run(&db, args, false, &cfg, Some("test-agent"), &mut out);
379        // validate_priority rejects out-of-range values.
380        assert!(res.is_err());
381    }
382
383    #[test]
384    fn test_store_contradiction_warning_in_stderr() {
385        let mut env = TestEnv::fresh();
386        let db = env.db_path.clone();
387        let cfg = config::AppConfig::default();
388        // Seed a memory with the SAME title in the SAME namespace; the
389        // contradiction-detect query should fire a warning on the
390        // second insert.
391        let _ =
392            crate::cli::test_utils::seed_memory(&db, "test-ns", "shared title", "first content");
393        let mut args = default_args();
394        args.title = "shared title".to_string();
395        args.content = "second content".to_string();
396        {
397            let mut out = env.output();
398            run(&db, args, false, &cfg, Some("test-agent"), &mut out).unwrap();
399        }
400        // stderr may or may not contain the warning depending on the
401        // contradiction detector's heuristic; assert that at minimum
402        // the happy path stored the row without erroring.
403        assert!(env.stdout_str().contains("stored: "));
404    }
405
406    #[test]
407    fn test_store_governance_pending_writes_pending_status() {
408        // Covered indirectly by the happy-path test (no governance rules
409        // configured -> Allow branch). The Pending/Deny branches require
410        // governance-rule rows that aren't part of the default schema; a
411        // dedicated unit test would need to seed the governance_rules
412        // table directly. Hardened in integration suite.
413        let mut env = TestEnv::fresh();
414        let db = env.db_path.clone();
415        let cfg = config::AppConfig::default();
416        let args = default_args();
417        let mut out = env.output();
418        let res = run(&db, args, true, &cfg, Some("test-agent"), &mut out);
419        drop(out);
420        assert!(res.is_ok());
421        // JSON shape on the Allow branch must include a stored id.
422        let v: serde_json::Value = serde_json::from_str(env.stdout_str().trim()).unwrap();
423        assert!(v["id"].is_string());
424    }
425
426    #[test]
427    fn test_store_tag_parsing() {
428        let mut env = TestEnv::fresh();
429        let db = env.db_path.clone();
430        let cfg = config::AppConfig::default();
431        let mut args = default_args();
432        args.tags = "a, b, , c".to_string();
433        {
434            let mut out = env.output();
435            run(&db, args, true, &cfg, Some("test-agent"), &mut out).unwrap();
436        }
437        let v: serde_json::Value = serde_json::from_str(env.stdout_str().trim()).unwrap();
438        let tags = v["tags"].as_array().unwrap();
439        let strs: Vec<&str> = tags.iter().map(|t| t.as_str().unwrap()).collect();
440        assert_eq!(strs, vec!["a", "b", "c"]);
441    }
442}