1use dashmap::DashMap;
2use serde_json::json;
3use std::sync::{Arc, OnceLock};
4use tokio::sync::Mutex;
5
6use bamboo_agent_core::{ToolError, ToolResult};
7use bamboo_memory::memory::DEFAULT_TOPIC;
8use bamboo_memory::memory_store::{count_chars, truncate_chars, MemoryStore};
9
10pub const MAX_SESSION_NOTE_CHARS: usize = 12_000;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum SessionMemoryAction {
14 Read,
15 Append,
16 Replace,
17 Clear,
18 ListTopics,
19}
20
21#[derive(Debug, Clone, Copy)]
22pub struct SessionMemoryActionNames {
23 pub tool_name: &'static str,
24 pub read: &'static str,
25 pub append: &'static str,
26 pub replace: &'static str,
27 pub clear: &'static str,
28 pub list_topics: &'static str,
29}
30
31pub const SESSION_NOTE_ACTION_NAMES: SessionMemoryActionNames = SessionMemoryActionNames {
32 tool_name: "session_note",
33 read: "read",
34 append: "append",
35 replace: "replace",
36 clear: "clear",
37 list_topics: "list_topics",
38};
39
40pub const MEMORY_SESSION_ACTION_NAMES: SessionMemoryActionNames = SessionMemoryActionNames {
41 tool_name: "memory",
42 read: "session_read",
43 append: "session_append",
44 replace: "session_replace",
45 clear: "session_clear",
46 list_topics: "session_list_topics",
47};
48
49fn note_locks() -> &'static DashMap<String, Arc<Mutex<()>>> {
50 static NOTE_LOCKS: OnceLock<DashMap<String, Arc<Mutex<()>>>> = OnceLock::new();
51 NOTE_LOCKS.get_or_init(DashMap::new)
52}
53
54pub fn session_memory_lock(session_id: &str) -> Arc<Mutex<()>> {
55 note_locks()
56 .entry(session_id.to_string())
57 .or_insert_with(|| Arc::new(Mutex::new(())))
58 .clone()
59}
60
61pub fn parse_session_note_action(action: &str) -> Result<SessionMemoryAction, ToolError> {
62 match action.trim().to_ascii_lowercase().as_str() {
63 "read" => Ok(SessionMemoryAction::Read),
64 "append" => Ok(SessionMemoryAction::Append),
65 "replace" => Ok(SessionMemoryAction::Replace),
66 "clear" => Ok(SessionMemoryAction::Clear),
67 "list_topics" => Ok(SessionMemoryAction::ListTopics),
68 _ => Err(ToolError::InvalidArguments(
69 "action must be one of: read, append, replace, clear, list_topics. Rewrite the session_note call with valid JSON.".to_string(),
70 )),
71 }
72}
73
74pub async fn execute_session_memory_action(
75 memory: &MemoryStore,
76 session_id: &str,
77 action: SessionMemoryAction,
78 topic: Option<&str>,
79 content: Option<&str>,
80 max_chars: Option<usize>,
81 names: SessionMemoryActionNames,
82) -> Result<ToolResult, ToolError> {
83 let topic = topic
84 .map(str::trim)
85 .filter(|value| !value.is_empty())
86 .unwrap_or(DEFAULT_TOPIC);
87 let session_guard = session_memory_lock(session_id);
88 let _guard = session_guard.lock().await;
89
90 match action {
91 SessionMemoryAction::Read => {
92 let max_chars = max_chars
93 .unwrap_or(MAX_SESSION_NOTE_CHARS)
94 .clamp(1, MAX_SESSION_NOTE_CHARS);
95 let content = memory.read_session_topic(session_id, topic).await.map_err(|error| {
96 ToolError::Execution(format!(
97 "Failed to read note: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\",\"topic\":\"{}\"}}.",
98 names.tool_name, names.read, topic,
99 ))
100 })?;
101 let exists = content.is_some();
102 let body = content.unwrap_or_default();
103 let length_chars = count_chars(&body);
104 let (snippet, truncated) = truncate_chars(&body, max_chars);
105 Ok(ToolResult {
106 success: true,
107 result: json!({
108 "action": names.read,
109 "session_id": session_id,
110 "topic": topic,
111 "exists": exists,
112 "content": snippet,
113 "length_chars": length_chars,
114 "body_truncated": truncated,
115 "max_chars": max_chars,
116 })
117 .to_string(),
118 display_preference: Some("json".to_string()),
119 images: Vec::new(),
120 })
121 }
122 SessionMemoryAction::Clear => {
123 let deleted = memory.delete_session_topic(session_id, topic).await.map_err(|error| {
124 ToolError::Execution(format!(
125 "Failed to delete note: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\",\"topic\":\"{}\"}}.",
126 names.tool_name, names.clear, topic,
127 ))
128 })?;
129 Ok(ToolResult {
130 success: true,
131 result: json!({
132 "action": names.clear,
133 "session_id": session_id,
134 "topic": topic,
135 "deleted": deleted,
136 })
137 .to_string(),
138 display_preference: Some("json".to_string()),
139 images: Vec::new(),
140 })
141 }
142 SessionMemoryAction::ListTopics => {
143 let topics = memory.list_session_topics(session_id).await.map_err(|error| {
144 ToolError::Execution(format!(
145 "Failed to list topics: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\"}}.",
146 names.tool_name, names.list_topics,
147 ))
148 })?;
149 Ok(ToolResult {
150 success: true,
151 result: json!({
152 "action": names.list_topics,
153 "session_id": session_id,
154 "topics": topics,
155 "count": topics.len(),
156 })
157 .to_string(),
158 display_preference: Some("json".to_string()),
159 images: Vec::new(),
160 })
161 }
162 SessionMemoryAction::Replace | SessionMemoryAction::Append => {
163 let content = content
164 .map(str::trim)
165 .filter(|value| !value.is_empty())
166 .ok_or_else(|| {
167 ToolError::InvalidArguments(format!(
168 "content is required for action={}|{}. Rewrite the {} call with valid JSON and include non-empty content.",
169 names.append, names.replace, names.tool_name,
170 ))
171 })?;
172
173 if action == SessionMemoryAction::Replace {
174 let length_chars = count_chars(content);
175 if length_chars > MAX_SESSION_NOTE_CHARS {
176 return Err(ToolError::Execution(format!(
177 "session note too long (>{} chars). Compress it (rewrite more concisely) and call {} with action={} again.",
178 MAX_SESSION_NOTE_CHARS, names.tool_name, names.replace,
179 )));
180 }
181
182 let path = memory
183 .write_session_topic(session_id, topic, content)
184 .await
185 .map_err(|error| {
186 ToolError::Execution(format!(
187 "Failed to write note: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\",\"topic\":\"{}\",\"content\":\"...\"}}.",
188 names.tool_name, names.replace, topic,
189 ))
190 })?;
191
192 Ok(ToolResult {
193 success: true,
194 result: json!({
195 "action": names.replace,
196 "session_id": session_id,
197 "topic": topic,
198 "path": path,
199 "length_chars": length_chars,
200 "max_chars": MAX_SESSION_NOTE_CHARS,
201 })
202 .to_string(),
203 display_preference: Some("json".to_string()),
204 images: Vec::new(),
205 })
206 } else {
207 let existing = memory.read_session_topic(session_id, topic).await.map_err(|error| {
208 ToolError::Execution(format!(
209 "Failed to read note: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\",\"topic\":\"{}\",\"content\":\"...\"}}.",
210 names.tool_name, names.append, topic,
211 ))
212 })?;
213
214 let mut next = existing.unwrap_or_default();
215 if !next.is_empty() {
216 next.push_str("\n\n");
217 }
218 next.push_str(content);
219
220 let next_len = count_chars(&next);
221 if next_len > MAX_SESSION_NOTE_CHARS {
222 return Err(ToolError::Execution(format!(
223 "session note would exceed the limit ({}>{} chars). Compress the existing note (use {} action={} topic={}), then call {} action={} with a shorter version, then append again if needed.",
224 next_len,
225 MAX_SESSION_NOTE_CHARS,
226 names.tool_name,
227 names.read,
228 topic,
229 names.tool_name,
230 names.replace,
231 )));
232 }
233
234 let path = memory
235 .write_session_topic(session_id, topic, &next)
236 .await
237 .map_err(|error| {
238 ToolError::Execution(format!(
239 "Failed to write note: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\",\"topic\":\"{}\",\"content\":\"...\"}}.",
240 names.tool_name, names.append, topic,
241 ))
242 })?;
243
244 Ok(ToolResult {
245 success: true,
246 result: json!({
247 "action": names.append,
248 "session_id": session_id,
249 "topic": topic,
250 "path": path,
251 "length_chars": next_len,
252 "max_chars": MAX_SESSION_NOTE_CHARS,
253 })
254 .to_string(),
255 display_preference: Some("json".to_string()),
256 images: Vec::new(),
257 })
258 }
259 }
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266 use bamboo_memory::memory_store::MemoryStore;
267
268 #[tokio::test]
269 async fn execute_read_reports_length_and_truncation() {
270 let dir = tempfile::tempdir().expect("tempdir");
271 let store = MemoryStore::new(dir.path());
272 store
273 .write_session_topic("session-1", "default", &"x".repeat(32))
274 .await
275 .expect("write session topic");
276
277 let result = execute_session_memory_action(
278 &store,
279 "session-1",
280 SessionMemoryAction::Read,
281 Some("default"),
282 None,
283 Some(8),
284 SESSION_NOTE_ACTION_NAMES,
285 )
286 .await
287 .expect("read should succeed");
288
289 let value: serde_json::Value = serde_json::from_str(&result.result).expect("valid json");
290 assert_eq!(value["action"], "read");
291 assert_eq!(value["length_chars"], 32);
292 assert_eq!(value["body_truncated"], true);
293 assert_eq!(value["content"].as_str().unwrap().chars().count(), 8);
294 }
295
296 #[tokio::test]
297 async fn execute_append_enforces_shared_limit() {
298 let dir = tempfile::tempdir().expect("tempdir");
299 let store = MemoryStore::new(dir.path());
300 store
301 .write_session_topic(
302 "session-1",
303 "default",
304 &"x".repeat(MAX_SESSION_NOTE_CHARS - 1),
305 )
306 .await
307 .expect("write session topic");
308
309 let error = execute_session_memory_action(
310 &store,
311 "session-1",
312 SessionMemoryAction::Append,
313 Some("default"),
314 Some("y"),
315 None,
316 MEMORY_SESSION_ACTION_NAMES,
317 )
318 .await
319 .expect_err("append should fail");
320
321 let message = error.to_string();
322 assert!(message.contains("session note would exceed the limit"));
323 assert!(message.contains("action=session_read"));
324 assert!(message.contains("action=session_replace"));
325 }
326}