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 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
47pub 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}