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/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 "write",
108 "merge",
109 "purge",
110 "inspect",
111 "rebuild"
112 ]
113 },
114 "scope": {"type": "string", "enum": ["session", "project", "global"]},
115 "project_key": {"type": "string"},
116 "topic": {"type": "string"},
117 "id": {"type": "string"},
118 "query": {"type": "string"},
119 "type": {"type": "string", "enum": ["user", "feedback", "project", "reference"]},
120 "title": {"type": "string"},
121 "content": {"type": "string"},
122 "tags": {"type": "array", "items": {"type": "string"}},
123 "filters": {"type": "object"},
124 "options": {"type": "object"},
125 "reason": {"type": "string"}
126 },
127 "required": ["action"]
128 })
129 }
130
131 fn call_mutability(&self, args: &serde_json::Value) -> bamboo_tools::ToolMutability {
132 let action = args
133 .get("action")
134 .and_then(|value| value.as_str())
135 .unwrap_or("")
136 .trim()
137 .to_ascii_lowercase();
138 match action.as_str() {
139 "session_read" | "session_list_topics" | "query" | "get" | "inspect" => {
140 bamboo_tools::ToolMutability::ReadOnly
141 }
142 _ => bamboo_tools::ToolMutability::Mutating,
143 }
144 }
145
146 fn call_concurrency_safe(&self, args: &serde_json::Value) -> bool {
147 matches!(
148 self.call_mutability(args),
149 bamboo_tools::ToolMutability::ReadOnly
150 )
151 }
152
153 async fn execute(&self, args: serde_json::Value) -> Result<ToolResult, ToolError> {
154 self.execute_with_context(args, ToolExecutionContext::none("tool_call"))
155 .await
156 }
157
158 async fn execute_with_context(
159 &self,
160 args: serde_json::Value,
161 ctx: ToolExecutionContext<'_>,
162 ) -> Result<ToolResult, ToolError> {
163 let session_id = ctx.session_id.ok_or_else(|| {
164 ToolError::Execution("memory requires a session_id in tool context".to_string())
165 })?;
166
167 let parsed: MemoryArgs = serde_json::from_value(args).map_err(|error| {
168 ToolError::InvalidArguments(format!("Invalid memory args: {error}"))
169 })?;
170
171 match parsed {
172 MemoryArgs::SessionRead { topic, options } => {
173 let max_chars = options.and_then(|value| value.max_chars);
174 execute_session_memory_action(
175 &self.memory_store,
176 session_id,
177 SessionMemoryAction::Read,
178 topic.as_deref(),
179 None,
180 max_chars,
181 MEMORY_SESSION_ACTION_NAMES,
182 )
183 .await
184 }
185 MemoryArgs::SessionAppend { topic, content } => {
186 execute_session_memory_action(
187 &self.memory_store,
188 session_id,
189 SessionMemoryAction::Append,
190 topic.as_deref(),
191 Some(content.as_str()),
192 None,
193 MEMORY_SESSION_ACTION_NAMES,
194 )
195 .await
196 }
197 MemoryArgs::SessionReplace { topic, content } => {
198 execute_session_memory_action(
199 &self.memory_store,
200 session_id,
201 SessionMemoryAction::Replace,
202 topic.as_deref(),
203 Some(content.as_str()),
204 None,
205 MEMORY_SESSION_ACTION_NAMES,
206 )
207 .await
208 }
209 MemoryArgs::SessionClear { topic } => {
210 execute_session_memory_action(
211 &self.memory_store,
212 session_id,
213 SessionMemoryAction::Clear,
214 topic.as_deref(),
215 None,
216 None,
217 MEMORY_SESSION_ACTION_NAMES,
218 )
219 .await
220 }
221 MemoryArgs::SessionListTopics => {
222 execute_session_memory_action(
223 &self.memory_store,
224 session_id,
225 SessionMemoryAction::ListTopics,
226 None,
227 None,
228 None,
229 MEMORY_SESSION_ACTION_NAMES,
230 )
231 .await
232 }
233 MemoryArgs::Query {
234 scope,
235 query,
236 filters,
237 project_key,
238 options,
239 } => {
240 let scope = Self::parse_scope(Some(&scope))?;
241 if scope == MemoryScope::Session {
242 return Err(ToolError::InvalidArguments(
243 "query supports durable scopes only; use session_read/session_list_topics for session scope"
244 .to_string(),
245 ));
246 }
247 let project_key = self
248 .resolve_project_key(project_key.as_deref(), Some(session_id))
249 .await;
250 let options = MemoryQueryOptions {
251 limit: options
252 .as_ref()
253 .and_then(|value| value.limit)
254 .map(|value| value.min(MAX_QUERY_LIMIT)),
255 max_chars: options
256 .as_ref()
257 .and_then(|value| value.max_chars)
258 .map(|value| value.min(MAX_MAX_CHARS)),
259 cursor: options.as_ref().and_then(|value| value.cursor.clone()),
260 include_related: options
261 .as_ref()
262 .and_then(|value| value.include_related)
263 .unwrap_or(false),
264 };
265 let (filter_types, filter_statuses) = Self::parse_query_filters(filters.as_ref())?;
266 let result = self
267 .memory_store
268 .query_scope(
269 scope,
270 project_key.as_deref(),
271 query.as_deref(),
272 filter_types.as_ref(),
273 filter_statuses.as_ref(),
274 &options,
275 )
276 .await
277 .map_err(|error| {
278 ToolError::Execution(format!("Failed to query memory: {error}"))
279 })?;
280 Ok(ToolResult {
281 success: true,
282 result: json!({
283 "action": "query",
284 "success": true,
285 "data": result,
286 "summary": bamboo_memory::memory_store::summary_json(result.returned_count, result.matched_count),
287 "warnings": [],
288 }).to_string(),
289 display_preference: Some("json".to_string()),
290 })
291 }
292 MemoryArgs::Get {
293 id,
294 project_key,
295 options,
296 } => {
297 let project_key = self
298 .resolve_project_key(project_key.as_deref(), Some(session_id))
299 .await;
300 let max_chars = options
301 .and_then(|value| value.max_chars)
302 .unwrap_or(MAX_MAX_CHARS)
303 .min(MAX_MAX_CHARS);
304 let Some(mut doc) = self
305 .memory_store
306 .get_memory(id.trim(), project_key.as_deref())
307 .await
308 .map_err(|error| {
309 ToolError::Execution(format!("Failed to get memory: {error}"))
310 })?
311 else {
312 return Err(ToolError::Execution(format!(
313 "memory not found: {}",
314 id.trim()
315 )));
316 };
317 let (body, truncated) =
318 bamboo_memory::memory_store::truncate_chars(&doc.body, max_chars);
319 doc.body = body;
320 Ok(ToolResult {
321 success: true,
322 result: json!({
323 "action": "get",
324 "id": doc.frontmatter.id,
325 "memory": {
326 "frontmatter": doc.frontmatter,
327 "body": doc.body,
328 "path": doc.path,
329 "body_truncated": truncated,
330 }
331 })
332 .to_string(),
333 display_preference: Some("json".to_string()),
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(true),
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 })
391 }
392 MemoryArgs::Merge {
393 id,
394 content,
395 tags,
396 project_key,
397 source_memory_ids,
398 mode,
399 reason,
400 } => {
401 let project_key = self
402 .resolve_project_key(project_key.as_deref(), Some(session_id))
403 .await;
404 let mode = Self::parse_merge_mode(mode.as_deref())?;
405 if matches!(mode.as_deref(), Some("contradict")) {
406 let Some(result) = self
407 .memory_store
408 .mark_memory_contradicted(
409 id.trim(),
410 project_key.as_deref(),
411 &source_memory_ids,
412 reason.as_deref().or(Some(content.trim())),
413 Some(session_id),
414 "main-model",
415 )
416 .await
417 .map_err(|error| {
418 ToolError::Execution(format!("Failed to contradict memory: {error}"))
419 })?
420 else {
421 return Err(ToolError::Execution(format!(
422 "memory not found: {}",
423 id.trim()
424 )));
425 };
426 Ok(ToolResult {
427 success: true,
428 result: json!({
429 "action": "merge",
430 "mode": "contradict",
431 "data": result,
432 })
433 .to_string(),
434 display_preference: Some("json".to_string()),
435 })
436 } else {
437 let Some(result) = self
438 .memory_store
439 .merge_memory(
440 id.trim(),
441 project_key.as_deref(),
442 &content,
443 &tags,
444 Some(session_id),
445 "main-model",
446 &source_memory_ids,
447 )
448 .await
449 .map_err(|error| {
450 ToolError::Execution(format!("Failed to merge memory: {error}"))
451 })?
452 else {
453 return Err(ToolError::Execution(format!(
454 "memory not found: {}",
455 id.trim()
456 )));
457 };
458 Ok(ToolResult {
459 success: true,
460 result: json!({
461 "action": "merge",
462 "mode": mode.unwrap_or_else(|| "merge".to_string()),
463 "data": result,
464 })
465 .to_string(),
466 display_preference: Some("json".to_string()),
467 })
468 }
469 }
470 MemoryArgs::Purge {
471 id,
472 scope,
473 reason,
474 project_key,
475 filters,
476 mode,
477 } => {
478 let mode = match mode
479 .as_deref()
480 .map(str::trim)
481 .filter(|value| !value.is_empty())
482 {
483 Some(value) => Self::parse_status(value)?,
484 None => DurableMemoryStatus::Archived,
485 };
486 let project_key = self
487 .resolve_project_key(project_key.as_deref(), Some(session_id))
488 .await;
489
490 if let Some(id) = id
491 .as_deref()
492 .map(str::trim)
493 .filter(|value| !value.is_empty())
494 {
495 let Some(doc) = self
496 .memory_store
497 .archive_memory(id, project_key.as_deref(), mode, reason.as_deref())
498 .await
499 .map_err(|error| {
500 ToolError::Execution(format!("Failed to purge memory: {error}"))
501 })?
502 else {
503 return Err(ToolError::Execution(format!("memory not found: {}", id)));
504 };
505 Ok(ToolResult {
506 success: true,
507 result: json!({
508 "action": "purge",
509 "id": doc.frontmatter.id,
510 "status": doc.frontmatter.status,
511 })
512 .to_string(),
513 display_preference: Some("json".to_string()),
514 })
515 } else {
516 let scope = Self::parse_scope(scope.as_deref())?;
517 if scope == MemoryScope::Session {
518 return Err(ToolError::InvalidArguments(
519 "purge supports durable scopes only in v1".to_string(),
520 ));
521 }
522 let (filter_types, filter_statuses) =
523 Self::parse_query_filters(filters.as_ref())?;
524 let result = self
525 .memory_store
526 .purge_memories(
527 scope,
528 project_key.as_deref(),
529 filter_types.as_ref(),
530 filter_statuses.as_ref(),
531 mode,
532 reason.as_deref(),
533 )
534 .await
535 .map_err(|error| {
536 ToolError::Execution(format!("Failed to purge memory: {error}"))
537 })?;
538 Ok(ToolResult {
539 success: true,
540 result: json!({
541 "action": "purge",
542 "data": result,
543 })
544 .to_string(),
545 display_preference: Some("json".to_string()),
546 })
547 }
548 }
549 MemoryArgs::Inspect { scope, project_key } => {
550 let scope = Self::parse_scope(Some(&scope))?;
551 if scope == MemoryScope::Session {
552 return Err(ToolError::InvalidArguments(
553 "inspect supports durable scopes only in v1".to_string(),
554 ));
555 }
556 let project_key = self
557 .resolve_project_key(project_key.as_deref(), Some(session_id))
558 .await;
559 let result = self
560 .memory_store
561 .inspect_scope(scope, project_key.as_deref())
562 .await
563 .map_err(|error| {
564 ToolError::Execution(format!("Failed to inspect memory: {error}"))
565 })?;
566 Ok(ToolResult {
567 success: true,
568 result: json!({
569 "action": "inspect",
570 "data": result,
571 })
572 .to_string(),
573 display_preference: Some("json".to_string()),
574 })
575 }
576 MemoryArgs::Rebuild { scope, project_key } => {
577 let scope = Self::parse_scope(Some(&scope))?;
578 if scope == MemoryScope::Session {
579 return Err(ToolError::InvalidArguments(
580 "rebuild supports durable scopes only in v1".to_string(),
581 ));
582 }
583 let project_key = self
584 .resolve_project_key(project_key.as_deref(), Some(session_id))
585 .await;
586 self.memory_store
587 .rebuild_scope(scope, project_key.as_deref())
588 .await
589 .map_err(|error| {
590 ToolError::Execution(format!("Failed to rebuild memory artifacts: {error}"))
591 })?;
592 let inspect = self
593 .memory_store
594 .inspect_scope(scope, project_key.as_deref())
595 .await
596 .map_err(|error| {
597 ToolError::Execution(format!("Failed to inspect rebuilt memory: {error}"))
598 })?;
599 Ok(ToolResult {
600 success: true,
601 result: json!({
602 "action": "rebuild",
603 "scope": scope,
604 "project_key": project_key,
605 "data": inspect,
606 })
607 .to_string(),
608 display_preference: Some("json".to_string()),
609 })
610 }
611 }
612 }
613}