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_caller: None,
172                    executing_supervisor: None,
173                    session_id: Some(std::sync::Arc::from("session-1")),
174                    tool_call_id: std::sync::Arc::from("tool_call_unknown"),
175                    event_tx: None,
176                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
177                    bypass_permissions: false,
178                    auto_approve_permissions: false,
179                    plan_read_only: false,
180                    can_async_resume: false,
181                    async_completion_sink: None,
182                    bash_completion_sink: None,
183                },
184            )
185            .await;
186        assert!(matches!(
187            unknown,
188            Err(ToolError::InvalidArguments(msg)) if msg.contains("action must be one of")
189        ));
190
191        let missing_content = tool
192            .invoke(
193                json!({"action": "replace"}),
194                ToolCtx {
195                    executing_caller: None,
196                    executing_supervisor: None,
197                    session_id: Some(std::sync::Arc::from("session-1")),
198                    tool_call_id: std::sync::Arc::from("tool_call_replace"),
199                    event_tx: None,
200                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
201                    bypass_permissions: false,
202                    auto_approve_permissions: false,
203                    plan_read_only: false,
204                    can_async_resume: false,
205                    async_completion_sink: None,
206                    bash_completion_sink: None,
207                },
208            )
209            .await;
210        assert!(matches!(
211            missing_content,
212            Err(ToolError::InvalidArguments(msg)) if msg.contains("content is required")
213        ));
214    }
215
216    #[tokio::test]
217    async fn session_note_uses_memory_store_session_topics() {
218        let dir = tempfile::tempdir().unwrap();
219        let tool = SessionNoteTool::with_memory_store(MemoryStore::new(dir.path()));
220
221        let append = tool
222            .invoke(
223                json!({"action": "append", "topic": "backend", "content": "API finalized"}),
224                ToolCtx {
225                    executing_caller: None,
226                    executing_supervisor: None,
227                    session_id: Some(std::sync::Arc::from("session-1")),
228                    tool_call_id: std::sync::Arc::from("tool_call_append"),
229                    event_tx: None,
230                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
231                    bypass_permissions: false,
232                    auto_approve_permissions: false,
233                    plan_read_only: false,
234                    can_async_resume: false,
235                    async_completion_sink: None,
236                    bash_completion_sink: None,
237                },
238            )
239            .await
240            .expect("append should succeed");
241        let ToolOutcome::Completed(append) = append else {
242            panic!("expected Completed")
243        };
244        let append_json: serde_json::Value = serde_json::from_str(&append.result).unwrap();
245        assert_eq!(append_json["action"], "append");
246        assert_eq!(append_json["length_chars"], "API finalized".chars().count());
247
248        let read = tool
249            .invoke(
250                json!({"action": "read", "topic": "backend"}),
251                ToolCtx {
252                    executing_caller: None,
253                    executing_supervisor: None,
254                    session_id: Some(std::sync::Arc::from("session-1")),
255                    tool_call_id: std::sync::Arc::from("tool_call_read"),
256                    event_tx: None,
257                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
258                    bypass_permissions: false,
259                    auto_approve_permissions: false,
260                    plan_read_only: false,
261                    can_async_resume: false,
262                    async_completion_sink: None,
263                    bash_completion_sink: None,
264                },
265            )
266            .await
267            .expect("read should succeed");
268        let ToolOutcome::Completed(read) = read else {
269            panic!("expected Completed")
270        };
271        let read_json: serde_json::Value = serde_json::from_str(&read.result).unwrap();
272        assert_eq!(read_json["action"], "read");
273        assert_eq!(read_json["content"], "API finalized");
274        assert_eq!(read_json["length_chars"], "API finalized".chars().count());
275        assert_eq!(read_json["body_truncated"], false);
276
277        let list = tool
278            .invoke(
279                json!({"action": "list_topics"}),
280                ToolCtx {
281                    executing_caller: None,
282                    executing_supervisor: None,
283                    session_id: Some(std::sync::Arc::from("session-1")),
284                    tool_call_id: std::sync::Arc::from("tool_call_list"),
285                    event_tx: None,
286                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
287                    bypass_permissions: false,
288                    auto_approve_permissions: false,
289                    plan_read_only: false,
290                    can_async_resume: false,
291                    async_completion_sink: None,
292                    bash_completion_sink: None,
293                },
294            )
295            .await
296            .expect("list should succeed");
297        let ToolOutcome::Completed(list) = list else {
298            panic!("expected Completed")
299        };
300        let list_json: serde_json::Value = serde_json::from_str(&list.result).unwrap();
301        assert_eq!(list_json["topics"][0], "backend");
302        assert_eq!(list_json["count"], 1);
303
304        let clear = tool
305            .invoke(
306                json!({"action": "clear", "topic": "backend"}),
307                ToolCtx {
308                    executing_caller: None,
309                    executing_supervisor: None,
310                    session_id: Some(std::sync::Arc::from("session-1")),
311                    tool_call_id: std::sync::Arc::from("tool_call_clear"),
312                    event_tx: None,
313                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
314                    bypass_permissions: false,
315                    auto_approve_permissions: false,
316                    plan_read_only: false,
317                    can_async_resume: false,
318                    async_completion_sink: None,
319                    bash_completion_sink: None,
320                },
321            )
322            .await
323            .expect("clear should succeed");
324        let ToolOutcome::Completed(clear) = clear else {
325            panic!("expected Completed")
326        };
327        let clear_json: serde_json::Value = serde_json::from_str(&clear.result).unwrap();
328        assert_eq!(clear_json["action"], "clear");
329        assert_eq!(clear_json["deleted"], true);
330    }
331
332    #[tokio::test]
333    async fn session_note_read_reports_truncation_and_append_enforces_limit() {
334        let dir = tempfile::tempdir().unwrap();
335        let tool = SessionNoteTool::with_memory_store(MemoryStore::new(dir.path()));
336        let long_content = "x".repeat(32);
337
338        tool.invoke(
339            json!({"action": "replace", "topic": "default", "content": long_content}),
340            ToolCtx {
341                executing_caller: None,
342                executing_supervisor: None,
343                session_id: Some(std::sync::Arc::from("session-2")),
344                tool_call_id: std::sync::Arc::from("tool_call_replace_long"),
345                event_tx: None,
346                available_tool_schemas: std::sync::Arc::from(Vec::new()),
347                bypass_permissions: false,
348                auto_approve_permissions: false,
349                plan_read_only: false,
350                can_async_resume: false,
351                async_completion_sink: None,
352                bash_completion_sink: None,
353            },
354        )
355        .await
356        .expect("replace should succeed");
357
358        let read = tool
359            .invoke(
360                json!({"action": "read", "topic": "default"}),
361                ToolCtx {
362                    executing_caller: None,
363                    executing_supervisor: None,
364                    session_id: Some(std::sync::Arc::from("session-2")),
365                    tool_call_id: std::sync::Arc::from("tool_call_read_long"),
366                    event_tx: None,
367                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
368                    bypass_permissions: false,
369                    auto_approve_permissions: false,
370                    plan_read_only: false,
371                    can_async_resume: false,
372                    async_completion_sink: None,
373                    bash_completion_sink: None,
374                },
375            )
376            .await
377            .expect("read should succeed");
378        let ToolOutcome::Completed(read) = read else {
379            panic!("expected Completed")
380        };
381        let read_json: serde_json::Value = serde_json::from_str(&read.result).unwrap();
382        assert_eq!(read_json["length_chars"], 32);
383        assert_eq!(read_json["body_truncated"], false);
384
385        tool.invoke(
386            json!({"action": "replace", "topic": "limit", "content": "x".repeat(crate::tools::session_memory::MAX_SESSION_NOTE_CHARS - 1)}),
387            ToolCtx {
388                executing_caller: None,
389                executing_supervisor: None,
390                session_id: Some(std::sync::Arc::from("session-3")),
391                tool_call_id: std::sync::Arc::from("tool_call_replace_limit"),
392                event_tx: None,
393                available_tool_schemas: std::sync::Arc::from(Vec::new()),
394                bypass_permissions: false,
395                auto_approve_permissions: false,
396                plan_read_only: false,
397                can_async_resume: false,
398                async_completion_sink: None,
399                bash_completion_sink: None,
400            },
401        )
402        .await
403        .expect("replace near limit should succeed");
404
405        let append_err = tool
406            .invoke(
407                json!({"action": "append", "topic": "limit", "content": "y"}),
408                ToolCtx {
409                    executing_caller: None,
410                    executing_supervisor: None,
411                    session_id: Some(std::sync::Arc::from("session-3")),
412                    tool_call_id: std::sync::Arc::from("tool_call_append_limit"),
413                    event_tx: None,
414                    available_tool_schemas: std::sync::Arc::from(Vec::new()),
415                    bypass_permissions: false,
416                    auto_approve_permissions: false,
417                    plan_read_only: false,
418                    can_async_resume: false,
419                    async_completion_sink: None,
420                    bash_completion_sink: None,
421                },
422            )
423            .await
424            .expect_err("append should exceed limit");
425        assert!(append_err
426            .to_string()
427            .contains("session note would exceed the limit"));
428    }
429}