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_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}