1use async_trait::async_trait;
2use serde_json::json;
3
4use bamboo_agent_core::tools::{Tool, ToolClass, ToolCtx, ToolError, ToolOutcome, ToolResult};
5use bamboo_agent_core::Session;
6use bamboo_memory::memory_store::{
7 DurableMemoryStatus, MemoryQueryOptions, MemoryScope, MemoryStore, MAX_MAX_CHARS,
8 MAX_QUERY_LIMIT,
9};
10use bamboo_tools::tools::session_memory::{
11 execute_session_memory_action, SessionMemoryAction, MEMORY_SESSION_ACTION_NAMES,
12};
13
14mod args;
15mod parsing;
16
17#[cfg(test)]
18mod tests;
19
20use args::MemoryArgs;
21
22#[derive(Clone)]
23pub struct MemoryTool {
24 session_repo: bamboo_engine::SessionRepository,
25 memory_store: MemoryStore,
26}
27
28impl MemoryTool {
29 pub fn new(
30 session_repo: bamboo_engine::SessionRepository,
31 data_dir: impl Into<std::path::PathBuf>,
32 ) -> Self {
33 Self {
34 session_repo,
35 memory_store: MemoryStore::new(data_dir),
36 }
37 }
38
39 async fn session_for_context(&self, session_id: Option<&str>) -> Option<Session> {
40 self.session_repo.load(session_id?).await
41 }
42
43 async fn resolve_project_key(
44 &self,
45 explicit: Option<&str>,
46 session_id: Option<&str>,
47 ) -> Option<String> {
48 if let Some(explicit) = explicit
49 .map(str::trim)
50 .filter(|value| !value.is_empty())
51 .map(ToString::to_string)
52 {
53 return Some(explicit);
54 }
55
56 if let Some(project_key) = self.memory_store.project_key_for_session(session_id) {
57 return Some(project_key);
58 }
59
60 self.session_for_context(session_id)
61 .await
62 .and_then(|session| session.metadata.get("workspace_path").cloned())
63 .map(std::path::PathBuf::from)
64 .map(|path| bamboo_memory::memory_store::project_key_from_path(&path))
65 }
66}
67
68#[async_trait]
69impl Tool for MemoryTool {
70 fn name(&self) -> &str {
71 "memory"
72 }
73
74 fn description(&self) -> &str {
75 "Unified memory management tool for Bamboo. Use session_* actions for session continuity notes, and query/get/write/merge/split/consolidate/purge/inspect/rebuild for durable project/global memory backed by canonical topic files and derived indexes."
76 }
77
78 fn parameters_schema(&self) -> serde_json::Value {
79 json!({
80 "type": "object",
81 "properties": {
82 "action": {
83 "type": "string",
84 "enum": [
85 "session_read",
86 "session_append",
87 "session_replace",
88 "session_clear",
89 "session_list_topics",
90 "query",
91 "get",
92 "find_duplicates",
93 "write",
94 "merge",
95 "split",
96 "consolidate",
97 "purge",
98 "inspect",
99 "rebuild",
100 "scan_blobs",
101 "scan_duplicates"
102 ]
103 },
104 "scope": {"type": "string", "enum": ["session", "project", "global"]},
105 "granularity": {
106 "type": "string",
107 "enum": ["day", "week", "month", "quarter", "year"],
108 "description": "Optional temporal granularity for `write`, orthogonal to scope: day (today's working context), week (sprint), month, quarter (direction), year (long-term goals). Omit if the memory has no time horizon. Coarser granularities are prefix-cache friendly and recalled ahead of finer ones at equal relevance."
109 },
110 "project_key": {"type": "string"},
111 "topic": {"type": "string"},
112 "id": {"type": "string"},
113 "query": {"type": "string"},
114 "type": {"type": "string", "enum": ["user", "feedback", "project", "reference"]},
115 "title": {"type": "string"},
116 "content": {"type": "string"},
117 "tags": {"type": "array", "items": {"type": "string"}},
118 "pieces": {"type": "array", "items": {"type": "object"}},
119 "ids": {"type": "array", "items": {"type": "string"}},
120 "min_score": {"type": "number"},
121 "filters": {"type": "object"},
122 "options": {"type": "object"},
123 "reason": {"type": "string"}
124 },
125 "required": ["action"]
126 })
127 }
128
129 fn classify(&self, args: &serde_json::Value) -> ToolClass {
130 let action = args
131 .get("action")
132 .and_then(|value| value.as_str())
133 .unwrap_or("")
134 .trim()
135 .to_ascii_lowercase();
136 match action.as_str() {
137 "session_read"
138 | "session_list_topics"
139 | "query"
140 | "get"
141 | "find_duplicates"
142 | "scan_blobs"
143 | "scan_duplicates"
144 | "inspect" => ToolClass::READONLY_PARALLEL,
145 _ => ToolClass::MUTATING_SERIAL,
146 }
147 }
148
149 async fn invoke(
150 &self,
151 args: serde_json::Value,
152 ctx: ToolCtx,
153 ) -> Result<ToolOutcome, ToolError> {
154 let session_id = ctx.session_id().ok_or_else(|| {
155 ToolError::Execution("memory requires a session_id in tool context".to_string())
156 })?;
157
158 let parsed: MemoryArgs = serde_json::from_value(args).map_err(|error| {
159 ToolError::InvalidArguments(format!("Invalid memory args: {error}"))
160 })?;
161
162 let result = match parsed {
163 MemoryArgs::SessionRead { topic, options } => {
164 let max_chars = options.and_then(|value| value.max_chars);
165 execute_session_memory_action(
166 &self.memory_store,
167 session_id,
168 SessionMemoryAction::Read,
169 topic.as_deref(),
170 None,
171 max_chars,
172 MEMORY_SESSION_ACTION_NAMES,
173 )
174 .await
175 }
176 MemoryArgs::SessionAppend { topic, content } => {
177 execute_session_memory_action(
178 &self.memory_store,
179 session_id,
180 SessionMemoryAction::Append,
181 topic.as_deref(),
182 Some(content.as_str()),
183 None,
184 MEMORY_SESSION_ACTION_NAMES,
185 )
186 .await
187 }
188 MemoryArgs::SessionReplace { topic, content } => {
189 execute_session_memory_action(
190 &self.memory_store,
191 session_id,
192 SessionMemoryAction::Replace,
193 topic.as_deref(),
194 Some(content.as_str()),
195 None,
196 MEMORY_SESSION_ACTION_NAMES,
197 )
198 .await
199 }
200 MemoryArgs::SessionClear { topic } => {
201 execute_session_memory_action(
202 &self.memory_store,
203 session_id,
204 SessionMemoryAction::Clear,
205 topic.as_deref(),
206 None,
207 None,
208 MEMORY_SESSION_ACTION_NAMES,
209 )
210 .await
211 }
212 MemoryArgs::SessionListTopics => {
213 execute_session_memory_action(
214 &self.memory_store,
215 session_id,
216 SessionMemoryAction::ListTopics,
217 None,
218 None,
219 None,
220 MEMORY_SESSION_ACTION_NAMES,
221 )
222 .await
223 }
224 MemoryArgs::Query {
225 scope,
226 query,
227 filters,
228 project_key,
229 options,
230 } => {
231 let scope = Self::parse_scope(Some(&scope))?;
232 if scope == MemoryScope::Session {
233 return Err(ToolError::InvalidArguments(
234 "query supports durable scopes only; use session_read/session_list_topics for session scope"
235 .to_string(),
236 ));
237 }
238 let project_key = self
239 .resolve_project_key(project_key.as_deref(), Some(session_id))
240 .await;
241 let options = MemoryQueryOptions {
242 limit: options
243 .as_ref()
244 .and_then(|value| value.limit)
245 .map(|value| value.min(MAX_QUERY_LIMIT)),
246 max_chars: options
247 .as_ref()
248 .and_then(|value| value.max_chars)
249 .map(|value| value.min(MAX_MAX_CHARS)),
250 cursor: options.as_ref().and_then(|value| value.cursor.clone()),
251 include_related: options
252 .as_ref()
253 .and_then(|value| value.include_related)
254 .unwrap_or(false),
255 };
256 let (filter_types, filter_statuses) = Self::parse_query_filters(filters.as_ref())?;
257 let result = self
258 .memory_store
259 .query_scope(
260 scope,
261 project_key.as_deref(),
262 query.as_deref(),
263 filter_types.as_ref(),
264 filter_statuses.as_ref(),
265 &options,
266 )
267 .await
268 .map_err(|error| {
269 ToolError::Execution(format!("Failed to query memory: {error}"))
270 })?;
271 Ok(ToolResult {
272 success: true,
273 result: json!({
274 "action": "query",
275 "success": true,
276 "data": result,
277 "summary": bamboo_memory::memory_store::summary_json(result.returned_count, result.matched_count),
278 "warnings": [],
279 }).to_string(),
280 display_preference: Some("json".to_string()),
281 images: Vec::new(),
282 })
283 }
284 MemoryArgs::Get {
285 id,
286 project_key,
287 options,
288 } => {
289 let project_key = self
290 .resolve_project_key(project_key.as_deref(), Some(session_id))
291 .await;
292 let max_chars = options
293 .and_then(|value| value.max_chars)
294 .unwrap_or(MAX_MAX_CHARS)
295 .min(MAX_MAX_CHARS);
296 let Some(mut doc) = self
297 .memory_store
298 .get_memory(id.trim(), project_key.as_deref())
299 .await
300 .map_err(|error| {
301 ToolError::Execution(format!("Failed to get memory: {error}"))
302 })?
303 else {
304 return Err(ToolError::Execution(format!(
305 "memory not found: {}",
306 id.trim()
307 )));
308 };
309 let (body, truncated) =
310 bamboo_memory::memory_store::truncate_chars(&doc.body, max_chars);
311 doc.body = body;
312 Ok(ToolResult {
313 success: true,
314 result: json!({
315 "action": "get",
316 "id": doc.frontmatter.id,
317 "memory": {
318 "frontmatter": doc.frontmatter,
319 "body": doc.body,
320 "path": doc.path,
321 "body_truncated": truncated,
322 }
323 })
324 .to_string(),
325 display_preference: Some("json".to_string()),
326 images: Vec::new(),
327 })
328 }
329 MemoryArgs::Write {
330 scope,
331 r#type,
332 title,
333 content,
334 tags,
335 project_key,
336 granularity,
337 options,
338 } => {
339 let scope = Self::parse_scope(Some(&scope))?;
340 if scope == MemoryScope::Session {
341 return Err(ToolError::InvalidArguments(
342 "write supports durable scopes only; use session_replace/session_append for session scope"
343 .to_string(),
344 ));
345 }
346 let granularity = Self::parse_granularity(granularity.as_deref())?;
347 let project_key = self
348 .resolve_project_key(project_key.as_deref(), Some(session_id))
349 .await;
350 let doc = self
351 .memory_store
352 .write_memory(
353 scope,
354 project_key.as_deref(),
355 Self::parse_type(&r#type)?,
356 &title,
357 &content,
358 &tags,
359 Some(session_id),
360 "main-model",
361 options
362 .and_then(|value| value.allow_merge_if_similar)
363 .unwrap_or(false),
364 granularity,
365 )
366 .await
367 .map_err(|error| {
368 ToolError::Execution(format!("Failed to write memory: {error}"))
369 })?;
370 Ok(ToolResult {
371 success: true,
372 result: json!({
373 "action": "write",
374 "memory": {
375 "id": doc.frontmatter.id,
376 "title": doc.frontmatter.title,
377 "type": doc.frontmatter.r#type,
378 "scope": doc.frontmatter.scope,
379 "status": doc.frontmatter.status,
380 "project_key": doc.frontmatter.project_key,
381 "path": doc.path,
382 }
383 })
384 .to_string(),
385 display_preference: Some("json".to_string()),
386 images: Vec::new(),
387 })
388 }
389 MemoryArgs::Merge {
390 id,
391 content,
392 tags,
393 project_key,
394 source_memory_ids,
395 mode,
396 reason,
397 } => {
398 let project_key = self
399 .resolve_project_key(project_key.as_deref(), Some(session_id))
400 .await;
401 let mode = Self::parse_merge_mode(mode.as_deref())?;
402 if matches!(mode.as_deref(), Some("contradict")) {
403 let Some(result) = self
404 .memory_store
405 .mark_memory_contradicted(
406 id.trim(),
407 project_key.as_deref(),
408 &source_memory_ids,
409 reason.as_deref().or(Some(content.trim())),
410 Some(session_id),
411 "main-model",
412 )
413 .await
414 .map_err(|error| {
415 ToolError::Execution(format!("Failed to contradict memory: {error}"))
416 })?
417 else {
418 return Err(ToolError::Execution(format!(
419 "memory not found: {}",
420 id.trim()
421 )));
422 };
423 Ok(ToolResult {
424 success: true,
425 result: json!({
426 "action": "merge",
427 "mode": "contradict",
428 "data": result,
429 })
430 .to_string(),
431 display_preference: Some("json".to_string()),
432 images: Vec::new(),
433 })
434 } else {
435 let Some(result) = self
436 .memory_store
437 .merge_memory(
438 id.trim(),
439 project_key.as_deref(),
440 &content,
441 &tags,
442 Some(session_id),
443 "main-model",
444 &source_memory_ids,
445 )
446 .await
447 .map_err(|error| {
448 ToolError::Execution(format!("Failed to merge memory: {error}"))
449 })?
450 else {
451 return Err(ToolError::Execution(format!(
452 "memory not found: {}",
453 id.trim()
454 )));
455 };
456 Ok(ToolResult {
457 success: true,
458 result: json!({
459 "action": "merge",
460 "mode": mode.unwrap_or_else(|| "merge".to_string()),
461 "data": result,
462 })
463 .to_string(),
464 display_preference: Some("json".to_string()),
465 images: Vec::new(),
466 })
467 }
468 }
469 MemoryArgs::FindDuplicates {
470 scope,
471 title,
472 content,
473 r#type,
474 tags,
475 project_key,
476 options,
477 } => {
478 let scope = Self::parse_scope(Some(&scope))?;
479 if scope == MemoryScope::Session {
480 return Err(ToolError::InvalidArguments(
481 "find_duplicates supports durable scopes only".to_string(),
482 ));
483 }
484 let r#type = match r#type.as_deref() {
485 Some(value) => Some(Self::parse_type(value)?),
486 None => None,
487 };
488 let project_key = self
489 .resolve_project_key(project_key.as_deref(), Some(session_id))
490 .await;
491 let limit = options
492 .and_then(|value| value.limit)
493 .unwrap_or(5)
494 .clamp(1, MAX_QUERY_LIMIT);
495 let candidates = self
496 .memory_store
497 .find_duplicate_candidates(
498 scope,
499 project_key.as_deref(),
500 r#type,
501 &title,
502 content.as_deref().unwrap_or(""),
503 &tags,
504 limit,
505 )
506 .await
507 .map_err(|error| {
508 ToolError::Execution(format!("Failed to find duplicates: {error}"))
509 })?;
510 Ok(ToolResult {
511 success: true,
512 result: json!({
513 "action": "find_duplicates",
514 "candidates": candidates,
515 })
516 .to_string(),
517 display_preference: Some("json".to_string()),
518 images: Vec::new(),
519 })
520 }
521 MemoryArgs::Split {
522 id,
523 project_key,
524 pieces,
525 } => {
526 if pieces.is_empty() {
527 return Err(ToolError::InvalidArguments(
528 "split requires at least one piece".to_string(),
529 ));
530 }
531 let project_key = self
532 .resolve_project_key(project_key.as_deref(), Some(session_id))
533 .await;
534 let mut split_pieces = Vec::with_capacity(pieces.len());
535 for piece in pieces {
536 let r#type = match piece.r#type.as_deref() {
537 Some(value) => Some(Self::parse_type(value)?),
538 None => None,
539 };
540 split_pieces.push(bamboo_memory::memory_store::MemorySplitPiece {
541 title: piece.title,
542 r#type,
543 content: piece.content,
544 tags: piece.tags,
545 });
546 }
547 let Some(result) = self
548 .memory_store
549 .split_memory(
550 id.trim(),
551 project_key.as_deref(),
552 &split_pieces,
553 Some(session_id),
554 "main-model",
555 )
556 .await
557 .map_err(|error| {
558 ToolError::Execution(format!("Failed to split memory: {error}"))
559 })?
560 else {
561 return Err(ToolError::Execution(format!(
562 "memory not found: {}",
563 id.trim()
564 )));
565 };
566 Ok(ToolResult {
567 success: true,
568 result: json!({
569 "action": "split",
570 "data": result,
571 })
572 .to_string(),
573 display_preference: Some("json".to_string()),
574 images: Vec::new(),
575 })
576 }
577 MemoryArgs::ScanBlobs {
578 scope,
579 project_key,
580 min_sections,
581 options,
582 } => {
583 let scope = Self::parse_scope(Some(&scope))?;
584 if scope == MemoryScope::Session {
585 return Err(ToolError::InvalidArguments(
586 "scan_blobs supports durable scopes only".to_string(),
587 ));
588 }
589 let project_key = self
590 .resolve_project_key(project_key.as_deref(), Some(session_id))
591 .await;
592 let min_sections = min_sections.unwrap_or(3);
593 let limit = options
594 .and_then(|value| value.limit)
595 .unwrap_or(20)
596 .clamp(1, 200);
597 let report = self
598 .memory_store
599 .scan_blob_candidates(scope, project_key.as_deref(), min_sections, limit)
600 .await
601 .map_err(|error| {
602 ToolError::Execution(format!("Failed to scan blobs: {error}"))
603 })?;
604 Ok(ToolResult {
605 success: true,
606 result: json!({
607 "action": "scan_blobs",
608 "report": report,
609 })
610 .to_string(),
611 display_preference: Some("json".to_string()),
612 images: Vec::new(),
613 })
614 }
615 MemoryArgs::ScanDuplicates {
616 scope,
617 project_key,
618 min_score,
619 options,
620 } => {
621 let scope = Self::parse_scope(Some(&scope))?;
622 if scope == MemoryScope::Session {
623 return Err(ToolError::InvalidArguments(
624 "scan_duplicates supports durable scopes only".to_string(),
625 ));
626 }
627 let project_key = self
628 .resolve_project_key(project_key.as_deref(), Some(session_id))
629 .await;
630 let min_score = min_score.unwrap_or(0.6);
631 let limit = options
632 .and_then(|value| value.limit)
633 .unwrap_or(20)
634 .clamp(1, 200);
635 let report = self
636 .memory_store
637 .scan_duplicate_clusters(scope, project_key.as_deref(), min_score, 5, limit)
638 .await
639 .map_err(|error| {
640 ToolError::Execution(format!("Failed to scan duplicates: {error}"))
641 })?;
642 Ok(ToolResult {
643 success: true,
644 result: json!({
645 "action": "scan_duplicates",
646 "report": report,
647 })
648 .to_string(),
649 display_preference: Some("json".to_string()),
650 images: Vec::new(),
651 })
652 }
653 MemoryArgs::Consolidate {
654 ids,
655 title,
656 content,
657 r#type,
658 tags,
659 project_key,
660 } => {
661 if ids.len() < 2 {
662 return Err(ToolError::InvalidArguments(
663 "consolidate requires at least two source memory ids".to_string(),
664 ));
665 }
666 let r#type = match r#type.as_deref() {
667 Some(value) => Some(Self::parse_type(value)?),
668 None => None,
669 };
670 let project_key = self
671 .resolve_project_key(project_key.as_deref(), Some(session_id))
672 .await;
673 let merged = bamboo_memory::memory_store::MemorySplitPiece {
674 title,
675 r#type,
676 content,
677 tags,
678 };
679 let ids: Vec<String> = ids.iter().map(|id| id.trim().to_string()).collect();
680 let Some(result) = self
681 .memory_store
682 .consolidate_memories(
683 &ids,
684 project_key.as_deref(),
685 &merged,
686 Some(session_id),
687 "main-model",
688 )
689 .await
690 .map_err(|error| {
691 ToolError::Execution(format!("Failed to consolidate memories: {error}"))
692 })?
693 else {
694 return Err(ToolError::Execution(
695 "one or more source memories not found".to_string(),
696 ));
697 };
698 Ok(ToolResult {
699 success: true,
700 result: json!({
701 "action": "consolidate",
702 "data": result,
703 })
704 .to_string(),
705 display_preference: Some("json".to_string()),
706 images: Vec::new(),
707 })
708 }
709 MemoryArgs::Purge {
710 id,
711 scope,
712 reason,
713 project_key,
714 filters,
715 mode,
716 } => {
717 let mode = match mode
718 .as_deref()
719 .map(str::trim)
720 .filter(|value| !value.is_empty())
721 {
722 Some(value) => Self::parse_status(value)?,
723 None => DurableMemoryStatus::Archived,
724 };
725 let project_key = self
726 .resolve_project_key(project_key.as_deref(), Some(session_id))
727 .await;
728
729 if let Some(id) = id
730 .as_deref()
731 .map(str::trim)
732 .filter(|value| !value.is_empty())
733 {
734 let Some(doc) = self
735 .memory_store
736 .archive_memory(id, project_key.as_deref(), mode, reason.as_deref())
737 .await
738 .map_err(|error| {
739 ToolError::Execution(format!("Failed to purge memory: {error}"))
740 })?
741 else {
742 return Err(ToolError::Execution(format!("memory not found: {}", id)));
743 };
744 Ok(ToolResult {
745 success: true,
746 result: json!({
747 "action": "purge",
748 "id": doc.frontmatter.id,
749 "status": doc.frontmatter.status,
750 })
751 .to_string(),
752 display_preference: Some("json".to_string()),
753 images: Vec::new(),
754 })
755 } else {
756 let scope = Self::parse_scope(scope.as_deref())?;
757 if scope == MemoryScope::Session {
758 return Err(ToolError::InvalidArguments(
759 "purge supports durable scopes only in v1".to_string(),
760 ));
761 }
762 let (filter_types, filter_statuses) =
763 Self::parse_query_filters(filters.as_ref())?;
764 let result = self
765 .memory_store
766 .purge_memories(
767 scope,
768 project_key.as_deref(),
769 filter_types.as_ref(),
770 filter_statuses.as_ref(),
771 mode,
772 reason.as_deref(),
773 )
774 .await
775 .map_err(|error| {
776 ToolError::Execution(format!("Failed to purge memory: {error}"))
777 })?;
778 Ok(ToolResult {
779 success: true,
780 result: json!({
781 "action": "purge",
782 "data": result,
783 })
784 .to_string(),
785 display_preference: Some("json".to_string()),
786 images: Vec::new(),
787 })
788 }
789 }
790 MemoryArgs::Inspect { scope, project_key } => {
791 let scope = Self::parse_scope(Some(&scope))?;
792 if scope == MemoryScope::Session {
793 return Err(ToolError::InvalidArguments(
794 "inspect supports durable scopes only in v1".to_string(),
795 ));
796 }
797 let project_key = self
798 .resolve_project_key(project_key.as_deref(), Some(session_id))
799 .await;
800 let result = self
801 .memory_store
802 .inspect_scope(scope, project_key.as_deref())
803 .await
804 .map_err(|error| {
805 ToolError::Execution(format!("Failed to inspect memory: {error}"))
806 })?;
807 Ok(ToolResult {
808 success: true,
809 result: json!({
810 "action": "inspect",
811 "data": result,
812 })
813 .to_string(),
814 display_preference: Some("json".to_string()),
815 images: Vec::new(),
816 })
817 }
818 MemoryArgs::Rebuild { scope, project_key } => {
819 let scope = Self::parse_scope(Some(&scope))?;
820 if scope == MemoryScope::Session {
821 return Err(ToolError::InvalidArguments(
822 "rebuild supports durable scopes only in v1".to_string(),
823 ));
824 }
825 let project_key = self
826 .resolve_project_key(project_key.as_deref(), Some(session_id))
827 .await;
828 self.memory_store
829 .rebuild_scope(scope, project_key.as_deref())
830 .await
831 .map_err(|error| {
832 ToolError::Execution(format!("Failed to rebuild memory artifacts: {error}"))
833 })?;
834 let inspect = self
835 .memory_store
836 .inspect_scope(scope, project_key.as_deref())
837 .await
838 .map_err(|error| {
839 ToolError::Execution(format!("Failed to inspect rebuilt memory: {error}"))
840 })?;
841 Ok(ToolResult {
842 success: true,
843 result: json!({
844 "action": "rebuild",
845 "scope": scope,
846 "project_key": project_key,
847 "data": inspect,
848 })
849 .to_string(),
850 display_preference: Some("json".to_string()),
851 images: Vec::new(),
852 })
853 }
854 }?;
855 Ok(ToolOutcome::Completed(result))
856 }
857}