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