use std::sync::Arc;
use async_trait::async_trait;
use serde_json::json;
use tokio::sync::RwLock;
use bamboo_agent_core::storage::Storage;
use bamboo_agent_core::tools::{Tool, ToolError, ToolExecutionContext, ToolResult};
use bamboo_agent_core::Session;
use bamboo_memory::memory_store::{
DurableMemoryStatus, MemoryQueryOptions, MemoryScope, MemoryStore, MAX_MAX_CHARS,
MAX_QUERY_LIMIT,
};
use bamboo_tools::tools::session_memory::{
execute_session_memory_action, SessionMemoryAction, MEMORY_SESSION_ACTION_NAMES,
};
mod args;
mod parsing;
#[cfg(test)]
mod tests;
use args::MemoryArgs;
#[derive(Clone)]
pub struct MemoryTool {
sessions: Arc<RwLock<std::collections::HashMap<String, Session>>>,
storage: Arc<dyn Storage>,
memory_store: MemoryStore,
}
impl MemoryTool {
pub fn new(
sessions: Arc<RwLock<std::collections::HashMap<String, Session>>>,
storage: Arc<dyn Storage>,
data_dir: impl Into<std::path::PathBuf>,
) -> Self {
Self {
sessions,
storage,
memory_store: MemoryStore::new(data_dir),
}
}
async fn session_for_context(&self, session_id: Option<&str>) -> Option<Session> {
let session_id = session_id?;
let in_memory = {
let sessions = self.sessions.read().await;
sessions.get(session_id).cloned()
};
match in_memory {
Some(session) => Some(session),
None => self.storage.load_session(session_id).await.ok().flatten(),
}
}
async fn resolve_project_key(
&self,
explicit: Option<&str>,
session_id: Option<&str>,
) -> Option<String> {
if let Some(explicit) = explicit
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
{
return Some(explicit);
}
if let Some(project_key) = self.memory_store.project_key_for_session(session_id) {
return Some(project_key);
}
self.session_for_context(session_id)
.await
.and_then(|session| session.metadata.get("workspace_path").cloned())
.map(std::path::PathBuf::from)
.map(|path| bamboo_memory::memory_store::project_key_from_path(&path))
}
}
#[async_trait]
impl Tool for MemoryTool {
fn name(&self) -> &str {
"memory"
}
fn description(&self) -> &str {
"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."
}
fn parameters_schema(&self) -> serde_json::Value {
json!({
"type": "object",
"properties": {
"action": {
"type": "string",
"enum": [
"session_read",
"session_append",
"session_replace",
"session_clear",
"session_list_topics",
"query",
"get",
"write",
"merge",
"purge",
"inspect",
"rebuild"
]
},
"scope": {"type": "string", "enum": ["session", "project", "global"]},
"project_key": {"type": "string"},
"topic": {"type": "string"},
"id": {"type": "string"},
"query": {"type": "string"},
"type": {"type": "string", "enum": ["user", "feedback", "project", "reference"]},
"title": {"type": "string"},
"content": {"type": "string"},
"tags": {"type": "array", "items": {"type": "string"}},
"filters": {"type": "object"},
"options": {"type": "object"},
"reason": {"type": "string"}
},
"required": ["action"]
})
}
fn call_mutability(&self, args: &serde_json::Value) -> bamboo_tools::ToolMutability {
let action = args
.get("action")
.and_then(|value| value.as_str())
.unwrap_or("")
.trim()
.to_ascii_lowercase();
match action.as_str() {
"session_read" | "session_list_topics" | "query" | "get" | "inspect" => {
bamboo_tools::ToolMutability::ReadOnly
}
_ => bamboo_tools::ToolMutability::Mutating,
}
}
fn call_concurrency_safe(&self, args: &serde_json::Value) -> bool {
matches!(
self.call_mutability(args),
bamboo_tools::ToolMutability::ReadOnly
)
}
async fn execute(&self, args: serde_json::Value) -> Result<ToolResult, ToolError> {
self.execute_with_context(args, ToolExecutionContext::none("tool_call"))
.await
}
async fn execute_with_context(
&self,
args: serde_json::Value,
ctx: ToolExecutionContext<'_>,
) -> Result<ToolResult, ToolError> {
let session_id = ctx.session_id.ok_or_else(|| {
ToolError::Execution("memory requires a session_id in tool context".to_string())
})?;
let parsed: MemoryArgs = serde_json::from_value(args).map_err(|error| {
ToolError::InvalidArguments(format!("Invalid memory args: {error}"))
})?;
match parsed {
MemoryArgs::SessionRead { topic, options } => {
let max_chars = options.and_then(|value| value.max_chars);
execute_session_memory_action(
&self.memory_store,
session_id,
SessionMemoryAction::Read,
topic.as_deref(),
None,
max_chars,
MEMORY_SESSION_ACTION_NAMES,
)
.await
}
MemoryArgs::SessionAppend { topic, content } => {
execute_session_memory_action(
&self.memory_store,
session_id,
SessionMemoryAction::Append,
topic.as_deref(),
Some(content.as_str()),
None,
MEMORY_SESSION_ACTION_NAMES,
)
.await
}
MemoryArgs::SessionReplace { topic, content } => {
execute_session_memory_action(
&self.memory_store,
session_id,
SessionMemoryAction::Replace,
topic.as_deref(),
Some(content.as_str()),
None,
MEMORY_SESSION_ACTION_NAMES,
)
.await
}
MemoryArgs::SessionClear { topic } => {
execute_session_memory_action(
&self.memory_store,
session_id,
SessionMemoryAction::Clear,
topic.as_deref(),
None,
None,
MEMORY_SESSION_ACTION_NAMES,
)
.await
}
MemoryArgs::SessionListTopics => {
execute_session_memory_action(
&self.memory_store,
session_id,
SessionMemoryAction::ListTopics,
None,
None,
None,
MEMORY_SESSION_ACTION_NAMES,
)
.await
}
MemoryArgs::Query {
scope,
query,
filters,
project_key,
options,
} => {
let scope = Self::parse_scope(Some(&scope))?;
if scope == MemoryScope::Session {
return Err(ToolError::InvalidArguments(
"query supports durable scopes only; use session_read/session_list_topics for session scope"
.to_string(),
));
}
let project_key = self
.resolve_project_key(project_key.as_deref(), Some(session_id))
.await;
let options = MemoryQueryOptions {
limit: options
.as_ref()
.and_then(|value| value.limit)
.map(|value| value.min(MAX_QUERY_LIMIT)),
max_chars: options
.as_ref()
.and_then(|value| value.max_chars)
.map(|value| value.min(MAX_MAX_CHARS)),
cursor: options.as_ref().and_then(|value| value.cursor.clone()),
include_related: options
.as_ref()
.and_then(|value| value.include_related)
.unwrap_or(false),
};
let (filter_types, filter_statuses) = Self::parse_query_filters(filters.as_ref())?;
let result = self
.memory_store
.query_scope(
scope,
project_key.as_deref(),
query.as_deref(),
filter_types.as_ref(),
filter_statuses.as_ref(),
&options,
)
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to query memory: {error}"))
})?;
Ok(ToolResult {
success: true,
result: json!({
"action": "query",
"success": true,
"data": result,
"summary": bamboo_memory::memory_store::summary_json(result.returned_count, result.matched_count),
"warnings": [],
}).to_string(),
display_preference: Some("json".to_string()),
})
}
MemoryArgs::Get {
id,
project_key,
options,
} => {
let project_key = self
.resolve_project_key(project_key.as_deref(), Some(session_id))
.await;
let max_chars = options
.and_then(|value| value.max_chars)
.unwrap_or(MAX_MAX_CHARS)
.min(MAX_MAX_CHARS);
let Some(mut doc) = self
.memory_store
.get_memory(id.trim(), project_key.as_deref())
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to get memory: {error}"))
})?
else {
return Err(ToolError::Execution(format!(
"memory not found: {}",
id.trim()
)));
};
let (body, truncated) =
bamboo_memory::memory_store::truncate_chars(&doc.body, max_chars);
doc.body = body;
Ok(ToolResult {
success: true,
result: json!({
"action": "get",
"id": doc.frontmatter.id,
"memory": {
"frontmatter": doc.frontmatter,
"body": doc.body,
"path": doc.path,
"body_truncated": truncated,
}
})
.to_string(),
display_preference: Some("json".to_string()),
})
}
MemoryArgs::Write {
scope,
r#type,
title,
content,
tags,
project_key,
options,
} => {
let scope = Self::parse_scope(Some(&scope))?;
if scope == MemoryScope::Session {
return Err(ToolError::InvalidArguments(
"write supports durable scopes only; use session_replace/session_append for session scope"
.to_string(),
));
}
let project_key = self
.resolve_project_key(project_key.as_deref(), Some(session_id))
.await;
let doc = self
.memory_store
.write_memory(
scope,
project_key.as_deref(),
Self::parse_type(&r#type)?,
&title,
&content,
&tags,
Some(session_id),
"main-model",
options
.and_then(|value| value.allow_merge_if_similar)
.unwrap_or(true),
)
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to write memory: {error}"))
})?;
Ok(ToolResult {
success: true,
result: json!({
"action": "write",
"memory": {
"id": doc.frontmatter.id,
"title": doc.frontmatter.title,
"type": doc.frontmatter.r#type,
"scope": doc.frontmatter.scope,
"status": doc.frontmatter.status,
"project_key": doc.frontmatter.project_key,
"path": doc.path,
}
})
.to_string(),
display_preference: Some("json".to_string()),
})
}
MemoryArgs::Merge {
id,
content,
tags,
project_key,
source_memory_ids,
mode,
reason,
} => {
let project_key = self
.resolve_project_key(project_key.as_deref(), Some(session_id))
.await;
let mode = Self::parse_merge_mode(mode.as_deref())?;
if matches!(mode.as_deref(), Some("contradict")) {
let Some(result) = self
.memory_store
.mark_memory_contradicted(
id.trim(),
project_key.as_deref(),
&source_memory_ids,
reason.as_deref().or(Some(content.trim())),
Some(session_id),
"main-model",
)
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to contradict memory: {error}"))
})?
else {
return Err(ToolError::Execution(format!(
"memory not found: {}",
id.trim()
)));
};
Ok(ToolResult {
success: true,
result: json!({
"action": "merge",
"mode": "contradict",
"data": result,
})
.to_string(),
display_preference: Some("json".to_string()),
})
} else {
let Some(result) = self
.memory_store
.merge_memory(
id.trim(),
project_key.as_deref(),
&content,
&tags,
Some(session_id),
"main-model",
&source_memory_ids,
)
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to merge memory: {error}"))
})?
else {
return Err(ToolError::Execution(format!(
"memory not found: {}",
id.trim()
)));
};
Ok(ToolResult {
success: true,
result: json!({
"action": "merge",
"mode": mode.unwrap_or_else(|| "merge".to_string()),
"data": result,
})
.to_string(),
display_preference: Some("json".to_string()),
})
}
}
MemoryArgs::Purge {
id,
scope,
reason,
project_key,
filters,
mode,
} => {
let mode = match mode
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
Some(value) => Self::parse_status(value)?,
None => DurableMemoryStatus::Archived,
};
let project_key = self
.resolve_project_key(project_key.as_deref(), Some(session_id))
.await;
if let Some(id) = id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
let Some(doc) = self
.memory_store
.archive_memory(id, project_key.as_deref(), mode, reason.as_deref())
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to purge memory: {error}"))
})?
else {
return Err(ToolError::Execution(format!("memory not found: {}", id)));
};
Ok(ToolResult {
success: true,
result: json!({
"action": "purge",
"id": doc.frontmatter.id,
"status": doc.frontmatter.status,
})
.to_string(),
display_preference: Some("json".to_string()),
})
} else {
let scope = Self::parse_scope(scope.as_deref())?;
if scope == MemoryScope::Session {
return Err(ToolError::InvalidArguments(
"purge supports durable scopes only in v1".to_string(),
));
}
let (filter_types, filter_statuses) =
Self::parse_query_filters(filters.as_ref())?;
let result = self
.memory_store
.purge_memories(
scope,
project_key.as_deref(),
filter_types.as_ref(),
filter_statuses.as_ref(),
mode,
reason.as_deref(),
)
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to purge memory: {error}"))
})?;
Ok(ToolResult {
success: true,
result: json!({
"action": "purge",
"data": result,
})
.to_string(),
display_preference: Some("json".to_string()),
})
}
}
MemoryArgs::Inspect { scope, project_key } => {
let scope = Self::parse_scope(Some(&scope))?;
if scope == MemoryScope::Session {
return Err(ToolError::InvalidArguments(
"inspect supports durable scopes only in v1".to_string(),
));
}
let project_key = self
.resolve_project_key(project_key.as_deref(), Some(session_id))
.await;
let result = self
.memory_store
.inspect_scope(scope, project_key.as_deref())
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to inspect memory: {error}"))
})?;
Ok(ToolResult {
success: true,
result: json!({
"action": "inspect",
"data": result,
})
.to_string(),
display_preference: Some("json".to_string()),
})
}
MemoryArgs::Rebuild { scope, project_key } => {
let scope = Self::parse_scope(Some(&scope))?;
if scope == MemoryScope::Session {
return Err(ToolError::InvalidArguments(
"rebuild supports durable scopes only in v1".to_string(),
));
}
let project_key = self
.resolve_project_key(project_key.as_deref(), Some(session_id))
.await;
self.memory_store
.rebuild_scope(scope, project_key.as_deref())
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to rebuild memory artifacts: {error}"))
})?;
let inspect = self
.memory_store
.inspect_scope(scope, project_key.as_deref())
.await
.map_err(|error| {
ToolError::Execution(format!("Failed to inspect rebuilt memory: {error}"))
})?;
Ok(ToolResult {
success: true,
result: json!({
"action": "rebuild",
"scope": scope,
"project_key": project_key,
"data": inspect,
})
.to_string(),
display_preference: Some("json".to_string()),
})
}
}
}
}