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