Skip to main content

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    pub fn with_memory_store(memory_store: MemoryStore) -> Self {
37        Self { memory_store }
38    }
39}
40
41impl Default for SessionNoteTool {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47/// Deprecated compatibility alias for older code paths. The canonical tool type
48/// and registered tool name is [`SessionNoteTool`] / `session_note`.
49pub type MemoryNoteTool = SessionNoteTool;
50
51#[async_trait]
52impl Tool for SessionNoteTool {
53    fn name(&self) -> &str {
54        TOOL_NAME
55    }
56
57    fn description(&self) -> &str {
58        TOOL_DESCRIPTION
59    }
60
61    fn classify(&self, args: &serde_json::Value) -> ToolClass {
62        let action = args
63            .get("action")
64            .and_then(|v| v.as_str())
65            .unwrap_or("")
66            .trim()
67            .to_ascii_lowercase();
68        match action.as_str() {
69            "read" | "list_topics" => ToolClass::READONLY_PARALLEL,
70            _ => ToolClass::MUTATING_SERIAL,
71        }
72    }
73
74    fn parameters_schema(&self) -> serde_json::Value {
75        json!({
76            "type": "object",
77            "properties": {
78                "action": {
79                    "type": "string",
80                    "description": "Operation to perform on the note.",
81                    "enum": ["read", "append", "replace", "clear", "list_topics"]
82                },
83                "content": {
84                    "type": "string",
85                    "description": "Note content to append/replace (markdown). Required for append/replace."
86                },
87                "topic": {
88                    "type": "string",
89                    "description": "Optional topic name (alphanumeric/dash/underscore, max 50 chars). Defaults to 'default'. Use separate topics for unrelated workstreams."
90                }
91            },
92            "required": ["action"]
93        })
94    }
95
96    async fn invoke(
97        &self,
98        args: serde_json::Value,
99        ctx: ToolCtx,
100    ) -> Result<ToolOutcome, ToolError> {
101        let Some(session_id) = ctx.session_id() else {
102            return Err(ToolError::Execution(
103                "missing session_id in tool context".to_string(),
104            ));
105        };
106
107        let action_raw = args.get("action").and_then(|v| v.as_str()).unwrap_or("");
108        let action = parse_session_note_action(action_raw)?;
109        let topic = args.get("topic").and_then(|v| v.as_str());
110        let content = args.get("content").and_then(|v| v.as_str());
111
112        let result = execute_session_memory_action(
113            &self.memory_store,
114            session_id,
115            action,
116            topic,
117            content,
118            None,
119            SESSION_NOTE_ACTION_NAMES,
120        )
121        .await?;
122        Ok(ToolOutcome::Completed(result))
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn session_note_schema_requires_action() {
132        let tool = SessionNoteTool::new();
133        let schema = tool.parameters_schema();
134        assert_eq!(schema["required"], json!(["action"]));
135        assert_eq!(tool.name(), TOOL_NAME);
136        assert_eq!(
137            schema["properties"]["action"]["enum"],
138            json!(["read", "append", "replace", "clear", "list_topics"])
139        );
140    }
141
142    #[test]
143    fn session_note_schema_has_topic_field() {
144        let tool = SessionNoteTool::new();
145        let schema = tool.parameters_schema();
146        assert!(schema["properties"]["topic"].is_object());
147        assert_eq!(schema["properties"]["topic"]["type"], "string");
148    }
149
150    #[tokio::test]
151    async fn session_note_requires_session_context() {
152        let tool = SessionNoteTool::new();
153        let result = tool
154            .invoke(json!({"action": "read"}), ToolCtx::none("tool_call"))
155            .await;
156
157        assert!(matches!(
158            result,
159            Err(ToolError::Execution(msg)) if msg.contains("session_id")
160        ));
161    }
162
163    #[tokio::test]
164    async fn session_note_validates_action_and_content_before_io() {
165        let tool = SessionNoteTool::new();
166
167        let unknown = tool
168            .invoke(
169                json!({"action": "unknown"}),
170                ToolCtx {
171                    executing_supervisor: None,
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                    auto_approve_permissions: false,
178                    plan_read_only: false,
179                    can_async_resume: false,
180                    async_completion_sink: None,
181                    bash_completion_sink: None,
182                },
183            )
184            .await;
185        assert!(matches!(
186            unknown,
187            Err(ToolError::InvalidArguments(msg)) if msg.contains("action must be one of")
188        ));
189
190        let missing_content = tool
191            .invoke(
192                json!({"action": "replace"}),
193                ToolCtx {
194                    executing_supervisor: None,
195                    session_id: Some(std::sync::Arc::from("session-1")),
196                    tool_call_id: std::sync::Arc::from("tool_call_replace"),
197                    event_tx: None,
198                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
199                    bypass_permissions: false,
200                    auto_approve_permissions: false,
201                    plan_read_only: false,
202                    can_async_resume: false,
203                    async_completion_sink: None,
204                    bash_completion_sink: None,
205                },
206            )
207            .await;
208        assert!(matches!(
209            missing_content,
210            Err(ToolError::InvalidArguments(msg)) if msg.contains("content is required")
211        ));
212    }
213
214    #[tokio::test]
215    async fn session_note_uses_memory_store_session_topics() {
216        let dir = tempfile::tempdir().unwrap();
217        let tool = SessionNoteTool::with_memory_store(MemoryStore::new(dir.path()));
218
219        let append = tool
220            .invoke(
221                json!({"action": "append", "topic": "backend", "content": "API finalized"}),
222                ToolCtx {
223                    executing_supervisor: None,
224                    session_id: Some(std::sync::Arc::from("session-1")),
225                    tool_call_id: std::sync::Arc::from("tool_call_append"),
226                    event_tx: None,
227                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
228                    bypass_permissions: false,
229                    auto_approve_permissions: false,
230                    plan_read_only: false,
231                    can_async_resume: false,
232                    async_completion_sink: None,
233                    bash_completion_sink: None,
234                },
235            )
236            .await
237            .expect("append should succeed");
238        let ToolOutcome::Completed(append) = append else {
239            panic!("expected Completed")
240        };
241        let append_json: serde_json::Value = serde_json::from_str(&append.result).unwrap();
242        assert_eq!(append_json["action"], "append");
243        assert_eq!(append_json["length_chars"], "API finalized".chars().count());
244
245        let read = tool
246            .invoke(
247                json!({"action": "read", "topic": "backend"}),
248                ToolCtx {
249                    executing_supervisor: None,
250                    session_id: Some(std::sync::Arc::from("session-1")),
251                    tool_call_id: std::sync::Arc::from("tool_call_read"),
252                    event_tx: None,
253                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
254                    bypass_permissions: false,
255                    auto_approve_permissions: false,
256                    plan_read_only: false,
257                    can_async_resume: false,
258                    async_completion_sink: None,
259                    bash_completion_sink: None,
260                },
261            )
262            .await
263            .expect("read should succeed");
264        let ToolOutcome::Completed(read) = read else {
265            panic!("expected Completed")
266        };
267        let read_json: serde_json::Value = serde_json::from_str(&read.result).unwrap();
268        assert_eq!(read_json["action"], "read");
269        assert_eq!(read_json["content"], "API finalized");
270        assert_eq!(read_json["length_chars"], "API finalized".chars().count());
271        assert_eq!(read_json["body_truncated"], false);
272
273        let list = tool
274            .invoke(
275                json!({"action": "list_topics"}),
276                ToolCtx {
277                    executing_supervisor: None,
278                    session_id: Some(std::sync::Arc::from("session-1")),
279                    tool_call_id: std::sync::Arc::from("tool_call_list"),
280                    event_tx: None,
281                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
282                    bypass_permissions: false,
283                    auto_approve_permissions: false,
284                    plan_read_only: false,
285                    can_async_resume: false,
286                    async_completion_sink: None,
287                    bash_completion_sink: None,
288                },
289            )
290            .await
291            .expect("list should succeed");
292        let ToolOutcome::Completed(list) = list else {
293            panic!("expected Completed")
294        };
295        let list_json: serde_json::Value = serde_json::from_str(&list.result).unwrap();
296        assert_eq!(list_json["topics"][0], "backend");
297        assert_eq!(list_json["count"], 1);
298
299        let clear = tool
300            .invoke(
301                json!({"action": "clear", "topic": "backend"}),
302                ToolCtx {
303                    executing_supervisor: None,
304                    session_id: Some(std::sync::Arc::from("session-1")),
305                    tool_call_id: std::sync::Arc::from("tool_call_clear"),
306                    event_tx: None,
307                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
308                    bypass_permissions: false,
309                    auto_approve_permissions: false,
310                    plan_read_only: false,
311                    can_async_resume: false,
312                    async_completion_sink: None,
313                    bash_completion_sink: None,
314                },
315            )
316            .await
317            .expect("clear should succeed");
318        let ToolOutcome::Completed(clear) = clear else {
319            panic!("expected Completed")
320        };
321        let clear_json: serde_json::Value = serde_json::from_str(&clear.result).unwrap();
322        assert_eq!(clear_json["action"], "clear");
323        assert_eq!(clear_json["deleted"], true);
324    }
325
326    #[tokio::test]
327    async fn session_note_read_reports_truncation_and_append_enforces_limit() {
328        let dir = tempfile::tempdir().unwrap();
329        let tool = SessionNoteTool::with_memory_store(MemoryStore::new(dir.path()));
330        let long_content = "x".repeat(32);
331
332        tool.invoke(
333            json!({"action": "replace", "topic": "default", "content": long_content}),
334            ToolCtx {
335                executing_supervisor: None,
336                session_id: Some(std::sync::Arc::from("session-2")),
337                tool_call_id: std::sync::Arc::from("tool_call_replace_long"),
338                event_tx: None,
339                available_tool_schemas: std::sync::Arc::from(Vec::new()),
340                bypass_permissions: false,
341                auto_approve_permissions: false,
342                plan_read_only: false,
343                can_async_resume: false,
344                async_completion_sink: None,
345                bash_completion_sink: None,
346            },
347        )
348        .await
349        .expect("replace should succeed");
350
351        let read = tool
352            .invoke(
353                json!({"action": "read", "topic": "default"}),
354                ToolCtx {
355                    executing_supervisor: None,
356                    session_id: Some(std::sync::Arc::from("session-2")),
357                    tool_call_id: std::sync::Arc::from("tool_call_read_long"),
358                    event_tx: None,
359                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
360                    bypass_permissions: false,
361                    auto_approve_permissions: false,
362                    plan_read_only: false,
363                    can_async_resume: false,
364                    async_completion_sink: None,
365                    bash_completion_sink: None,
366                },
367            )
368            .await
369            .expect("read should succeed");
370        let ToolOutcome::Completed(read) = read else {
371            panic!("expected Completed")
372        };
373        let read_json: serde_json::Value = serde_json::from_str(&read.result).unwrap();
374        assert_eq!(read_json["length_chars"], 32);
375        assert_eq!(read_json["body_truncated"], false);
376
377        tool.invoke(
378            json!({"action": "replace", "topic": "limit", "content": "x".repeat(crate::tools::session_memory::MAX_SESSION_NOTE_CHARS - 1)}),
379            ToolCtx {
380                executing_supervisor: None,
381                session_id: Some(std::sync::Arc::from("session-3")),
382                tool_call_id: std::sync::Arc::from("tool_call_replace_limit"),
383                event_tx: None,
384                available_tool_schemas: std::sync::Arc::from(Vec::new()),
385                bypass_permissions: false,
386                auto_approve_permissions: false,
387                plan_read_only: false,
388                can_async_resume: false,
389                async_completion_sink: None,
390                bash_completion_sink: None,
391            },
392        )
393        .await
394        .expect("replace near limit should succeed");
395
396        let append_err = tool
397            .invoke(
398                json!({"action": "append", "topic": "limit", "content": "y"}),
399                ToolCtx {
400                    executing_supervisor: None,
401                    session_id: Some(std::sync::Arc::from("session-3")),
402                    tool_call_id: std::sync::Arc::from("tool_call_append_limit"),
403                    event_tx: None,
404                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
405                    bypass_permissions: false,
406                    auto_approve_permissions: false,
407                    plan_read_only: false,
408                    can_async_resume: false,
409                    async_completion_sink: None,
410                    bash_completion_sink: None,
411                },
412            )
413            .await
414            .expect_err("append should exceed limit");
415        assert!(append_err
416            .to_string()
417            .contains("session note would exceed the limit"));
418    }
419}