bamboo_tools/tools/
memory_note.rs

1//! Persistent session-scoped note tool.
2//!
3//! Canonical tool name: `session_note`.
4//! The legacy `memory_note` name is accepted via executor-level alias routing.
5//!
6//! This tool lets the model store (and later retrieve) per-session notes that
7//! are loaded into the system prompt at the start of each round.
8//!
9//! Supports multiple **topics** per session so the model can track separate
10//! workstreams without clobbering each other.
11
12use async_trait::async_trait;
13use serde_json::json;
14
15use crate::tools::session_memory::{
16    execute_session_memory_action, parse_session_note_action, SESSION_NOTE_ACTION_NAMES,
17};
18use bamboo_agent_core::{Tool, ToolClass, ToolCtx, ToolError, ToolOutcome};
19use bamboo_memory::memory_store::MemoryStore;
20
21const TOOL_NAME: &str = "session_note";
22const TOOL_DESCRIPTION: &str = "Read or update the persistent session-scoped note (markdown). Use this for durable local context, user preferences, constraints, and compression-resistant reminders within the current session/workstream. Do not use it as the primary long-term knowledge base. Hard limit: 12000 characters; compress before append/replace if needed.";
23
24#[derive(Debug, Clone)]
25pub struct SessionNoteTool {
26    memory_store: MemoryStore,
27}
28
29impl SessionNoteTool {
30    pub fn new() -> Self {
31        Self {
32            memory_store: MemoryStore::with_defaults(),
33        }
34    }
35
36    #[cfg(test)]
37    fn with_memory_store(memory_store: MemoryStore) -> Self {
38        Self { memory_store }
39    }
40}
41
42impl Default for SessionNoteTool {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48/// Deprecated compatibility alias for older code paths. The canonical tool type
49/// and registered tool name is [`SessionNoteTool`] / `session_note`.
50pub type MemoryNoteTool = SessionNoteTool;
51
52#[async_trait]
53impl Tool for SessionNoteTool {
54    fn name(&self) -> &str {
55        TOOL_NAME
56    }
57
58    fn description(&self) -> &str {
59        TOOL_DESCRIPTION
60    }
61
62    fn classify(&self, args: &serde_json::Value) -> ToolClass {
63        let action = args
64            .get("action")
65            .and_then(|v| v.as_str())
66            .unwrap_or("")
67            .trim()
68            .to_ascii_lowercase();
69        match action.as_str() {
70            "read" | "list_topics" => ToolClass::READONLY_PARALLEL,
71            _ => ToolClass::MUTATING_SERIAL,
72        }
73    }
74
75    fn parameters_schema(&self) -> serde_json::Value {
76        json!({
77            "type": "object",
78            "properties": {
79                "action": {
80                    "type": "string",
81                    "description": "Operation to perform on the note.",
82                    "enum": ["read", "append", "replace", "clear", "list_topics"]
83                },
84                "content": {
85                    "type": "string",
86                    "description": "Note content to append/replace (markdown). Required for append/replace."
87                },
88                "topic": {
89                    "type": "string",
90                    "description": "Optional topic name (alphanumeric/dash/underscore, max 50 chars). Defaults to 'default'. Use separate topics for unrelated workstreams."
91                }
92            },
93            "required": ["action"]
94        })
95    }
96
97    async fn invoke(
98        &self,
99        args: serde_json::Value,
100        ctx: ToolCtx,
101    ) -> Result<ToolOutcome, ToolError> {
102        let Some(session_id) = ctx.session_id() else {
103            return Err(ToolError::Execution(
104                "missing session_id in tool context".to_string(),
105            ));
106        };
107
108        let action_raw = args.get("action").and_then(|v| v.as_str()).unwrap_or("");
109        let action = parse_session_note_action(action_raw)?;
110        let topic = args.get("topic").and_then(|v| v.as_str());
111        let content = args.get("content").and_then(|v| v.as_str());
112
113        let result = execute_session_memory_action(
114            &self.memory_store,
115            session_id,
116            action,
117            topic,
118            content,
119            None,
120            SESSION_NOTE_ACTION_NAMES,
121        )
122        .await?;
123        Ok(ToolOutcome::Completed(result))
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn session_note_schema_requires_action() {
133        let tool = SessionNoteTool::new();
134        let schema = tool.parameters_schema();
135        assert_eq!(schema["required"], json!(["action"]));
136        assert_eq!(tool.name(), TOOL_NAME);
137        assert_eq!(
138            schema["properties"]["action"]["enum"],
139            json!(["read", "append", "replace", "clear", "list_topics"])
140        );
141    }
142
143    #[test]
144    fn session_note_schema_has_topic_field() {
145        let tool = SessionNoteTool::new();
146        let schema = tool.parameters_schema();
147        assert!(schema["properties"]["topic"].is_object());
148        assert_eq!(schema["properties"]["topic"]["type"], "string");
149    }
150
151    #[tokio::test]
152    async fn session_note_requires_session_context() {
153        let tool = SessionNoteTool::new();
154        let result = tool
155            .invoke(json!({"action": "read"}), ToolCtx::none("tool_call"))
156            .await;
157
158        assert!(matches!(
159            result,
160            Err(ToolError::Execution(msg)) if msg.contains("session_id")
161        ));
162    }
163
164    #[tokio::test]
165    async fn session_note_validates_action_and_content_before_io() {
166        let tool = SessionNoteTool::new();
167
168        let unknown = tool
169            .invoke(
170                json!({"action": "unknown"}),
171                ToolCtx {
172                    session_id: Some(std::sync::Arc::from("session-1")),
173                    tool_call_id: std::sync::Arc::from("tool_call_unknown"),
174                    event_tx: None,
175                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
176                    bypass_permissions: false,
177                    can_async_resume: false,
178                    async_completion_sink: None,
179                    bash_completion_sink: None,
180                },
181            )
182            .await;
183        assert!(matches!(
184            unknown,
185            Err(ToolError::InvalidArguments(msg)) if msg.contains("action must be one of")
186        ));
187
188        let missing_content = tool
189            .invoke(
190                json!({"action": "replace"}),
191                ToolCtx {
192                    session_id: Some(std::sync::Arc::from("session-1")),
193                    tool_call_id: std::sync::Arc::from("tool_call_replace"),
194                    event_tx: None,
195                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
196                    bypass_permissions: false,
197                    can_async_resume: false,
198                    async_completion_sink: None,
199                    bash_completion_sink: None,
200                },
201            )
202            .await;
203        assert!(matches!(
204            missing_content,
205            Err(ToolError::InvalidArguments(msg)) if msg.contains("content is required")
206        ));
207    }
208
209    #[tokio::test]
210    async fn session_note_uses_memory_store_session_topics() {
211        let dir = tempfile::tempdir().unwrap();
212        let tool = SessionNoteTool::with_memory_store(MemoryStore::new(dir.path()));
213
214        let append = tool
215            .invoke(
216                json!({"action": "append", "topic": "backend", "content": "API finalized"}),
217                ToolCtx {
218                    session_id: Some(std::sync::Arc::from("session-1")),
219                    tool_call_id: std::sync::Arc::from("tool_call_append"),
220                    event_tx: None,
221                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
222                    bypass_permissions: false,
223                    can_async_resume: false,
224                    async_completion_sink: None,
225                    bash_completion_sink: None,
226                },
227            )
228            .await
229            .expect("append should succeed");
230        let ToolOutcome::Completed(append) = append else {
231            panic!("expected Completed")
232        };
233        let append_json: serde_json::Value = serde_json::from_str(&append.result).unwrap();
234        assert_eq!(append_json["action"], "append");
235        assert_eq!(append_json["length_chars"], "API finalized".chars().count());
236
237        let read = tool
238            .invoke(
239                json!({"action": "read", "topic": "backend"}),
240                ToolCtx {
241                    session_id: Some(std::sync::Arc::from("session-1")),
242                    tool_call_id: std::sync::Arc::from("tool_call_read"),
243                    event_tx: None,
244                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
245                    bypass_permissions: false,
246                    can_async_resume: false,
247                    async_completion_sink: None,
248                    bash_completion_sink: None,
249                },
250            )
251            .await
252            .expect("read should succeed");
253        let ToolOutcome::Completed(read) = read else {
254            panic!("expected Completed")
255        };
256        let read_json: serde_json::Value = serde_json::from_str(&read.result).unwrap();
257        assert_eq!(read_json["action"], "read");
258        assert_eq!(read_json["content"], "API finalized");
259        assert_eq!(read_json["length_chars"], "API finalized".chars().count());
260        assert_eq!(read_json["body_truncated"], false);
261
262        let list = tool
263            .invoke(
264                json!({"action": "list_topics"}),
265                ToolCtx {
266                    session_id: Some(std::sync::Arc::from("session-1")),
267                    tool_call_id: std::sync::Arc::from("tool_call_list"),
268                    event_tx: None,
269                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
270                    bypass_permissions: false,
271                    can_async_resume: false,
272                    async_completion_sink: None,
273                    bash_completion_sink: None,
274                },
275            )
276            .await
277            .expect("list should succeed");
278        let ToolOutcome::Completed(list) = list else {
279            panic!("expected Completed")
280        };
281        let list_json: serde_json::Value = serde_json::from_str(&list.result).unwrap();
282        assert_eq!(list_json["topics"][0], "backend");
283        assert_eq!(list_json["count"], 1);
284
285        let clear = tool
286            .invoke(
287                json!({"action": "clear", "topic": "backend"}),
288                ToolCtx {
289                    session_id: Some(std::sync::Arc::from("session-1")),
290                    tool_call_id: std::sync::Arc::from("tool_call_clear"),
291                    event_tx: None,
292                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
293                    bypass_permissions: false,
294                    can_async_resume: false,
295                    async_completion_sink: None,
296                    bash_completion_sink: None,
297                },
298            )
299            .await
300            .expect("clear should succeed");
301        let ToolOutcome::Completed(clear) = clear else {
302            panic!("expected Completed")
303        };
304        let clear_json: serde_json::Value = serde_json::from_str(&clear.result).unwrap();
305        assert_eq!(clear_json["action"], "clear");
306        assert_eq!(clear_json["deleted"], true);
307    }
308
309    #[tokio::test]
310    async fn session_note_read_reports_truncation_and_append_enforces_limit() {
311        let dir = tempfile::tempdir().unwrap();
312        let tool = SessionNoteTool::with_memory_store(MemoryStore::new(dir.path()));
313        let long_content = "x".repeat(32);
314
315        tool.invoke(
316            json!({"action": "replace", "topic": "default", "content": long_content}),
317            ToolCtx {
318                session_id: Some(std::sync::Arc::from("session-2")),
319                tool_call_id: std::sync::Arc::from("tool_call_replace_long"),
320                event_tx: None,
321                available_tool_schemas: std::sync::Arc::from(Vec::new()),
322                bypass_permissions: false,
323                can_async_resume: false,
324                async_completion_sink: None,
325                bash_completion_sink: None,
326            },
327        )
328        .await
329        .expect("replace should succeed");
330
331        let read = tool
332            .invoke(
333                json!({"action": "read", "topic": "default"}),
334                ToolCtx {
335                    session_id: Some(std::sync::Arc::from("session-2")),
336                    tool_call_id: std::sync::Arc::from("tool_call_read_long"),
337                    event_tx: None,
338                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
339                    bypass_permissions: false,
340                    can_async_resume: false,
341                    async_completion_sink: None,
342                    bash_completion_sink: None,
343                },
344            )
345            .await
346            .expect("read should succeed");
347        let ToolOutcome::Completed(read) = read else {
348            panic!("expected Completed")
349        };
350        let read_json: serde_json::Value = serde_json::from_str(&read.result).unwrap();
351        assert_eq!(read_json["length_chars"], 32);
352        assert_eq!(read_json["body_truncated"], false);
353
354        tool.invoke(
355            json!({"action": "replace", "topic": "limit", "content": "x".repeat(crate::tools::session_memory::MAX_SESSION_NOTE_CHARS - 1)}),
356            ToolCtx {
357                session_id: Some(std::sync::Arc::from("session-3")),
358                tool_call_id: std::sync::Arc::from("tool_call_replace_limit"),
359                event_tx: None,
360                available_tool_schemas: std::sync::Arc::from(Vec::new()),
361                bypass_permissions: false,
362                can_async_resume: false,
363                async_completion_sink: None,
364                bash_completion_sink: None,
365            },
366        )
367        .await
368        .expect("replace near limit should succeed");
369
370        let append_err = tool
371            .invoke(
372                json!({"action": "append", "topic": "limit", "content": "y"}),
373                ToolCtx {
374                    session_id: Some(std::sync::Arc::from("session-3")),
375                    tool_call_id: std::sync::Arc::from("tool_call_append_limit"),
376                    event_tx: None,
377                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
378                    bypass_permissions: false,
379                    can_async_resume: false,
380                    async_completion_sink: None,
381                    bash_completion_sink: None,
382                },
383            )
384            .await
385            .expect_err("append should exceed limit");
386        assert!(append_err
387            .to_string()
388            .contains("session note would exceed the limit"));
389    }
390}
bamboo_tools/tools/memory_note.rs

bamboo_tools/tools/
memory_note.rs
