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, 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#[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}