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 })
120 }
121 SessionMemoryAction::Clear => {
122 let deleted = memory.delete_session_topic(session_id, topic).await.map_err(|error| {
123 ToolError::Execution(format!(
124 "Failed to delete note: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\",\"topic\":\"{}\"}}.",
125 names.tool_name, names.clear, topic,
126 ))
127 })?;
128 Ok(ToolResult {
129 success: true,
130 result: json!({
131 "action": names.clear,
132 "session_id": session_id,
133 "topic": topic,
134 "deleted": deleted,
135 })
136 .to_string(),
137 display_preference: Some("json".to_string()),
138 })
139 }
140 SessionMemoryAction::ListTopics => {
141 let topics = memory.list_session_topics(session_id).await.map_err(|error| {
142 ToolError::Execution(format!(
143 "Failed to list topics: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\"}}.",
144 names.tool_name, names.list_topics,
145 ))
146 })?;
147 Ok(ToolResult {
148 success: true,
149 result: json!({
150 "action": names.list_topics,
151 "session_id": session_id,
152 "topics": topics,
153 "count": topics.len(),
154 })
155 .to_string(),
156 display_preference: Some("json".to_string()),
157 })
158 }
159 SessionMemoryAction::Replace | SessionMemoryAction::Append => {
160 let content = content
161 .map(str::trim)
162 .filter(|value| !value.is_empty())
163 .ok_or_else(|| {
164 ToolError::InvalidArguments(format!(
165 "content is required for action={}|{}. Rewrite the {} call with valid JSON and include non-empty content.",
166 names.append, names.replace, names.tool_name,
167 ))
168 })?;
169
170 if action == SessionMemoryAction::Replace {
171 let length_chars = count_chars(content);
172 if length_chars > MAX_SESSION_NOTE_CHARS {
173 return Err(ToolError::Execution(format!(
174 "session note too long (>{} chars). Compress it (rewrite more concisely) and call {} with action={} again.",
175 MAX_SESSION_NOTE_CHARS, names.tool_name, names.replace,
176 )));
177 }
178
179 let path = memory
180 .write_session_topic(session_id, topic, content)
181 .await
182 .map_err(|error| {
183 ToolError::Execution(format!(
184 "Failed to write note: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\",\"topic\":\"{}\",\"content\":\"...\"}}.",
185 names.tool_name, names.replace, topic,
186 ))
187 })?;
188
189 Ok(ToolResult {
190 success: true,
191 result: json!({
192 "action": names.replace,
193 "session_id": session_id,
194 "topic": topic,
195 "path": path,
196 "length_chars": length_chars,
197 "max_chars": MAX_SESSION_NOTE_CHARS,
198 })
199 .to_string(),
200 display_preference: Some("json".to_string()),
201 })
202 } else {
203 let existing = memory.read_session_topic(session_id, topic).await.map_err(|error| {
204 ToolError::Execution(format!(
205 "Failed to read note: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\",\"topic\":\"{}\",\"content\":\"...\"}}.",
206 names.tool_name, names.append, topic,
207 ))
208 })?;
209
210 let mut next = existing.unwrap_or_default();
211 if !next.is_empty() {
212 next.push_str("\n\n");
213 }
214 next.push_str(content);
215
216 let next_len = count_chars(&next);
217 if next_len > MAX_SESSION_NOTE_CHARS {
218 return Err(ToolError::Execution(format!(
219 "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.",
220 next_len,
221 MAX_SESSION_NOTE_CHARS,
222 names.tool_name,
223 names.read,
224 topic,
225 names.tool_name,
226 names.replace,
227 )));
228 }
229
230 let path = memory
231 .write_session_topic(session_id, topic, &next)
232 .await
233 .map_err(|error| {
234 ToolError::Execution(format!(
235 "Failed to write note: {error}. Rewrite and retry {} with valid JSON, e.g. {{\"action\":\"{}\",\"topic\":\"{}\",\"content\":\"...\"}}.",
236 names.tool_name, names.append, topic,
237 ))
238 })?;
239
240 Ok(ToolResult {
241 success: true,
242 result: json!({
243 "action": names.append,
244 "session_id": session_id,
245 "topic": topic,
246 "path": path,
247 "length_chars": next_len,
248 "max_chars": MAX_SESSION_NOTE_CHARS,
249 })
250 .to_string(),
251 display_preference: Some("json".to_string()),
252 })
253 }
254 }
255 }
256}
257
258#[cfg(test)]
259mod tests {
260 use super::*;
261 use bamboo_memory::memory_store::MemoryStore;
262
263 #[tokio::test]
264 async fn execute_read_reports_length_and_truncation() {
265 let dir = tempfile::tempdir().expect("tempdir");
266 let store = MemoryStore::new(dir.path());
267 store
268 .write_session_topic("session-1", "default", &"x".repeat(32))
269 .await
270 .expect("write session topic");
271
272 let result = execute_session_memory_action(
273 &store,
274 "session-1",
275 SessionMemoryAction::Read,
276 Some("default"),
277 None,
278 Some(8),
279 SESSION_NOTE_ACTION_NAMES,
280 )
281 .await
282 .expect("read should succeed");
283
284 let value: serde_json::Value = serde_json::from_str(&result.result).expect("valid json");
285 assert_eq!(value["action"], "read");
286 assert_eq!(value["length_chars"], 32);
287 assert_eq!(value["body_truncated"], true);
288 assert_eq!(value["content"].as_str().unwrap().chars().count(), 8);
289 }
290
291 #[tokio::test]
292 async fn execute_append_enforces_shared_limit() {
293 let dir = tempfile::tempdir().expect("tempdir");
294 let store = MemoryStore::new(dir.path());
295 store
296 .write_session_topic(
297 "session-1",
298 "default",
299 &"x".repeat(MAX_SESSION_NOTE_CHARS - 1),
300 )
301 .await
302 .expect("write session topic");
303
304 let error = execute_session_memory_action(
305 &store,
306 "session-1",
307 SessionMemoryAction::Append,
308 Some("default"),
309 Some("y"),
310 None,
311 MEMORY_SESSION_ACTION_NAMES,
312 )
313 .await
314 .expect_err("append should fail");
315
316 let message = error.to_string();
317 assert!(message.contains("session note would exceed the limit"));
318 assert!(message.contains("action=session_read"));
319 assert!(message.contains("action=session_replace"));
320 }
321}