1use std::collections::BTreeMap;
17use std::path::{Path, PathBuf};
18
19use chrono::{DateTime, Utc};
20use serde::{Deserialize, Serialize};
21
22use crate::config::ZigConfig;
23use crate::error::ZigError;
24use crate::paths;
25use crate::workflow::model::MemoryMode;
26
27fn validate_entry_filename(name: &str) -> Result<(), ZigError> {
32 if name.is_empty() {
33 return Err(ZigError::Validation("name must not be empty".into()));
34 }
35 if name.contains('/')
36 || name.contains('\\')
37 || name.contains('\0')
38 || name == "."
39 || name == ".."
40 || name.starts_with('-')
41 {
42 return Err(ZigError::Validation(format!(
43 "name '{name}' must not contain path separators or traversal segments"
44 )));
45 }
46 Ok(())
47}
48
49#[derive(Debug, Clone, Default, Serialize, Deserialize)]
55pub struct Manifest {
56 pub next_id: u64,
58 pub entries: BTreeMap<String, MemoryEntry>,
60}
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
64pub struct MemoryEntry {
65 pub name: String,
67 pub file: String,
69 #[serde(default, skip_serializing_if = "Option::is_none")]
71 pub description: Option<String>,
72 #[serde(default, skip_serializing_if = "Vec::is_empty")]
74 pub tags: Vec<String>,
75 #[serde(default, skip_serializing_if = "Option::is_none")]
77 pub step: Option<String>,
78 pub source: String,
80 pub added: DateTime<Utc>,
82}
83
84#[derive(Debug, Clone)]
86pub enum MemoryTarget {
87 GlobalShared,
89 GlobalWorkflow(String),
91 Cwd,
93}
94
95impl MemoryTarget {
96 pub fn from_flags(workflow: Option<&str>, global: bool, cwd: bool) -> Result<Self, ZigError> {
98 if let Some(name) = workflow {
99 if cwd {
100 return Err(ZigError::Validation(
101 "--workflow cannot be combined with --cwd".into(),
102 ));
103 }
104 return Ok(MemoryTarget::GlobalWorkflow(name.to_string()));
105 }
106 if cwd {
107 return Ok(MemoryTarget::Cwd);
108 }
109 if global {
110 return Ok(MemoryTarget::GlobalShared);
111 }
112 Ok(MemoryTarget::Cwd)
114 }
115
116 pub fn ensure_dir(&self) -> Result<PathBuf, ZigError> {
118 match self {
119 MemoryTarget::GlobalShared => paths::ensure_global_memory_dir(Some("_shared")),
120 MemoryTarget::GlobalWorkflow(name) => paths::ensure_global_memory_dir(Some(name)),
121 MemoryTarget::Cwd => ensure_cwd_memory_dir(),
122 }
123 }
124
125 pub fn existing_dir(&self) -> Option<PathBuf> {
127 match self {
128 MemoryTarget::GlobalShared => paths::global_shared_memory_dir(),
129 MemoryTarget::GlobalWorkflow(name) => paths::global_memory_for(name),
130 MemoryTarget::Cwd => paths::cwd_memory_dir().or_else(|| {
131 std::env::current_dir()
132 .ok()
133 .map(|p| p.join(".zig").join("memory"))
134 }),
135 }
136 }
137
138 pub fn label(&self) -> String {
140 match self {
141 MemoryTarget::GlobalShared => "global:_shared".to_string(),
142 MemoryTarget::GlobalWorkflow(n) => format!("global:{n}"),
143 MemoryTarget::Cwd => "cwd".to_string(),
144 }
145 }
146}
147
148#[derive(Debug, Clone, Copy, PartialEq, Eq)]
150pub enum SearchScope {
151 Sentence,
153 Paragraph,
155 Section,
157 File,
159}
160
161fn manifest_path(dir: &Path) -> PathBuf {
166 dir.join(".manifest")
167}
168
169pub fn load_manifest(dir: &Path) -> Result<Manifest, ZigError> {
172 let path = manifest_path(dir);
173 if !path.exists() {
174 return Ok(Manifest {
175 next_id: 1,
176 entries: BTreeMap::new(),
177 });
178 }
179 let content = std::fs::read_to_string(&path)
180 .map_err(|e| ZigError::Io(format!("failed to read {}: {e}", path.display())))?;
181 serde_json::from_str(&content).map_err(|e| {
182 ZigError::Io(format!(
183 "failed to parse manifest at {}: {e}",
184 path.display()
185 ))
186 })
187}
188
189pub fn save_manifest(dir: &Path, manifest: &Manifest) -> Result<(), ZigError> {
191 let path = manifest_path(dir);
192 let content = serde_json::to_string_pretty(manifest)
193 .map_err(|e| ZigError::Serialize(format!("failed to serialize manifest: {e}")))?;
194 std::fs::write(&path, content)
195 .map_err(|e| ZigError::Io(format!("failed to write {}: {e}", path.display())))
196}
197
198fn ensure_cwd_memory_dir() -> Result<PathBuf, ZigError> {
199 if let Some(existing) = paths::cwd_memory_dir() {
200 return Ok(existing);
201 }
202 let cwd = std::env::current_dir()
203 .map_err(|e| ZigError::Io(format!("failed to read current directory: {e}")))?;
204 let dir = cwd.join(".zig").join("memory");
205 std::fs::create_dir_all(&dir)
206 .map_err(|e| ZigError::Io(format!("failed to create {}: {e}", dir.display())))?;
207 Ok(dir)
208}
209
210fn candidate_dirs(workflow: Option<&str>) -> Vec<(String, PathBuf)> {
216 let mut out: Vec<(String, PathBuf)> = Vec::new();
217 if let Some(d) = paths::global_shared_memory_dir() {
218 out.push(("global:_shared".into(), d));
219 }
220 if let Some(name) = workflow {
221 if let Some(d) = paths::global_memory_for(name) {
222 out.push((format!("global:{name}"), d));
223 }
224 }
225 if let Some(d) = paths::cwd_memory_dir() {
226 out.push(("cwd".into(), d));
227 } else if let Ok(cwd) = std::env::current_dir() {
228 out.push(("cwd".into(), cwd.join(".zig").join("memory")));
229 }
230 out
231}
232
233fn find_entry_across_tiers(
236 id: u64,
237 workflow: Option<&str>,
238) -> Result<(PathBuf, String, Manifest, MemoryEntry), ZigError> {
239 let id_str = id.to_string();
240 let dirs = candidate_dirs(workflow);
242 for (label, dir) in dirs.iter().rev() {
243 if !dir.is_dir() {
244 continue;
245 }
246 let manifest = load_manifest(dir)?;
247 if let Some(entry) = manifest.entries.get(&id_str).cloned() {
248 return Ok((dir.clone(), label.clone(), manifest, entry));
249 }
250 }
251 Err(ZigError::Io(format!(
252 "memory entry with id {id} not found in any tier"
253 )))
254}
255
256pub fn add(
265 file_path: &str,
266 target: MemoryTarget,
267 step: Option<&str>,
268 name: Option<&str>,
269 description: Option<&str>,
270 tags: &[String],
271) -> Result<u64, ZigError> {
272 let src = Path::new(file_path);
273 if !src.exists() {
274 return Err(ZigError::Io(format!("source file not found: {file_path}")));
275 }
276 if !src.is_file() {
277 return Err(ZigError::Io(format!("not a regular file: {file_path}")));
278 }
279
280 let dir = target.ensure_dir()?;
281 let mut manifest = load_manifest(&dir)?;
282
283 let id = manifest.next_id;
284 manifest.next_id += 1;
285
286 let file_name = name
287 .map(str::to_string)
288 .or_else(|| src.file_name().map(|n| n.to_string_lossy().into_owned()))
289 .ok_or_else(|| ZigError::Io(format!("could not derive a name from {}", src.display())))?;
290 validate_entry_filename(&file_name)?;
291
292 let dest = dir.join(&file_name);
293 if dest.exists() {
294 return Err(ZigError::Io(format!(
295 "file '{}' already exists in {} — remove it first or use --name to rename",
296 file_name,
297 dir.display()
298 )));
299 }
300
301 std::fs::copy(src, &dest).map_err(|e| {
302 ZigError::Io(format!(
303 "failed to copy {} → {}: {e}",
304 src.display(),
305 dest.display()
306 ))
307 })?;
308
309 let source_abs = std::fs::canonicalize(src)
310 .unwrap_or_else(|_| src.to_path_buf())
311 .display()
312 .to_string();
313
314 let entry = MemoryEntry {
315 name: file_name.clone(),
316 file: file_name,
317 description: description.map(str::to_string),
318 tags: tags.to_vec(),
319 step: step.map(str::to_string),
320 source: source_abs,
321 added: Utc::now(),
322 };
323
324 manifest.entries.insert(id.to_string(), entry);
325 save_manifest(&dir, &manifest)?;
326
327 println!(
328 "added memory entry id={id} '{}' to {}",
329 manifest.entries[&id.to_string()].name,
330 target.label()
331 );
332
333 if description.is_none() {
334 eprintln!("hint: add a description with `zig memory update {id} --description \"...\"`");
335 }
336
337 Ok(id)
338}
339
340pub fn update(
342 id: u64,
343 workflow: Option<&str>,
344 name: Option<&str>,
345 description: Option<&str>,
346 tags: Option<&[String]>,
347) -> Result<(), ZigError> {
348 let (dir, label, mut manifest, _entry) = find_entry_across_tiers(id, workflow)?;
349 let id_str = id.to_string();
350
351 let entry = manifest
352 .entries
353 .get_mut(&id_str)
354 .ok_or_else(|| ZigError::Io(format!("memory entry {id} vanished during update")))?;
355
356 if let Some(n) = name {
357 validate_entry_filename(n)?;
358 let old_path = dir.join(&entry.file);
360 let new_path = dir.join(n);
361 if old_path != new_path {
362 if new_path.exists() {
363 return Err(ZigError::Io(format!(
364 "cannot rename: '{}' already exists in {}",
365 n,
366 dir.display()
367 )));
368 }
369 std::fs::rename(&old_path, &new_path).map_err(|e| {
370 ZigError::Io(format!(
371 "failed to rename {} → {}: {e}",
372 old_path.display(),
373 new_path.display()
374 ))
375 })?;
376 entry.file = n.to_string();
377 }
378 entry.name = n.to_string();
379 }
380 if let Some(d) = description {
381 entry.description = Some(d.to_string());
382 }
383 if let Some(t) = tags {
384 entry.tags = t.to_vec();
385 }
386
387 save_manifest(&dir, &manifest)?;
388 println!("updated memory entry id={id} in {label}");
389 Ok(())
390}
391
392pub fn delete(id: u64, workflow: Option<&str>) -> Result<(), ZigError> {
394 let (dir, label, mut manifest, entry) = find_entry_across_tiers(id, workflow)?;
395 let id_str = id.to_string();
396
397 let file_path = dir.join(&entry.file);
398 if file_path.is_file() {
399 std::fs::remove_file(&file_path)
400 .map_err(|e| ZigError::Io(format!("failed to remove {}: {e}", file_path.display())))?;
401 }
402
403 manifest.entries.remove(&id_str);
404 save_manifest(&dir, &manifest)?;
405 println!("deleted memory entry id={id} '{}' from {label}", entry.name);
406 Ok(())
407}
408
409pub fn show(id: u64, workflow: Option<&str>) -> Result<(), ZigError> {
411 let (dir, label, _manifest, entry) = find_entry_across_tiers(id, workflow)?;
412 let file_path = dir.join(&entry.file);
413
414 println!("id: {id}");
415 println!("name: {}", entry.name);
416 println!("tier: {label}");
417 println!("source: {}", entry.source);
418 println!(
419 "added: {}",
420 entry.added.format("%Y-%m-%d %H:%M:%S UTC")
421 );
422 if let Some(ref desc) = entry.description {
423 println!("description: {desc}");
424 }
425 if !entry.tags.is_empty() {
426 println!("tags: {}", entry.tags.join(", "));
427 }
428 if let Some(ref step) = entry.step {
429 println!("step: {step}");
430 }
431
432 if file_path.is_file() {
433 let contents = std::fs::read_to_string(&file_path)
434 .map_err(|e| ZigError::Io(format!("failed to read {}: {e}", file_path.display())))?;
435 println!("\n--- contents ({}) ---", file_path.display());
436 print!("{contents}");
437 if !contents.ends_with('\n') {
438 println!();
439 }
440 } else {
441 println!("\n(file not found: {})", file_path.display());
442 }
443
444 Ok(())
445}
446
447pub fn list(workflow: Option<&str>) -> Result<(), ZigError> {
449 let mut rows: Vec<(String, String, String, String, String, String)> = Vec::new();
450
451 let dirs = candidate_dirs(workflow);
452 for (label, dir) in &dirs {
453 if !dir.is_dir() {
454 continue;
455 }
456 let manifest = load_manifest(dir)?;
457 for (id_str, entry) in &manifest.entries {
458 let desc = entry
459 .description
460 .as_deref()
461 .unwrap_or("")
462 .chars()
463 .take(50)
464 .collect::<String>();
465 let tags = entry.tags.join(", ");
466 rows.push((
467 id_str.clone(),
468 entry.name.clone(),
469 tags,
470 desc,
471 label.clone(),
472 entry.step.clone().unwrap_or_default(),
473 ));
474 }
475 }
476
477 if rows.is_empty() {
478 println!("No memory entries found.");
479 println!("Hint: add one with `zig memory add <file> [--workflow <name>]`");
480 return Ok(());
481 }
482
483 let id_w = rows.iter().map(|r| r.0.len()).max().unwrap_or(2).max(2);
484 let name_w = rows.iter().map(|r| r.1.len()).max().unwrap_or(4).max(4);
485 let tags_w = rows.iter().map(|r| r.2.len()).max().unwrap_or(4).max(4);
486 let tier_w = rows.iter().map(|r| r.4.len()).max().unwrap_or(4).max(4);
487
488 println!(
489 "{:<id_w$} {:<name_w$} {:<tags_w$} {:<tier_w$} DESCRIPTION",
490 "ID", "NAME", "TAGS", "TIER",
491 );
492 for (id, name, tags, desc, tier, _step) in &rows {
493 println!(
494 "{:<id_w$} {:<name_w$} {:<tags_w$} {:<tier_w$} {desc}",
495 id, name, tags, tier,
496 );
497 }
498
499 Ok(())
500}
501
502pub fn search(query: &str, scope: SearchScope, workflow: Option<&str>) -> Result<(), ZigError> {
508 let query_lower = query.to_lowercase();
509 let mut found = false;
510
511 let dirs = candidate_dirs(workflow);
512 for (label, dir) in &dirs {
513 if !dir.is_dir() {
514 continue;
515 }
516 let manifest = load_manifest(dir)?;
517 for (id_str, entry) in &manifest.entries {
518 let file_path = dir.join(&entry.file);
519 if !file_path.is_file() {
520 continue;
521 }
522 let content = match std::fs::read_to_string(&file_path) {
523 Ok(c) => c,
524 Err(_) => continue,
525 };
526 if !content.to_lowercase().contains(&query_lower) {
527 continue;
528 }
529
530 let matches = extract_matches(&content, &query_lower, scope);
531 for m in matches {
532 if !found {
533 found = true;
534 }
535 println!(
536 "[id:{} {} {}:{}]",
537 id_str,
538 label,
539 entry.name,
540 m.line_number.unwrap_or(0)
541 );
542 println!("{}", m.text);
543 println!();
544 }
545 }
546 }
547
548 if !found {
549 println!("No matches found for '{query}'.");
550 }
551
552 Ok(())
553}
554
555struct MatchFragment {
556 text: String,
557 line_number: Option<usize>,
558}
559
560fn extract_matches(content: &str, query_lower: &str, scope: SearchScope) -> Vec<MatchFragment> {
561 match scope {
562 SearchScope::Sentence => extract_sentences(content, query_lower),
563 SearchScope::Paragraph => extract_paragraphs(content, query_lower),
564 SearchScope::Section => extract_sections(content, query_lower),
565 SearchScope::File => extract_file(content, query_lower),
566 }
567}
568
569fn extract_sentences(content: &str, query_lower: &str) -> Vec<MatchFragment> {
570 let mut results = Vec::new();
571 let line_starts: Vec<usize> = std::iter::once(0)
573 .chain(content.match_indices('\n').map(|(i, _)| i + 1))
574 .collect();
575
576 let find_line = |byte_offset: usize| -> usize {
577 match line_starts.binary_search(&byte_offset) {
578 Ok(i) => i + 1,
579 Err(i) => i,
580 }
581 };
582
583 let chars: Vec<char> = content.chars().collect();
585 let mut byte_pos = 0;
586 let mut sentence_start_byte = 0;
587
588 for (i, &ch) in chars.iter().enumerate() {
589 let ch_len = ch.len_utf8();
590 if (ch == '.' || ch == '!' || ch == '?')
591 && (i + 1 >= chars.len() || chars[i + 1].is_whitespace())
592 {
593 let sentence_end_byte = byte_pos + ch_len;
594 let sentence = &content[sentence_start_byte..sentence_end_byte];
595 if sentence.to_lowercase().contains(query_lower) {
596 results.push(MatchFragment {
597 text: sentence.trim().to_string(),
598 line_number: Some(find_line(sentence_start_byte)),
599 });
600 }
601 sentence_start_byte = sentence_end_byte;
602 }
603 byte_pos += ch_len;
604 }
605
606 if sentence_start_byte < content.len() {
608 let sentence = &content[sentence_start_byte..];
609 if sentence.to_lowercase().contains(query_lower) {
610 results.push(MatchFragment {
611 text: sentence.trim().to_string(),
612 line_number: Some(find_line(sentence_start_byte)),
613 });
614 }
615 }
616
617 results
618}
619
620fn extract_paragraphs(content: &str, query_lower: &str) -> Vec<MatchFragment> {
621 let mut results = Vec::new();
622 let mut line_num = 1;
623
624 for paragraph in content.split("\n\n") {
625 if paragraph.to_lowercase().contains(query_lower) {
626 results.push(MatchFragment {
627 text: paragraph.trim().to_string(),
628 line_number: Some(line_num),
629 });
630 }
631 line_num += paragraph.matches('\n').count() + 2;
633 }
634
635 results
636}
637
638fn extract_sections(content: &str, query_lower: &str) -> Vec<MatchFragment> {
639 let mut results = Vec::new();
640 let mut sections: Vec<(usize, String)> = Vec::new();
641
642 let mut current_start_line = 1;
643 let mut current_section = String::new();
644 let mut line_num = 0;
645
646 for line in content.lines() {
647 line_num += 1;
648 if line.starts_with("## ") && !current_section.is_empty() {
649 sections.push((current_start_line, current_section.clone()));
650 current_section.clear();
651 current_start_line = line_num;
652 }
653 if !current_section.is_empty() {
654 current_section.push('\n');
655 }
656 current_section.push_str(line);
657 }
658 if !current_section.is_empty() {
659 sections.push((current_start_line, current_section));
660 }
661
662 for (start_line, section) in sections {
663 if section.to_lowercase().contains(query_lower) {
664 results.push(MatchFragment {
665 text: section.trim().to_string(),
666 line_number: Some(start_line),
667 });
668 }
669 }
670
671 results
672}
673
674fn extract_file(content: &str, query_lower: &str) -> Vec<MatchFragment> {
675 if content.to_lowercase().contains(query_lower) {
676 vec![MatchFragment {
677 text: content.trim().to_string(),
678 line_number: Some(1),
679 }]
680 } else {
681 vec![]
682 }
683}
684
685pub struct MemoryCollector {
691 pub global_shared_dir: Option<PathBuf>,
692 pub global_workflow_dir: Option<PathBuf>,
693 pub cwd_memory_dir: Option<PathBuf>,
694 pub workflow_mode: MemoryMode,
696 pub local_enabled: bool,
698 pub disabled: bool,
700}
701
702impl MemoryCollector {
703 pub fn from_env(
705 workflow_name: &str,
706 workflow_mode: MemoryMode,
707 config: &ZigConfig,
708 disabled: bool,
709 ) -> Self {
710 Self {
711 global_shared_dir: paths::global_shared_memory_dir(),
712 global_workflow_dir: paths::global_memory_for(workflow_name),
713 cwd_memory_dir: paths::cwd_memory_dir(),
714 workflow_mode,
715 local_enabled: config.memory.local,
716 disabled,
717 }
718 }
719
720 pub fn collect_for_step(
724 &self,
725 step_memory: Option<&str>,
726 ) -> Result<Vec<(PathBuf, String, MemoryEntry)>, ZigError> {
727 if self.disabled {
728 return Ok(Vec::new());
729 }
730
731 let effective_mode = if step_memory.is_some() {
733 MemoryMode::from_str_opt(step_memory)
734 } else {
735 self.workflow_mode
736 };
737
738 if effective_mode == MemoryMode::None {
739 return Ok(Vec::new());
740 }
741
742 let mut entries = Vec::new();
743 let include_local = effective_mode == MemoryMode::All && self.local_enabled;
744
745 if let Some(dir) = self.global_shared_dir.as_deref() {
747 collect_from_dir(dir, &mut entries)?;
748 }
749
750 if let Some(dir) = self.global_workflow_dir.as_deref() {
752 collect_from_dir(dir, &mut entries)?;
753 }
754
755 if include_local {
757 if let Some(dir) = self.cwd_memory_dir.as_deref() {
758 collect_from_dir(dir, &mut entries)?;
759 }
760 }
761
762 Ok(entries)
763 }
764}
765
766fn collect_from_dir(
767 dir: &Path,
768 out: &mut Vec<(PathBuf, String, MemoryEntry)>,
769) -> Result<(), ZigError> {
770 if !dir.is_dir() {
771 return Ok(());
772 }
773 let manifest = load_manifest(dir)?;
774 for (id_str, entry) in &manifest.entries {
775 let abs_path = dir.join(&entry.file);
776 if abs_path.is_file() {
777 out.push((abs_path, id_str.clone(), entry.clone()));
778 }
779 }
780 Ok(())
781}
782
783pub fn render_memory_block(
787 entries: &[(PathBuf, String, MemoryEntry)],
788 workflow_name: &str,
789 step_name: Option<&str>,
790) -> String {
791 if entries.is_empty() {
792 return String::new();
793 }
794
795 let mut out = String::from("<memory>\n");
796 out.push_str(
797 "You have access to the following memory files — a scratch pad of accumulated knowledge. \
798 Read them with your file tools when relevant.\n",
799 );
800
801 let step_flag = step_name
803 .map(|s| format!(" --step {s}"))
804 .unwrap_or_default();
805 out.push_str(&format!(
806 "To add new memories: `zig memory add <path> --workflow {workflow_name}{step_flag}`\n"
807 ));
808 out.push_str(
809 "To update metadata: `zig memory update <id> --description \"...\" --tags \"...\"`\n\n",
810 );
811
812 for (path, id, entry) in entries {
813 out.push_str("- ");
814 out.push_str(&path.display().to_string());
815 if let Some(desc) = &entry.description {
816 out.push_str(&format!(" (id: {id}) — {desc}"));
817 } else {
818 out.push_str(&format!(
819 " (id: {id}, no description — run: zig memory update {id} --description \"...\")"
820 ));
821 }
822 if !entry.tags.is_empty() {
823 out.push_str(&format!(" [{}]", entry.tags.join(", ")));
824 }
825 out.push('\n');
826 }
827 out.push_str("</memory>\n\n");
828 out
829}
830
831#[cfg(test)]
832#[path = "memory_tests.rs"]
833mod tests;