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