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