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, ToolError, ToolExecutionContext, ToolResult};
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`.
50#[allow(dead_code)]
51pub type MemoryNoteTool = SessionNoteTool;
52
53#[async_trait]
54impl Tool for SessionNoteTool {
55    fn name(&self) -> &str {
56        TOOL_NAME
57    }
58
59    fn description(&self) -> &str {
60        TOOL_DESCRIPTION
61    }
62
63    fn mutability(&self) -> crate::ToolMutability {
64        crate::ToolMutability::Mutating
65    }
66
67    fn call_mutability(&self, args: &serde_json::Value) -> crate::ToolMutability {
68        let action = args
69            .get("action")
70            .and_then(|v| v.as_str())
71            .unwrap_or("")
72            .trim()
73            .to_ascii_lowercase();
74        match action.as_str() {
75            "read" | "list_topics" => crate::ToolMutability::ReadOnly,
76            _ => crate::ToolMutability::Mutating,
77        }
78    }
79
80    fn call_concurrency_safe(&self, args: &serde_json::Value) -> bool {
81        matches!(self.call_mutability(args), crate::ToolMutability::ReadOnly)
82    }
83
84    fn parameters_schema(&self) -> serde_json::Value {
85        json!({
86            "type": "object",
87            "properties": {
88                "action": {
89                    "type": "string",
90                    "description": "Operation to perform on the note.",
91                    "enum": ["read", "append", "replace", "clear", "list_topics"]
92                },
93                "content": {
94                    "type": "string",
95                    "description": "Note content to append/replace (markdown). Required for append/replace."
96                },
97                "topic": {
98                    "type": "string",
99                    "description": "Optional topic name (alphanumeric/dash/underscore, max 50 chars). Defaults to 'default'. Use separate topics for unrelated workstreams."
100                }
101            },
102            "required": ["action"]
103        })
104    }
105
106    async fn execute(&self, _args: serde_json::Value) -> Result<ToolResult, ToolError> {
107        Err(ToolError::Execution(format!(
108            "{TOOL_NAME} must be executed with ToolExecutionContext (session_id required)"
109        )))
110    }
111
112    async fn execute_with_context(
113        &self,
114        args: serde_json::Value,
115        ctx: ToolExecutionContext<'_>,
116    ) -> Result<ToolResult, ToolError> {
117        let Some(session_id) = ctx.session_id else {
118            return Err(ToolError::Execution(
119                "missing session_id in tool context".to_string(),
120            ));
121        };
122
123        let action_raw = args.get("action").and_then(|v| v.as_str()).unwrap_or("");
124        let action = parse_session_note_action(action_raw)?;
125        let topic = args.get("topic").and_then(|v| v.as_str());
126        let content = args.get("content").and_then(|v| v.as_str());
127
128        execute_session_memory_action(
129            &self.memory_store,
130            session_id,
131            action,
132            topic,
133            content,
134            None,
135            SESSION_NOTE_ACTION_NAMES,
136        )
137        .await
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn session_note_schema_requires_action() {
147        let tool = SessionNoteTool::new();
148        let schema = tool.parameters_schema();
149        assert_eq!(schema["required"], json!(["action"]));
150        assert_eq!(tool.name(), TOOL_NAME);
151        assert_eq!(
152            schema["properties"]["action"]["enum"],
153            json!(["read", "append", "replace", "clear", "list_topics"])
154        );
155    }
156
157    #[test]
158    fn session_note_schema_has_topic_field() {
159        let tool = SessionNoteTool::new();
160        let schema = tool.parameters_schema();
161        assert!(schema["properties"]["topic"].is_object());
162        assert_eq!(schema["properties"]["topic"]["type"], "string");
163    }
164
165    #[tokio::test]
166    async fn session_note_requires_session_context() {
167        let tool = SessionNoteTool::new();
168        let result = tool
169            .execute_with_context(
170                json!({"action": "read"}),
171                ToolExecutionContext::none("tool_call"),
172            )
173            .await;
174
175        assert!(matches!(
176            result,
177            Err(ToolError::Execution(msg)) if msg.contains("session_id")
178        ));
179    }
180
181    #[tokio::test]
182    async fn session_note_validates_action_and_content_before_io() {
183        let tool = SessionNoteTool::new();
184
185        let unknown = tool
186            .execute_with_context(
187                json!({"action": "unknown"}),
188                ToolExecutionContext {
189                    session_id: Some("session-1"),
190                    tool_call_id: "tool_call_unknown",
191                    event_tx: None,
192                    available_tool_schemas: None,
193                },
194            )
195            .await;
196        assert!(matches!(
197            unknown,
198            Err(ToolError::InvalidArguments(msg)) if msg.contains("action must be one of")
199        ));
200
201        let missing_content = tool
202            .execute_with_context(
203                json!({"action": "replace"}),
204                ToolExecutionContext {
205                    session_id: Some("session-1"),
206                    tool_call_id: "tool_call_replace",
207                    event_tx: None,
208                    available_tool_schemas: None,
209                },
210            )
211            .await;
212        assert!(matches!(
213            missing_content,
214            Err(ToolError::InvalidArguments(msg)) if msg.contains("content is required")
215        ));
216    }
217
218    #[tokio::test]
219    async fn session_note_uses_memory_store_session_topics() {
220        let dir = tempfile::tempdir().unwrap();
221        let tool = SessionNoteTool::with_memory_store(MemoryStore::new(dir.path()));
222
223        let append = tool
224            .execute_with_context(
225                json!({"action": "append", "topic": "backend", "content": "API finalized"}),
226                ToolExecutionContext {
227                    session_id: Some("session-1"),
228                    tool_call_id: "tool_call_append",
229                    event_tx: None,
230                    available_tool_schemas: None,
231                },
232            )
233            .await
234            .expect("append should succeed");
235        let append_json: serde_json::Value = serde_json::from_str(&append.result).unwrap();
236        assert_eq!(append_json["action"], "append");
237        assert_eq!(append_json["length_chars"], "API finalized".chars().count());
238
239        let read = tool
240            .execute_with_context(
241                json!({"action": "read", "topic": "backend"}),
242                ToolExecutionContext {
243                    session_id: Some("session-1"),
244                    tool_call_id: "tool_call_read",
245                    event_tx: None,
246                    available_tool_schemas: None,
247                },
248            )
249            .await
250            .expect("read should succeed");
251        let read_json: serde_json::Value = serde_json::from_str(&read.result).unwrap();
252        assert_eq!(read_json["action"], "read");
253        assert_eq!(read_json["content"], "API finalized");
254        assert_eq!(read_json["length_chars"], "API finalized".chars().count());
255        assert_eq!(read_json["body_truncated"], false);
256
257        let list = tool
258            .execute_with_context(
259                json!({"action": "list_topics"}),
260                ToolExecutionContext {
261                    session_id: Some("session-1"),
262                    tool_call_id: "tool_call_list",
263                    event_tx: None,
264                    available_tool_schemas: None,
265                },
266            )
267            .await
268            .expect("list should succeed");
269        let list_json: serde_json::Value = serde_json::from_str(&list.result).unwrap();
270        assert_eq!(list_json["topics"][0], "backend");
271        assert_eq!(list_json["count"], 1);
272
273        let clear = tool
274            .execute_with_context(
275                json!({"action": "clear", "topic": "backend"}),
276                ToolExecutionContext {
277                    session_id: Some("session-1"),
278                    tool_call_id: "tool_call_clear",
279                    event_tx: None,
280                    available_tool_schemas: None,
281                },
282            )
283            .await
284            .expect("clear should succeed");
285        let clear_json: serde_json::Value = serde_json::from_str(&clear.result).unwrap();
286        assert_eq!(clear_json["action"], "clear");
287        assert_eq!(clear_json["deleted"], true);
288    }
289
290    #[tokio::test]
291    async fn session_note_read_reports_truncation_and_append_enforces_limit() {
292        let dir = tempfile::tempdir().unwrap();
293        let tool = SessionNoteTool::with_memory_store(MemoryStore::new(dir.path()));
294        let long_content = "x".repeat(32);
295
296        tool.execute_with_context(
297            json!({"action": "replace", "topic": "default", "content": long_content}),
298            ToolExecutionContext {
299                session_id: Some("session-2"),
300                tool_call_id: "tool_call_replace_long",
301                event_tx: None,
302                available_tool_schemas: None,
303            },
304        )
305        .await
306        .expect("replace should succeed");
307
308        let read = tool
309            .execute_with_context(
310                json!({"action": "read", "topic": "default"}),
311                ToolExecutionContext {
312                    session_id: Some("session-2"),
313                    tool_call_id: "tool_call_read_long",
314                    event_tx: None,
315                    available_tool_schemas: None,
316                },
317            )
318            .await
319            .expect("read should succeed");
320        let read_json: serde_json::Value = serde_json::from_str(&read.result).unwrap();
321        assert_eq!(read_json["length_chars"], 32);
322        assert_eq!(read_json["body_truncated"], false);
323
324        tool.execute_with_context(
325            json!({"action": "replace", "topic": "limit", "content": "x".repeat(crate::tools::session_memory::MAX_SESSION_NOTE_CHARS - 1)}),
326            ToolExecutionContext {
327                session_id: Some("session-3"),
328                tool_call_id: "tool_call_replace_limit",
329                event_tx: None,
330                available_tool_schemas: None,
331            },
332        )
333        .await
334        .expect("replace near limit should succeed");
335
336        let append_err = tool
337            .execute_with_context(
338                json!({"action": "append", "topic": "limit", "content": "y"}),
339                ToolExecutionContext {
340                    session_id: Some("session-3"),
341                    tool_call_id: "tool_call_append_limit",
342                    event_tx: None,
343                    available_tool_schemas: None,
344                },
345            )
346            .await
347            .expect_err("append should exceed limit");
348        assert!(append_err
349            .to_string()
350            .contains("session note would exceed the limit"));
351    }
352}