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