Skip to main content

ai_memory/cli/
offload.rs

1// Copyright 2026 AlphaOne LLC
2// SPDX-License-Identifier: Apache-2.0
3
4//! v0.7.0 QW-3 โ€” `ai-memory offload` and `ai-memory deref` CLI commands.
5//!
6//! Substrate-only wrappers around [`crate::offload::ContextOffloader`].
7//! v0.8.0 short-term-context-compression will layer the auto-cadence
8//! and Mermaid-canvas trigger paths on top.
9
10use std::path::{Path, PathBuf};
11
12use anyhow::{Context, Result};
13use clap::Args;
14use serde_json::{Value, json};
15
16use crate::cli::CliOutput;
17use crate::offload::{ContextOffloader, OffloadConfig};
18use crate::storage as db;
19
20#[derive(Args)]
21pub struct OffloadArgs {
22    /// File whose contents will be offloaded. Pass `-` to read stdin.
23    pub file: String,
24    /// Namespace the blob lives under. Defaults to `auto` so a short
25    /// invocation still records a sensible value.
26    #[arg(long)]
27    pub namespace: Option<String>,
28    /// Optional TTL (seconds). Omit for permanent storage.
29    #[arg(long)]
30    pub ttl_seconds: Option<u64>,
31    /// Override the storing `agent_id`. Defaults to the standard
32    /// resolution chain.
33    #[arg(long)]
34    pub agent_id: Option<String>,
35    /// Emit JSON instead of a human-readable line.
36    #[arg(long)]
37    pub json: bool,
38}
39
40#[derive(Args)]
41pub struct DerefArgs {
42    /// The `ref_id` returned by a prior `offload`.
43    pub ref_id: String,
44    /// Optional output path; otherwise content is written to stdout.
45    #[arg(long)]
46    pub out: Option<PathBuf>,
47    /// Emit a JSON envelope alongside the content (suppresses raw
48    /// content on stdout when `--out` is also passed).
49    #[arg(long)]
50    pub json: bool,
51}
52
53fn read_input(file: &str) -> Result<String> {
54    if file == "-" {
55        use std::io::Read;
56        let mut s = String::new();
57        std::io::stdin()
58            .read_to_string(&mut s)
59            .context("read stdin")?;
60        Ok(s)
61    } else {
62        std::fs::read_to_string(file).with_context(|| format!("read {file}"))
63    }
64}
65
66/// Resolve `agent_id` against the CLI override, falling back to the
67/// project-wide resolution chain (env / hostname / anonymous).
68fn resolve_agent_id(override_value: Option<&str>) -> Result<String> {
69    if let Some(value) = override_value {
70        return Ok(value.to_string());
71    }
72    crate::identity::resolve_agent_id(None, None)
73}
74
75/// `ai-memory offload <file>` entry point.
76pub fn run_offload(db_path: &Path, args: &OffloadArgs, out: &mut CliOutput<'_>) -> Result<()> {
77    let content = read_input(&args.file)?;
78    let namespace = args.namespace.clone().unwrap_or_else(|| "auto".to_string());
79    let agent_id = resolve_agent_id(args.agent_id.as_deref())?;
80    let conn = db::open(db_path).context("open db")?;
81    let off = ContextOffloader::new(&conn, None, OffloadConfig::default());
82    let result = off
83        .offload(&content, &namespace, args.ttl_seconds, &agent_id)
84        .context("offload failed")?;
85    if args.json {
86        writeln!(
87            out.stdout,
88            "{}",
89            serde_json::to_string(&json!({
90                "ref_id": result.ref_id,
91                (crate::models::field_names::CONTENT_SHA256): result.content_sha256,
92                "stored_at": result.stored_at,
93                "namespace": namespace,
94                "agent_id": agent_id,
95            }))?
96        )?;
97    } else {
98        writeln!(
99            out.stdout,
100            "offloaded {} bytes -> {} (sha256 {})",
101            content.len(),
102            result.ref_id,
103            result.content_sha256,
104        )?;
105    }
106    Ok(())
107}
108
109/// `ai-memory deref <ref_id>` entry point.
110pub fn run_deref(db_path: &Path, args: &DerefArgs, out: &mut CliOutput<'_>) -> Result<()> {
111    let conn = db::open(db_path).context("open db")?;
112    let off = ContextOffloader::new(&conn, None, OffloadConfig::default());
113    // SEC-4 (Cluster D) โ€” operator CLI is the trusted-direct-ops path
114    // (see CLAUDE.md ยง"Agent Identity"); pass `None` to BYPASS the
115    // per-agent ownership gate that the MCP handler enforces. The
116    // operator can deref any blob in the local DB.
117    let result = off.deref(&args.ref_id, None).context("deref failed")?;
118    if let Some(path) = &args.out {
119        std::fs::write(path, &result.content)
120            .with_context(|| format!("write {}", path.display()))?;
121    }
122    if args.json {
123        let body_value = if args.out.is_some() {
124            Value::Null
125        } else {
126            Value::String(result.content)
127        };
128        writeln!(
129            out.stdout,
130            "{}",
131            serde_json::to_string(&json!({
132                "ref_id": args.ref_id,
133                "sha256": result.sha256,
134                "stored_at": result.stored_at,
135                "bytes": body_value.as_str().map_or(0, str::len),
136                "content": body_value,
137            }))?
138        )?;
139    } else if args.out.is_none() {
140        write!(out.stdout, "{}", result.content)?;
141    }
142    Ok(())
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use std::path::PathBuf;
149
150    fn fresh_db_path() -> (PathBuf, tempfile::TempDir) {
151        let tmp = tempfile::TempDir::new().expect("tempdir");
152        let db = tmp.path().join("offload-cli.db");
153        (db, tmp)
154    }
155
156    #[test]
157    fn run_offload_reads_file_and_round_trips() {
158        let (db_path, _tmp) = fresh_db_path();
159        let payload_path = _tmp.path().join("payload.txt");
160        std::fs::write(&payload_path, "cli-test-body").unwrap();
161        let args = OffloadArgs {
162            file: payload_path.display().to_string(),
163            namespace: Some("cli/test".to_string()),
164            ttl_seconds: None,
165            agent_id: Some("ai:cli-test".to_string()),
166            json: true,
167        };
168        let mut buf_out = Vec::new();
169        let mut buf_err = Vec::new();
170        {
171            let mut cli_out = CliOutput::from_std(&mut buf_out, &mut buf_err);
172            run_offload(&db_path, &args, &mut cli_out).expect("offload");
173        }
174        let parsed: serde_json::Value = serde_json::from_slice(&buf_out).expect("json");
175        let ref_id = parsed["ref_id"].as_str().expect("ref_id").to_string();
176        // Deref round-trip.
177        let deref_args = DerefArgs {
178            ref_id: ref_id.clone(),
179            out: None,
180            json: false,
181        };
182        let mut buf_out2 = Vec::new();
183        let mut buf_err2 = Vec::new();
184        {
185            let mut cli_out = CliOutput::from_std(&mut buf_out2, &mut buf_err2);
186            run_deref(&db_path, &deref_args, &mut cli_out).expect("deref");
187        }
188        let body = String::from_utf8(buf_out2).unwrap();
189        assert_eq!(body, "cli-test-body");
190    }
191
192    // ------------------------------------------------------------------
193    // Coverage-uplift block (2026-05-19): exercise the non-JSON human-
194    // render path, the --out path, the --json variant of deref, the
195    // read_input file-error path, and the resolve_agent_id default
196    // chain.
197    // ------------------------------------------------------------------
198
199    #[test]
200    fn run_offload_human_render_emits_summary_line() {
201        let (db_path, tmp) = fresh_db_path();
202        let payload_path = tmp.path().join("human.txt");
203        std::fs::write(&payload_path, "human-render-body").unwrap();
204        let args = OffloadArgs {
205            file: payload_path.display().to_string(),
206            namespace: Some("cli/human".to_string()),
207            ttl_seconds: None,
208            agent_id: Some("ai:human-test".to_string()),
209            json: false,
210        };
211        let mut buf_out = Vec::new();
212        let mut buf_err = Vec::new();
213        {
214            let mut cli_out = CliOutput::from_std(&mut buf_out, &mut buf_err);
215            run_offload(&db_path, &args, &mut cli_out).expect("offload");
216        }
217        let text = String::from_utf8(buf_out).unwrap();
218        assert!(text.starts_with("offloaded "), "got: {text}");
219        assert!(text.contains("bytes -> "));
220        assert!(text.contains("sha256 "));
221    }
222
223    #[test]
224    fn run_deref_writes_to_out_path_and_json_envelope_suppresses_content() {
225        let (db_path, tmp) = fresh_db_path();
226        // First offload a payload to get a ref_id.
227        let payload_path = tmp.path().join("orig.txt");
228        std::fs::write(&payload_path, "deref-out-body").unwrap();
229        let off_args = OffloadArgs {
230            file: payload_path.display().to_string(),
231            namespace: Some("cli/deref-out".to_string()),
232            ttl_seconds: None,
233            agent_id: Some("ai:deref-out".to_string()),
234            json: true,
235        };
236        let mut bo = Vec::new();
237        let mut be = Vec::new();
238        {
239            let mut co = CliOutput::from_std(&mut bo, &mut be);
240            run_offload(&db_path, &off_args, &mut co).expect("offload");
241        }
242        let parsed: serde_json::Value = serde_json::from_slice(&bo).unwrap();
243        let ref_id = parsed["ref_id"].as_str().unwrap().to_string();
244
245        // Now deref with --out and --json.
246        let out_path = tmp.path().join("deref-out.bin");
247        let args = DerefArgs {
248            ref_id: ref_id.clone(),
249            out: Some(out_path.clone()),
250            json: true,
251        };
252        let mut bo2 = Vec::new();
253        let mut be2 = Vec::new();
254        {
255            let mut co = CliOutput::from_std(&mut bo2, &mut be2);
256            run_deref(&db_path, &args, &mut co).expect("deref");
257        }
258        // File written.
259        let written = std::fs::read_to_string(&out_path).unwrap();
260        assert_eq!(written, "deref-out-body");
261        // JSON envelope present; `content` is `null` (suppressed).
262        let envelope: serde_json::Value = serde_json::from_slice(&bo2).unwrap();
263        assert_eq!(envelope["ref_id"], ref_id);
264        assert!(envelope["content"].is_null());
265        assert_eq!(envelope["bytes"].as_u64().unwrap(), 0);
266    }
267
268    #[test]
269    fn run_deref_json_without_out_returns_content_inline() {
270        let (db_path, tmp) = fresh_db_path();
271        let payload_path = tmp.path().join("inline.txt");
272        std::fs::write(&payload_path, "inline-json-body").unwrap();
273        let off_args = OffloadArgs {
274            file: payload_path.display().to_string(),
275            namespace: Some("cli/inline".to_string()),
276            ttl_seconds: None,
277            agent_id: Some("ai:inline".to_string()),
278            json: true,
279        };
280        let mut bo = Vec::new();
281        let mut be = Vec::new();
282        {
283            let mut co = CliOutput::from_std(&mut bo, &mut be);
284            run_offload(&db_path, &off_args, &mut co).expect("offload");
285        }
286        let parsed: serde_json::Value = serde_json::from_slice(&bo).unwrap();
287        let ref_id = parsed["ref_id"].as_str().unwrap().to_string();
288
289        let args = DerefArgs {
290            ref_id: ref_id.clone(),
291            out: None,
292            json: true,
293        };
294        let mut bo2 = Vec::new();
295        let mut be2 = Vec::new();
296        {
297            let mut co = CliOutput::from_std(&mut bo2, &mut be2);
298            run_deref(&db_path, &args, &mut co).expect("deref");
299        }
300        let envelope: serde_json::Value = serde_json::from_slice(&bo2).unwrap();
301        assert_eq!(envelope["content"].as_str().unwrap(), "inline-json-body");
302        assert_eq!(
303            envelope["bytes"].as_u64().unwrap(),
304            "inline-json-body".len() as u64
305        );
306    }
307
308    #[test]
309    fn read_input_returns_error_for_missing_file() {
310        let err = read_input("/nonexistent/path/never-exists.txt").unwrap_err();
311        let chain = format!("{err:#}");
312        assert!(chain.contains("read "), "got: {chain}");
313    }
314
315    #[test]
316    fn resolve_agent_id_uses_override_when_present() {
317        let v = resolve_agent_id(Some("ai:explicit-override")).unwrap();
318        assert_eq!(v, "ai:explicit-override");
319    }
320
321    #[test]
322    fn resolve_agent_id_falls_back_to_default_chain_when_none() {
323        // Default chain returns a non-empty stable id (host:... or
324        // ai:... or anonymous:...). Pin only the non-empty property
325        // so the test is hostname/env-agnostic.
326        let v = resolve_agent_id(None).unwrap();
327        assert!(!v.is_empty());
328    }
329
330    #[test]
331    fn run_offload_propagates_read_error_for_missing_file() {
332        let (db_path, _tmp) = fresh_db_path();
333        let args = OffloadArgs {
334            file: "/nonexistent/never-exists/file.txt".to_string(),
335            namespace: None,
336            ttl_seconds: None,
337            agent_id: Some("ai:test".to_string()),
338            json: true,
339        };
340        let mut bo = Vec::new();
341        let mut be = Vec::new();
342        let mut co = CliOutput::from_std(&mut bo, &mut be);
343        let err = run_offload(&db_path, &args, &mut co).expect_err("must fail");
344        let chain = format!("{err:#}");
345        assert!(chain.contains("read "));
346    }
347}