1use async_trait::async_trait;
6use chrono::{DateTime, Utc};
7use everruns_core::AgentCapabilityConfig;
8use everruns_core::error::{AgentLoopError, Result};
9use everruns_core::session::Session;
10use everruns_core::session_file::{
11 FileInfo, FileStat, GrepMatch, GrepOptions, GrepSearchResult, InitialFile, SessionFile,
12 build_grep_search_result,
13};
14use everruns_core::traits::{
15 KeyInfo, SecretInfo, SessionFileSystem, SessionFileSystemFactory,
16 SessionFileSystemFactoryContext, SessionMutator, SessionStorageStore, SessionStore,
17};
18use everruns_core::typed_id::SessionId;
19use regex::Regex;
20use std::collections::{BTreeSet, HashMap};
21use std::sync::Arc;
22use tokio::sync::RwLock;
23use uuid::Uuid;
24
25#[derive(Debug, Default, Clone)]
29pub struct InMemorySessionStore {
30 sessions: Arc<RwLock<HashMap<SessionId, Session>>>,
31}
32
33impl InMemorySessionStore {
34 pub fn new() -> Self {
36 Self {
37 sessions: Arc::new(RwLock::new(HashMap::new())),
38 }
39 }
40
41 pub async fn add_session(&self, session: Session) {
43 self.sessions.write().await.insert(session.id, session);
44 }
45}
46
47#[async_trait]
48impl SessionStore for InMemorySessionStore {
49 async fn get_session(&self, session_id: SessionId) -> Result<Option<Session>> {
50 Ok(self.sessions.read().await.get(&session_id).cloned())
51 }
52}
53
54#[async_trait]
55impl SessionMutator for InMemorySessionStore {
56 async fn update_session_title(&self, session_id: SessionId, title: String) -> Result<Session> {
57 let mut sessions = self.sessions.write().await;
58 let session = sessions
59 .get_mut(&session_id)
60 .ok_or_else(|| AgentLoopError::store(format!("session not found: {session_id}")))?;
61 session.title = Some(title);
62 session.updated_at = Utc::now();
63 Ok(session.clone())
64 }
65
66 async fn upsert_session_capability(
67 &self,
68 session_id: SessionId,
69 capability: AgentCapabilityConfig,
70 ) -> Result<Session> {
71 let mut sessions = self.sessions.write().await;
72 let session = sessions
73 .get_mut(&session_id)
74 .ok_or_else(|| AgentLoopError::store(format!("session not found: {session_id}")))?;
75 if let Some(existing) = session
76 .capabilities
77 .iter_mut()
78 .find(|existing| existing.capability_id() == capability.capability_id())
79 {
80 *existing = capability;
81 } else {
82 session.capabilities.push(capability);
83 }
84 session.updated_at = Utc::now();
85 Ok(session.clone())
86 }
87
88 async fn remove_session_capability(
89 &self,
90 session_id: SessionId,
91 capability_id: &str,
92 ) -> Result<Session> {
93 let mut sessions = self.sessions.write().await;
94 let session = sessions
95 .get_mut(&session_id)
96 .ok_or_else(|| AgentLoopError::store(format!("session not found: {session_id}")))?;
97 session
98 .capabilities
99 .retain(|capability| capability.capability_id() != capability_id);
100 session.updated_at = Utc::now();
101 Ok(session.clone())
102 }
103}
104
105#[derive(Debug, Clone)]
106struct FileEntry {
107 file: SessionFile,
108}
109
110#[derive(Debug, Default, Clone)]
115pub struct InMemorySessionFileStore {
116 files: Arc<RwLock<HashMap<(SessionId, String), FileEntry>>>,
117}
118
119#[derive(Debug, Clone, Default)]
121pub struct InMemorySessionFileSystemFactory;
122
123#[async_trait]
124impl SessionFileSystemFactory for InMemorySessionFileSystemFactory {
125 fn name(&self) -> &'static str {
126 "InMemorySessionFileSystemFactory"
127 }
128
129 async fn create_session_file_system(
130 &self,
131 _context: SessionFileSystemFactoryContext,
132 ) -> Result<Arc<dyn SessionFileSystem>> {
133 Ok(Arc::new(InMemorySessionFileStore::new()))
134 }
135}
136
137impl InMemorySessionFileStore {
138 pub fn new() -> Self {
140 Self {
141 files: Arc::new(RwLock::new(HashMap::new())),
142 }
143 }
144
145 pub async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
147 let path = normalize_path(&file.path);
148 self.ensure_parent_directories(session_id, &path).await?;
149 self.upsert_file(
150 session_id,
151 &path,
152 &file.content,
153 &file.encoding,
154 file.is_readonly,
155 )
156 .await
157 .map(|_| ())
158 }
159
160 async fn ensure_parent_directories(&self, session_id: SessionId, path: &str) -> Result<()> {
161 let mut current = String::new();
162 for segment in path.trim_start_matches('/').split('/').collect::<Vec<_>>() {
163 if segment.is_empty() {
164 continue;
165 }
166 current.push('/');
167 current.push_str(segment);
168 let is_leaf = current == path;
169 if is_leaf {
170 break;
171 }
172 self.insert_directory_if_missing(session_id, ¤t)
173 .await?;
174 }
175 Ok(())
176 }
177
178 async fn insert_directory_if_missing(&self, session_id: SessionId, path: &str) -> Result<()> {
179 let path = normalize_path(path);
180 if path == "/" {
181 return Ok(());
182 }
183
184 let mut files = self.files.write().await;
185 files
186 .entry((session_id, path.clone()))
187 .or_insert_with(|| FileEntry {
188 file: SessionFile {
189 id: Uuid::now_v7(),
190 session_id: session_id.uuid(),
191 path: path.clone(),
192 name: FileInfo::name_from_path(&path),
193 content: None,
194 encoding: "text".to_string(),
195 is_directory: true,
196 is_readonly: false,
197 size_bytes: 0,
198 created_at: Utc::now(),
199 updated_at: Utc::now(),
200 },
201 });
202 Ok(())
203 }
204
205 async fn upsert_file(
206 &self,
207 session_id: SessionId,
208 path: &str,
209 content: &str,
210 encoding: &str,
211 is_readonly: bool,
212 ) -> Result<SessionFile> {
213 let now = Utc::now();
214 let normalized = normalize_path(path);
215 let mut files = self.files.write().await;
216 let key = (session_id, normalized.clone());
217
218 let file = files
219 .entry(key)
220 .and_modify(|entry| {
221 entry.file.content = Some(content.to_string());
222 entry.file.encoding = encoding.to_string();
223 entry.file.is_directory = false;
224 entry.file.is_readonly = is_readonly;
225 entry.file.size_bytes = content.len() as i64;
226 entry.file.updated_at = now;
227 })
228 .or_insert_with(|| FileEntry {
229 file: SessionFile {
230 id: Uuid::now_v7(),
231 session_id: session_id.uuid(),
232 path: normalized.clone(),
233 name: FileInfo::name_from_path(&normalized),
234 content: Some(content.to_string()),
235 encoding: encoding.to_string(),
236 is_directory: false,
237 is_readonly,
238 size_bytes: content.len() as i64,
239 created_at: now,
240 updated_at: now,
241 },
242 })
243 .file
244 .clone();
245
246 Ok(file)
247 }
248
249 pub async fn read_text(&self, session_id: SessionId, path: &str) -> Option<String> {
251 self.read_file(session_id, path)
252 .await
253 .ok()
254 .flatten()
255 .and_then(|file| file.content)
256 }
257}
258
259#[async_trait]
260impl SessionFileSystem for InMemorySessionFileStore {
261 async fn seed_initial_file(&self, session_id: SessionId, file: &InitialFile) -> Result<()> {
262 InMemorySessionFileStore::seed_initial_file(self, session_id, file).await
263 }
264
265 async fn read_file(&self, session_id: SessionId, path: &str) -> Result<Option<SessionFile>> {
266 let normalized = normalize_path(path);
267 if normalized == "/" {
268 return Ok(Some(root_directory(session_id)));
269 }
270
271 Ok(self
272 .files
273 .read()
274 .await
275 .get(&(session_id, normalized))
276 .map(|entry| entry.file.clone()))
277 }
278
279 async fn write_file(
280 &self,
281 session_id: SessionId,
282 path: &str,
283 content: &str,
284 encoding: &str,
285 ) -> Result<SessionFile> {
286 let normalized = normalize_path(path);
287 self.ensure_parent_directories(session_id, &normalized)
288 .await?;
289
290 if let Some(existing) = self.read_file(session_id, &normalized).await?
291 && existing.is_readonly
292 {
293 return Err(AgentLoopError::tool(format!(
294 "file is read-only: {}",
295 normalized
296 )));
297 }
298
299 self.upsert_file(session_id, &normalized, content, encoding, false)
300 .await
301 }
302
303 async fn delete_file(
304 &self,
305 session_id: SessionId,
306 path: &str,
307 recursive: bool,
308 ) -> Result<bool> {
309 let normalized = normalize_path(path);
310 if normalized == "/" {
311 return Ok(false);
312 }
313
314 let mut files = self.files.write().await;
315 let key = (session_id, normalized.clone());
316 let Some(existing) = files.get(&key).cloned() else {
317 return Ok(false);
318 };
319
320 if existing.file.is_readonly {
321 return Ok(false);
322 }
323
324 if existing.file.is_directory {
325 let prefix = format!("{normalized}/");
326 let has_children = files
327 .keys()
328 .any(|(sid, candidate)| *sid == session_id && candidate.starts_with(&prefix));
329 if has_children && !recursive {
330 return Ok(false);
331 }
332 files.retain(|(sid, candidate), _| {
333 !(*sid == session_id
334 && (candidate == &normalized || candidate.starts_with(&prefix)))
335 });
336 return Ok(true);
337 }
338
339 Ok(files.remove(&key).is_some())
340 }
341
342 async fn list_directory(&self, session_id: SessionId, path: &str) -> Result<Vec<FileInfo>> {
343 let normalized = normalize_path(path);
344 if normalized != "/" {
345 let Some(dir) = self.read_file(session_id, &normalized).await? else {
346 return Ok(vec![]);
347 };
348 if !dir.is_directory {
349 return Ok(vec![]);
350 }
351 }
352
353 let files = self.files.read().await;
354 let mut infos = Vec::new();
355 let mut seen = BTreeSet::new();
356
357 for ((sid, candidate), entry) in files.iter() {
358 if *sid != session_id {
359 continue;
360 }
361 if FileInfo::parent_path(candidate).as_deref() != Some(normalized.as_str()) {
362 continue;
363 }
364 if seen.insert(candidate.clone()) {
365 infos.push(file_info(&entry.file));
366 }
367 }
368
369 infos.sort_by(|a, b| a.path.cmp(&b.path));
370 Ok(infos)
371 }
372
373 async fn stat_file(&self, session_id: SessionId, path: &str) -> Result<Option<FileStat>> {
374 let normalized = normalize_path(path);
375 if normalized == "/" {
376 let root = root_directory(session_id);
377 return Ok(Some(FileStat {
378 path: root.path,
379 name: root.name,
380 is_directory: true,
381 is_readonly: false,
382 size_bytes: 0,
383 created_at: root.created_at,
384 updated_at: root.updated_at,
385 }));
386 }
387
388 Ok(self
389 .files
390 .read()
391 .await
392 .get(&(session_id, normalized))
393 .map(|entry| FileStat {
394 path: entry.file.path.clone(),
395 name: entry.file.name.clone(),
396 is_directory: entry.file.is_directory,
397 is_readonly: entry.file.is_readonly,
398 size_bytes: entry.file.size_bytes,
399 created_at: entry.file.created_at,
400 updated_at: entry.file.updated_at,
401 }))
402 }
403
404 async fn grep_files(
405 &self,
406 session_id: SessionId,
407 pattern: &str,
408 path_pattern: Option<&str>,
409 ) -> Result<Vec<GrepMatch>> {
410 let regex = Regex::new(pattern)
411 .map_err(|error| AgentLoopError::tool(format!("Invalid regex pattern: {error}")))?;
412 let path_pattern = path_pattern
413 .map(everruns_core::session_path::GrepPathPattern::new)
414 .transpose()?;
415 let files = self.files.read().await;
416 let mut matches = Vec::new();
417
418 for ((sid, path), entry) in files.iter() {
419 if *sid != session_id || entry.file.is_directory || entry.file.encoding != "text" {
420 continue;
421 }
422 if let Some(path_pattern) = &path_pattern
423 && !path_pattern.is_match(path)
424 {
425 continue;
426 }
427
428 let Some(content) = &entry.file.content else {
429 continue;
430 };
431
432 for (idx, line) in content.lines().enumerate() {
433 if regex.is_match(line) {
434 matches.push(GrepMatch {
435 path: path.clone(),
436 line_number: idx + 1,
437 line: line.to_string(),
438 });
439 }
440 }
441 }
442
443 Ok(matches)
444 }
445
446 async fn grep_files_with_options(
447 &self,
448 session_id: SessionId,
449 pattern: &str,
450 options: &GrepOptions,
451 ) -> Result<GrepSearchResult> {
452 let regex = Regex::new(pattern)
453 .map_err(|error| AgentLoopError::tool(format!("Invalid regex pattern: {error}")))?;
454 let path_pattern = options
455 .path_pattern
456 .as_deref()
457 .map(everruns_core::session_path::GrepPathPattern::new)
458 .transpose()?;
459 let files = self.files.read().await;
460 let text_files = files
461 .iter()
462 .filter(|((sid, path), entry)| {
463 *sid == session_id
464 && !entry.file.is_directory
465 && entry.file.encoding == "text"
466 && path_pattern
467 .as_ref()
468 .is_none_or(|matcher| matcher.is_match(path))
469 })
470 .filter_map(|((_, path), entry)| {
471 entry
472 .file
473 .content
474 .as_ref()
475 .map(|content| (path.clone(), content.clone()))
476 })
477 .collect();
478 Ok(build_grep_search_result(text_files, ®ex, options))
479 }
480
481 async fn create_directory(&self, session_id: SessionId, path: &str) -> Result<FileInfo> {
482 let normalized = normalize_path(path);
483 self.ensure_parent_directories(session_id, &normalized)
484 .await?;
485 self.insert_directory_if_missing(session_id, &normalized)
486 .await?;
487 let file = self
488 .read_file(session_id, &normalized)
489 .await?
490 .ok_or_else(|| AgentLoopError::store(format!("directory not found: {normalized}")))?;
491 Ok(file_info(&file))
492 }
493
494 fn is_mount_resolver(&self) -> bool {
495 false
496 }
497}
498
499#[derive(Debug, Default, Clone)]
501pub struct InMemorySessionStorageStore {
502 values: Arc<RwLock<HashMap<(SessionId, String), StorageValue>>>,
503 secrets: Arc<RwLock<HashMap<(SessionId, String), StorageValue>>>,
504}
505
506#[derive(Debug, Clone)]
507struct StorageValue {
508 value: String,
509 created_at: DateTime<Utc>,
510 updated_at: DateTime<Utc>,
511}
512
513impl InMemorySessionStorageStore {
514 pub fn new() -> Self {
516 Self {
517 values: Arc::new(RwLock::new(HashMap::new())),
518 secrets: Arc::new(RwLock::new(HashMap::new())),
519 }
520 }
521}
522
523#[async_trait]
524impl SessionStorageStore for InMemorySessionStorageStore {
525 async fn set_value(&self, session_id: SessionId, key: &str, value: &str) -> Result<()> {
526 upsert_storage(&self.values, session_id, key, value).await;
527 Ok(())
528 }
529
530 async fn get_value(&self, session_id: SessionId, key: &str) -> Result<Option<String>> {
531 Ok(self
532 .values
533 .read()
534 .await
535 .get(&(session_id, key.to_string()))
536 .map(|value| value.value.clone()))
537 }
538
539 async fn delete_value(&self, session_id: SessionId, key: &str) -> Result<bool> {
540 Ok(self
541 .values
542 .write()
543 .await
544 .remove(&(session_id, key.to_string()))
545 .is_some())
546 }
547
548 async fn list_keys(&self, session_id: SessionId) -> Result<Vec<KeyInfo>> {
549 Ok(list_storage(&self.values, session_id)
550 .await
551 .into_iter()
552 .map(|(key, value)| KeyInfo {
553 key,
554 created_at: value.created_at,
555 updated_at: value.updated_at,
556 })
557 .collect())
558 }
559
560 async fn set_secret(&self, session_id: SessionId, name: &str, value: &str) -> Result<()> {
561 upsert_storage(&self.secrets, session_id, name, value).await;
562 Ok(())
563 }
564
565 async fn get_secret(&self, session_id: SessionId, name: &str) -> Result<Option<String>> {
566 Ok(self
567 .secrets
568 .read()
569 .await
570 .get(&(session_id, name.to_string()))
571 .map(|value| value.value.clone()))
572 }
573
574 async fn delete_secret(&self, session_id: SessionId, name: &str) -> Result<bool> {
575 Ok(self
576 .secrets
577 .write()
578 .await
579 .remove(&(session_id, name.to_string()))
580 .is_some())
581 }
582
583 async fn list_secrets(&self, session_id: SessionId) -> Result<Vec<SecretInfo>> {
584 Ok(list_storage(&self.secrets, session_id)
585 .await
586 .into_iter()
587 .map(|(name, value)| SecretInfo {
588 name,
589 created_at: value.created_at,
590 updated_at: value.updated_at,
591 })
592 .collect())
593 }
594}
595
596async fn upsert_storage(
597 map: &Arc<RwLock<HashMap<(SessionId, String), StorageValue>>>,
598 session_id: SessionId,
599 key: &str,
600 value: &str,
601) {
602 let mut map = map.write().await;
603 let now = Utc::now();
604 map.entry((session_id, key.to_string()))
605 .and_modify(|stored| {
606 stored.value = value.to_string();
607 stored.updated_at = now;
608 })
609 .or_insert_with(|| StorageValue {
610 value: value.to_string(),
611 created_at: now,
612 updated_at: now,
613 });
614}
615
616async fn list_storage(
617 map: &Arc<RwLock<HashMap<(SessionId, String), StorageValue>>>,
618 session_id: SessionId,
619) -> Vec<(String, StorageValue)> {
620 let mut values: Vec<_> = map
621 .read()
622 .await
623 .iter()
624 .filter(|((sid, _), _)| *sid == session_id)
625 .map(|((_, key), value)| (key.clone(), value.clone()))
626 .collect();
627 values.sort_by(|a, b| a.0.cmp(&b.0));
628 values
629}
630
631fn normalize_path(path: &str) -> String {
632 everruns_core::session_path::to_session_path(path)
635}
636
637fn root_directory(session_id: SessionId) -> SessionFile {
638 let now = Utc::now();
639 SessionFile {
640 id: Uuid::nil(),
641 session_id: session_id.uuid(),
642 path: "/".to_string(),
643 name: "/".to_string(),
644 content: None,
645 encoding: "text".to_string(),
646 is_directory: true,
647 is_readonly: false,
648 size_bytes: 0,
649 created_at: now,
650 updated_at: now,
651 }
652}
653
654fn file_info(file: &SessionFile) -> FileInfo {
655 FileInfo {
656 id: file.id,
657 session_id: file.session_id,
658 path: file.path.clone(),
659 name: file.name.clone(),
660 is_directory: file.is_directory,
661 is_readonly: file.is_readonly,
662 size_bytes: file.size_bytes,
663 created_at: file.created_at,
664 updated_at: file.updated_at,
665 }
666}
667
668#[cfg(test)]
669mod tests {
670 use super::*;
671
672 #[tokio::test]
673 async fn grep_path_pattern_supports_globs_and_substring_compatibility() {
674 let store = InMemorySessionFileStore::new();
675 let session = SessionId::from_seed(1);
676 for path in [
677 "/src/lib.rs",
678 "/src/nested/mod.rs",
679 "/docs/readme.md",
680 "/notes.txt",
681 ] {
682 store
683 .write_file(session, path, "needle", "text")
684 .await
685 .unwrap();
686 }
687
688 let mut glob_paths: Vec<_> = store
689 .grep_files(session, "needle", Some("src/**/*.rs"))
690 .await
691 .unwrap()
692 .into_iter()
693 .map(|hit| hit.path)
694 .collect();
695 glob_paths.sort();
696 assert_eq!(glob_paths, vec!["/src/lib.rs", "/src/nested/mod.rs"]);
697
698 let substring_paths: Vec<_> = store
699 .grep_files(session, "needle", Some("docs"))
700 .await
701 .unwrap()
702 .into_iter()
703 .map(|hit| hit.path)
704 .collect();
705 assert_eq!(substring_paths, vec!["/docs/readme.md"]);
706 }
707}