1use std::collections::{HashMap, HashSet};
21use std::io;
22use std::path::{Path, PathBuf};
23
24use base64::Engine;
25use chrono::{DateTime, Utc};
26use fs2::FileExt;
27use serde::{Deserialize, Serialize};
28use tokio::fs;
29use tokio::sync::{Mutex, OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock};
30use uuid::Uuid;
31
32use bamboo_domain::ProviderModelRef;
33use bamboo_domain::ReasoningEffort;
34use bamboo_domain::{ProjectId, Role, Session, SessionKind, TokenBudgetUsage};
35
36use crate::search_index::{should_index_session, SessionSearchIndex};
37use bamboo_domain::AttachmentReader;
38use bamboo_domain::Storage;
39
40pub(crate) fn other_io_error(message: impl Into<String>) -> io::Error {
41 io::Error::other(message.into())
42}
43
44const RUNTIME_SIDECAR_FILE: &str = "runtime.json";
47const SESSIONS_INDEX_VERSION: u32 = 4;
48
49const TOKEN_USAGE_FILE: &str = "token-usage.jsonl";
52
53const RUNTIME_SIDECAR_MIGRATION_MARKER: &str = ".runtime_sidecar_migrated";
56const SESSION_LIFECYCLE_LOCK_FILE: &str = ".session-lifecycle.lock";
57
58pub(crate) struct SessionLifecycleReadGuard {
59 _process: OwnedRwLockReadGuard<()>,
60 file: std::fs::File,
61}
62
63impl Drop for SessionLifecycleReadGuard {
64 fn drop(&mut self) {
65 let _ = FileExt::unlock(&self.file);
66 }
67}
68
69pub(crate) struct SessionLifecycleWriteGuard {
70 _process: OwnedRwLockWriteGuard<()>,
71 file: std::fs::File,
72}
73
74impl Drop for SessionLifecycleWriteGuard {
75 fn drop(&mut self) {
76 let _ = FileExt::unlock(&self.file);
77 }
78}
79
80fn runtime_sidecar_snapshot(session: &Session) -> Session {
84 let mut snapshot = session.clone();
85 snapshot.messages.clear();
86 if let Some(metadata) = snapshot.runtime_metadata.as_mut() {
87 metadata.session_inbox_admission = None;
92 }
93 snapshot
94}
95
96fn overlay_runtime_sidecar(main: Session, sidecar: Option<Session>) -> Session {
103 match sidecar {
104 Some(mut side) => {
105 let admission = main
106 .runtime_metadata
107 .as_ref()
108 .and_then(|metadata| metadata.session_inbox_admission.clone());
109 side.messages = main.messages;
110 if let Some(admission) = admission {
111 side.runtime_metadata
112 .get_or_insert_with(Default::default)
113 .session_inbox_admission = Some(admission);
114 } else if let Some(metadata) = side.runtime_metadata.as_mut() {
115 metadata.session_inbox_admission = None;
116 }
117 side
118 }
119 None => main,
120 }
121}
122
123fn normalized_project_id(session: &Session) -> Option<String> {
128 let raw = session.project_id_meta()?;
129 match raw.trim().parse::<ProjectId>() {
130 Ok(project_id) => Some(project_id.into_string()),
131 Err(error) => {
132 tracing::warn!(
133 session_id = %session.id,
134 %error,
135 "ignoring malformed Project id while updating session index"
136 );
137 None
138 }
139 }
140}
141
142pub(crate) fn validate_session_id(session_id: &str) -> io::Result<()> {
146 if session_id.is_empty()
147 || session_id.contains('/')
148 || session_id.contains('\\')
149 || session_id.contains("..")
150 {
151 return Err(other_io_error(format!("invalid session id: {session_id}")));
152 }
153 Ok(())
154}
155
156#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
160pub struct SessionPlacement {
161 pub kind: String,
164 pub host: String,
167}
168
169#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct SessionIndexEntry {
171 pub id: String,
172 pub kind: SessionKind,
173 pub rel_path: String,
175 pub title: String,
176 #[serde(default)]
177 pub title_version: u64,
178 pub pinned: bool,
179 pub parent_session_id: Option<String>,
180 pub root_session_id: String,
181 pub spawn_depth: u32,
182 #[serde(default)]
183 pub model: String,
184 #[serde(default, skip_serializing_if = "Option::is_none")]
185 pub model_ref: Option<ProviderModelRef>,
186 #[serde(default, skip_serializing_if = "Option::is_none")]
187 pub reasoning_effort: Option<ReasoningEffort>,
188 #[serde(default, skip_serializing_if = "Option::is_none")]
190 pub workspace_path: Option<String>,
191 #[serde(default, skip_serializing_if = "Option::is_none")]
193 pub project_id: Option<String>,
194 #[serde(default, skip_serializing_if = "Option::is_none")]
197 pub gold_config_json: Option<String>,
198 #[serde(default, skip_serializing_if = "Option::is_none")]
200 pub created_by_schedule_id: Option<String>,
201 #[serde(default, skip_serializing_if = "Option::is_none")]
203 pub schedule_run_id: Option<String>,
204 pub created_at: DateTime<Utc>,
205 pub updated_at: DateTime<Utc>,
206 pub last_activity_at: DateTime<Utc>,
207 pub message_count: usize,
208 pub has_attachments: bool,
209 #[serde(default)]
213 pub has_pending_question: bool,
214 #[serde(default, skip_serializing_if = "Option::is_none")]
218 pub plan_mode: Option<bamboo_domain::PlanModeState>,
219 #[serde(default)]
223 pub bypass_permissions: bool,
224 #[serde(default)]
226 pub permission_mode: bamboo_domain::SessionPermissionMode,
227 #[serde(default, skip_serializing_if = "Option::is_none")]
230 pub last_run_status: Option<String>,
231 #[serde(default, skip_serializing_if = "Option::is_none")]
233 pub last_run_error: Option<String>,
234 #[serde(default, skip_serializing_if = "Option::is_none")]
239 pub token_usage: Option<TokenBudgetUsage>,
240 #[serde(default, skip_serializing_if = "Option::is_none")]
247 pub subagent_type: Option<String>,
248 #[serde(default, skip_serializing_if = "Option::is_none")]
252 pub lifecycle: Option<String>,
253 #[serde(default, skip_serializing_if = "Option::is_none")]
257 pub resident_name: Option<String>,
258 #[serde(default, skip_serializing_if = "Option::is_none")]
263 pub placement: Option<SessionPlacement>,
264}
265
266#[derive(Debug, Clone, Serialize, Deserialize)]
267pub struct SessionsIndex {
268 pub version: u32,
269 pub updated_at: DateTime<Utc>,
270 pub sessions: HashMap<String, SessionIndexEntry>,
271}
272
273impl SessionsIndex {
274 fn empty() -> Self {
275 Self {
276 version: SESSIONS_INDEX_VERSION,
277 updated_at: Utc::now(),
278 sessions: HashMap::new(),
279 }
280 }
281}
282
283#[derive(Debug)]
284pub struct SessionStoreV2 {
285 bamboo_home_dir: PathBuf,
286 sessions_dir: PathBuf,
287 index_path: PathBuf,
288 search_index: SessionSearchIndex,
289 index: RwLock<SessionsIndex>,
290 write_lock: Mutex<()>,
292 session_lifecycle_lock: std::sync::Arc<RwLock<()>>,
296}
297
298impl SessionStoreV2 {
299 pub async fn new(bamboo_home_dir: PathBuf) -> io::Result<Self> {
310 let sessions_dir = bamboo_home_dir.join("sessions");
311 let index_path = bamboo_home_dir.join("sessions.json");
312 let search_index = SessionSearchIndex::new(bamboo_home_dir.join("session_search.db"));
313
314 fs::create_dir_all(&sessions_dir).await?;
315 search_index.init().await?;
316
317 let mut needs_rebuild = false;
321 let index = if index_path.exists() {
322 let raw = fs::read_to_string(&index_path).await?;
323 match serde_json::from_str::<SessionsIndex>(&raw) {
324 Ok(index) if index.version >= SESSIONS_INDEX_VERSION => index,
325 Ok(index) => {
326 tracing::info!(
327 "migrating sessions index from version {} to version {} by rebuilding from session.json",
328 index.version,
329 SESSIONS_INDEX_VERSION,
330 );
331 needs_rebuild = true;
332 let mut rebuilding = SessionsIndex::empty();
333 rebuilding.version = index.version.min(SESSIONS_INDEX_VERSION - 1);
337 rebuilding
338 }
339 Err(error) => {
340 let backup_path = bamboo_home_dir.join("sessions.json.bak");
349 match fs::rename(&index_path, &backup_path).await {
350 Ok(()) => tracing::error!(
351 "sessions.json is corrupt ({error}); backed up to {} and rebuilding \
352 the index by scanning the session tree",
353 backup_path.display()
354 ),
355 Err(rename_error) => tracing::error!(
356 "sessions.json is corrupt ({error}); failed to back it up to {} \
357 ({rename_error}); rebuilding the index from disk anyway",
358 backup_path.display()
359 ),
360 }
361 needs_rebuild = true;
362 let mut rebuilding = SessionsIndex::empty();
363 rebuilding.version = 0;
364 rebuilding
365 }
366 }
367 } else {
368 let index = SessionsIndex::empty();
369 let tmp = index_path.with_extension(format!("json.tmp.{}", Uuid::new_v4()));
371 fs::write(
372 &tmp,
373 serde_json::to_vec_pretty(&index).map_err(|e| other_io_error(e.to_string()))?,
374 )
375 .await?;
376 atomic_rename(&tmp, &index_path).await?;
377 index
378 };
379
380 let storage = Self {
381 bamboo_home_dir,
382 sessions_dir,
383 index_path,
384 search_index,
385 index: RwLock::new(index),
386 write_lock: Mutex::new(()),
387 session_lifecycle_lock: std::sync::Arc::new(RwLock::new(())),
388 };
389
390 if needs_rebuild {
391 storage.rebuild_index_from_disk().await?;
392 }
393
394 Ok(storage)
395 }
396
397 async fn rebuild_index_from_disk(&self) -> io::Result<()> {
419 let mut recovered = 0usize;
420
421 let mut root_dirs = match fs::read_dir(&self.sessions_dir).await {
422 Ok(rd) => rd,
423 Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(()),
425 Err(error) => return Err(error),
426 };
427
428 loop {
429 let root_entry = match root_dirs.next_entry().await {
433 Ok(Some(entry)) => entry,
434 Ok(None) => break,
435 Err(error) => {
436 tracing::warn!("index rebuild: error scanning sessions dir: {error}");
437 break;
438 }
439 };
440 if !root_entry
441 .file_type()
442 .await
443 .map(|t| t.is_dir())
444 .unwrap_or(false)
445 {
446 continue;
447 }
448 let Ok(root_id) = root_entry.file_name().into_string() else {
449 continue;
451 };
452
453 if let Some(session) = Self::load_session_from_dir(&root_entry.path(), &root_id).await {
455 let rel_path = Self::root_rel_path(&root_id);
456 match self.upsert_index_from_session(&session, rel_path).await {
457 Ok(()) => recovered += 1,
458 Err(error) => {
459 tracing::warn!("index rebuild: failed to index root {root_id}: {error}")
460 }
461 }
462 }
463
464 let children_dir = root_entry.path().join("children");
466 let mut child_dirs = match fs::read_dir(&children_dir).await {
467 Ok(rd) => rd,
468 Err(error) if error.kind() == io::ErrorKind::NotFound => continue,
469 Err(error) => {
470 tracing::warn!("index rebuild: cannot read children of {root_id}: {error}");
471 continue;
472 }
473 };
474 loop {
475 let child_entry = match child_dirs.next_entry().await {
476 Ok(Some(entry)) => entry,
477 Ok(None) => break,
478 Err(error) => {
479 tracing::warn!(
480 "index rebuild: error scanning children of {root_id}: {error}"
481 );
482 break;
483 }
484 };
485 if !child_entry
486 .file_type()
487 .await
488 .map(|t| t.is_dir())
489 .unwrap_or(false)
490 {
491 continue;
492 }
493 let Ok(child_id) = child_entry.file_name().into_string() else {
494 continue;
495 };
496 if let Some(session) =
497 Self::load_session_from_dir(&child_entry.path(), &child_id).await
498 {
499 let rel_path = Self::child_rel_path(&root_id, &child_id);
500 match self.upsert_index_from_session(&session, rel_path).await {
501 Ok(()) => recovered += 1,
502 Err(error) => tracing::warn!(
503 "index rebuild: failed to index child {child_id}: {error}"
504 ),
505 }
506 }
507 }
508 }
509
510 self.update_index(|index| {
514 index.version = SESSIONS_INDEX_VERSION;
517 Ok(())
518 })
519 .await?;
520
521 tracing::info!("index rebuild from disk complete: recovered {recovered} session(s)");
522
523 if let Err(error) = self.rebuild_search_index().await {
525 tracing::warn!("index rebuild: failed to rebuild search index: {error}");
526 }
527 Ok(())
528 }
529
530 async fn load_session_from_dir(abs_dir: &Path, id: &str) -> Option<Session> {
540 let raw = match fs::read_to_string(abs_dir.join("session.json")).await {
541 Ok(raw) => raw,
542 Err(error) if error.kind() == io::ErrorKind::NotFound => return None,
543 Err(error) => {
544 tracing::warn!("index rebuild: skipping unreadable session {id}: {error}");
545 return None;
546 }
547 };
548 let main: Session = match serde_json::from_str(&raw) {
549 Ok(session) => session,
550 Err(error) => {
551 tracing::warn!("index rebuild: skipping corrupt session {id}: {error}");
552 return None;
553 }
554 };
555 let sidecar =
556 match Self::read_runtime_sidecar_at(&abs_dir.join(RUNTIME_SIDECAR_FILE), id).await {
557 Ok(sidecar) => sidecar,
558 Err(error) => {
559 tracing::warn!("index rebuild: cannot read runtime sidecar for {id}: {error}");
560 None
561 }
562 };
563 let mut session = overlay_runtime_sidecar(main, sidecar);
564 session.clear_stale_root_token_budget();
565 Some(session)
566 }
567
568 pub fn search_index(&self) -> &SessionSearchIndex {
569 &self.search_index
570 }
571
572 pub fn bamboo_home_dir(&self) -> &Path {
573 &self.bamboo_home_dir
574 }
575
576 async fn open_session_lifecycle_file(&self, exclusive: bool) -> io::Result<std::fs::File> {
577 let path = self.bamboo_home_dir.join(SESSION_LIFECYCLE_LOCK_FILE);
578 tokio::task::spawn_blocking(move || {
579 let file = std::fs::OpenOptions::new()
580 .create(true)
581 .truncate(false)
582 .read(true)
583 .write(true)
584 .open(&path)?;
585 if exclusive {
586 FileExt::lock_exclusive(&file)?;
587 } else {
588 FileExt::lock_shared(&file)?;
589 }
590 Ok(file)
591 })
592 .await
593 .map_err(|error| other_io_error(format!("join session lifecycle lock task: {error}")))?
594 }
595
596 pub(crate) async fn lock_session_lifecycle_shared(
597 &self,
598 ) -> io::Result<SessionLifecycleReadGuard> {
599 let process = self.session_lifecycle_lock.clone().read_owned().await;
600 let file = self.open_session_lifecycle_file(false).await?;
601 Ok(SessionLifecycleReadGuard {
602 _process: process,
603 file,
604 })
605 }
606
607 async fn lock_session_lifecycle_exclusive(&self) -> io::Result<SessionLifecycleWriteGuard> {
608 let process = self.session_lifecycle_lock.clone().write_owned().await;
609 let file = self.open_session_lifecycle_file(true).await?;
610 Ok(SessionLifecycleWriteGuard {
611 _process: process,
612 file,
613 })
614 }
615
616 pub fn index_path(&self) -> &Path {
617 &self.index_path
618 }
619
620 pub async fn rebuild_search_index(&self) -> io::Result<()> {
621 let session_ids = {
622 let index = self.index.read().await;
623 index.sessions.keys().cloned().collect::<Vec<_>>()
624 };
625 for session_id in session_ids {
626 if let Some(session) = self.load_session(&session_id).await? {
627 if !should_index_session(session.updated_at) {
628 continue;
629 }
630 if let Err(error) = self.search_index.upsert_session(&session).await {
631 tracing::warn!(
632 "failed to rebuild search index entry for {}: {}",
633 session_id,
634 error
635 );
636 }
637 }
638 }
639 Ok(())
640 }
641
642 pub fn sessions_root_dir(&self) -> &Path {
643 &self.sessions_dir
644 }
645
646 fn root_rel_path(session_id: &str) -> String {
647 format!("sessions/{session_id}")
648 }
649
650 fn child_rel_path(root_id: &str, child_id: &str) -> String {
651 format!("sessions/{root_id}/children/{child_id}")
652 }
653
654 fn abs_path_from_rel(&self, rel: &str) -> PathBuf {
655 self.bamboo_home_dir.join(rel)
656 }
657
658 async fn persist_index_locked(&self, index: &SessionsIndex) -> io::Result<()> {
659 let tmp = self
660 .index_path
661 .with_extension(format!("json.tmp.{}", Uuid::new_v4()));
662 let bytes = serde_json::to_vec_pretty(index).map_err(|e| other_io_error(e.to_string()))?;
663 fs::write(&tmp, bytes).await?;
664 atomic_rename(&tmp, &self.index_path).await?;
665 Ok(())
666 }
667
668 async fn update_index<F, T>(&self, f: F) -> io::Result<T>
669 where
670 F: FnOnce(&mut SessionsIndex) -> io::Result<T>,
671 {
672 let _guard = self.write_lock.lock().await;
673 let mut index = self.index.write().await;
674 let out = f(&mut index)?;
675 index.updated_at = Utc::now();
676 self.persist_index_locked(&index).await?;
677 Ok(out)
678 }
679
680 pub async fn list_index_entries(&self) -> Vec<SessionIndexEntry> {
681 let index = self.index.read().await;
682 let mut items: Vec<_> = index.sessions.values().cloned().collect();
683 items.sort_by_key(|b| std::cmp::Reverse(b.updated_at));
684 items
685 }
686
687 pub async fn get_index_entry(&self, session_id: &str) -> Option<SessionIndexEntry> {
688 let index = self.index.read().await;
689 index.sessions.get(session_id).cloned()
690 }
691
692 pub async fn resolve_rel_path(&self, session_id: &str) -> Option<String> {
693 self.get_index_entry(session_id).await.map(|e| e.rel_path)
694 }
695
696 async fn ensure_session_dirs(&self, session: &Session) -> io::Result<String> {
697 validate_session_id(&session.id)?;
698
699 let rel_path = match session.kind {
700 SessionKind::Root => Self::root_rel_path(&session.id),
701 SessionKind::Child => {
702 let root_id = session.root_session_id.trim();
703 let parent_id = session.parent_session_id.as_deref().unwrap_or("").trim();
704 if root_id.is_empty() || parent_id.is_empty() {
705 return Err(other_io_error(
706 "child session missing root_session_id/parent_session_id",
707 ));
708 }
709 validate_session_id(root_id)?;
714 Self::child_rel_path(root_id, &session.id)
715 }
716 };
717
718 let abs_dir = self.abs_path_from_rel(&rel_path);
719 fs::create_dir_all(&abs_dir).await?;
720 fs::create_dir_all(abs_dir.join("attachments")).await?;
722 if session.kind == SessionKind::Root {
723 fs::create_dir_all(abs_dir.join("children")).await?;
724 }
725 Ok(rel_path)
726 }
727
728 async fn session_json_path(&self, session_id: &str) -> io::Result<Option<PathBuf>> {
729 if let Some(rel) = self.resolve_rel_path(session_id).await {
730 Ok(Some(self.abs_path_from_rel(&rel).join("session.json")))
731 } else {
732 Ok(None)
733 }
734 }
735
736 async fn runtime_json_path(&self, session_id: &str) -> io::Result<Option<PathBuf>> {
737 if let Some(rel) = self.resolve_rel_path(session_id).await {
738 Ok(Some(
739 self.abs_path_from_rel(&rel).join(RUNTIME_SIDECAR_FILE),
740 ))
741 } else {
742 Ok(None)
743 }
744 }
745
746 async fn write_runtime_sidecar(&self, abs_dir: &Path, session: &Session) -> io::Result<()> {
750 let path = abs_dir.join(RUNTIME_SIDECAR_FILE);
751 let snapshot = runtime_sidecar_snapshot(session);
752 let tmp = path.with_extension(format!("json.tmp.{}", Uuid::new_v4()));
753 let bytes =
754 serde_json::to_vec_pretty(&snapshot).map_err(|e| other_io_error(e.to_string()))?;
755 fs::write(&tmp, bytes).await?;
756 atomic_rename(&tmp, &path).await?;
757 Ok(())
758 }
759
760 pub async fn migrate_runtime_sidecars(&self) -> io::Result<usize> {
773 let marker = self.bamboo_home_dir.join(RUNTIME_SIDECAR_MIGRATION_MARKER);
774 if fs::try_exists(&marker).await.unwrap_or(false) {
775 return Ok(0);
776 }
777
778 let entries = self.list_index_entries().await;
779 let mut migrated = 0usize;
780 for entry in entries {
781 let abs_dir = self.abs_path_from_rel(&entry.rel_path);
782 let sidecar_path = abs_dir.join(RUNTIME_SIDECAR_FILE);
783 if fs::try_exists(&sidecar_path).await.unwrap_or(false) {
784 continue;
785 }
786 let session_path = abs_dir.join("session.json");
787 let raw = match fs::read_to_string(&session_path).await {
790 Ok(raw) => raw,
791 Err(error) if error.kind() == io::ErrorKind::NotFound => continue,
792 Err(error) => return Err(error),
793 };
794 let session: Session = match serde_json::from_str(&raw) {
795 Ok(session) => session,
796 Err(error) => {
797 tracing::warn!(
798 "runtime sidecar migration: skipping unreadable session {}: {}",
799 entry.id,
800 error
801 );
802 continue;
803 }
804 };
805 self.write_runtime_sidecar(&abs_dir, &session).await?;
806 migrated += 1;
807 }
808
809 let tmp = marker.with_extension(format!("tmp.{}", Uuid::new_v4()));
812 fs::write(&tmp, b"runtime-sidecar-v1\n").await?;
813 atomic_rename(&tmp, &marker).await?;
814
815 if migrated > 0 {
816 tracing::info!("runtime sidecar migration: created {migrated} sidecar(s)");
817 }
818 Ok(migrated)
819 }
820
821 async fn read_runtime_sidecar(&self, session_id: &str) -> io::Result<Option<Session>> {
825 let Some(path) = self.runtime_json_path(session_id).await? else {
826 return Ok(None);
827 };
828 Self::read_runtime_sidecar_at(&path, session_id).await
829 }
830
831 async fn read_runtime_sidecar_at(path: &Path, id: &str) -> io::Result<Option<Session>> {
837 if !path.exists() {
838 return Ok(None);
839 }
840 let raw = fs::read_to_string(path).await?;
841 match serde_json::from_str::<Session>(&raw) {
842 Ok(mut side) => {
843 side.clear_stale_root_token_budget();
846 Ok(Some(side))
847 }
848 Err(error) => {
849 tracing::warn!("ignoring corrupt runtime sidecar for {id}: {error}");
852 Ok(None)
853 }
854 }
855 }
856
857 async fn attachments_dir(&self, session_id: &str) -> io::Result<Option<PathBuf>> {
858 if let Some(rel) = self.resolve_rel_path(session_id).await {
859 Ok(Some(self.abs_path_from_rel(&rel).join("attachments")))
860 } else {
861 Ok(None)
862 }
863 }
864
865 async fn compute_has_attachments(&self, session_id: &str) -> bool {
866 let Ok(Some(dir)) = self.attachments_dir(session_id).await else {
867 return false;
868 };
869 let Ok(mut rd) = fs::read_dir(dir).await else {
870 return false;
871 };
872 rd.next_entry().await.ok().flatten().is_some()
873 }
874
875 async fn upsert_index_from_session(
876 &self,
877 session: &Session,
878 rel_path: String,
879 ) -> io::Result<()> {
880 let has_attachments = self.compute_has_attachments(&session.id).await;
881 let last_run_status = session
884 .last_run_status()
885 .filter(|value| !value.trim().is_empty());
886 let last_run_error = session
887 .last_run_error()
888 .filter(|value| !value.trim().is_empty());
889 let created_by_schedule_id = session
890 .metadata
891 .get("created_by_schedule_id")
892 .cloned()
893 .filter(|v| !v.trim().is_empty());
894 let schedule_run_id = session
895 .metadata
896 .get("schedule_run_id")
897 .cloned()
898 .filter(|v| !v.trim().is_empty());
899 let subagent_type = session.subagent_type().filter(|v| !v.trim().is_empty());
900 let lifecycle = session
901 .metadata
902 .get("lifecycle")
903 .cloned()
904 .filter(|v| !v.trim().is_empty());
905 let resident_name = session
906 .metadata
907 .get("resident_name")
908 .cloned()
909 .filter(|v| !v.trim().is_empty());
910 let gold_config_json = session
911 .metadata
912 .get("gold_config")
913 .cloned()
914 .filter(|v| !v.trim().is_empty());
915 let plan_mode = session
916 .agent_runtime_state
917 .as_ref()
918 .and_then(|state| state.plan_mode.clone());
919 let permission_mode = session
920 .agent_runtime_state
921 .as_ref()
922 .map(|state| state.effective_permission_mode())
923 .unwrap_or_default();
924 let bypass_permissions = permission_mode != bamboo_domain::SessionPermissionMode::Default;
925 let placement = session
930 .metadata
931 .get("placement")
932 .and_then(|v| serde_json::from_str::<SessionPlacement>(v).ok());
933 let workspace_path = session
934 .workspace_path_meta()
935 .map(|value| value.trim().to_string())
936 .filter(|value| !value.is_empty());
937 let project_id = normalized_project_id(session);
938 self.update_index(|index| {
939 index.sessions.insert(
940 session.id.clone(),
941 SessionIndexEntry {
942 id: session.id.clone(),
943 kind: session.kind,
944 rel_path,
945 title: session.title.clone(),
946 title_version: session.title_version,
947 pinned: session.pinned,
948 parent_session_id: session.parent_session_id.clone(),
949 root_session_id: session.root_session_id.clone(),
950 spawn_depth: session.spawn_depth,
951 model: session.model.clone(),
952 model_ref: session.model_ref.clone(),
953 reasoning_effort: session.reasoning_effort,
954 workspace_path,
955 project_id,
956 gold_config_json,
957 created_by_schedule_id,
958 schedule_run_id,
959 created_at: session.created_at,
960 updated_at: session.updated_at,
961 last_activity_at: session.updated_at,
962 message_count: session.messages.len(),
963 has_attachments,
964 has_pending_question: session.has_pending_question(),
965 plan_mode,
966 bypass_permissions,
967 permission_mode,
968 last_run_status,
969 last_run_error,
970 token_usage: session.token_usage.clone(),
971 subagent_type,
972 lifecycle,
973 resident_name,
974 placement,
975 },
976 );
977 Ok(())
978 })
979 .await?;
980 Ok(())
981 }
982
983 pub async fn write_image_attachment(
984 &self,
985 session: &Session,
986 raw_base64_or_data_url: &str,
987 mime_hint: Option<&str>,
988 ) -> io::Result<(String, String)> {
989 let (mime, base64_data) =
990 parse_data_url_base64(raw_base64_or_data_url).unwrap_or_else(|| {
991 (
992 mime_hint.unwrap_or("image/png").trim().to_string(),
993 raw_base64_or_data_url.trim().to_string(),
994 )
995 });
996
997 let bytes = base64::engine::general_purpose::STANDARD
998 .decode(base64_data.as_bytes())
999 .map_err(|e| other_io_error(format!("invalid base64 image data: {e}")))?;
1000
1001 let attachment_id = Uuid::new_v4().to_string();
1002 let ext = mime_to_extension(mime.as_str()).unwrap_or("bin");
1003
1004 let rel_path = self.ensure_session_dirs(session).await?;
1005 let abs_dir = self.abs_path_from_rel(&rel_path);
1006 let attachments_dir = abs_dir.join("attachments");
1007 fs::create_dir_all(&attachments_dir).await?;
1008
1009 let path = attachments_dir.join(format!("{attachment_id}.{ext}"));
1010 let tmp = path.with_extension(format!("{ext}.tmp.{}", Uuid::new_v4()));
1011 fs::write(&tmp, &bytes).await?;
1012 atomic_rename(&tmp, &path).await?;
1013
1014 Ok((
1015 attachment_id.clone(),
1016 format!("bamboo-attachment://{}/{}", session.id, attachment_id),
1017 ))
1018 }
1019
1020 pub async fn read_attachment(
1022 &self,
1023 session_id: &str,
1024 attachment_id: &str,
1025 ) -> io::Result<Option<(Vec<u8>, String)>> {
1026 validate_session_id(session_id)?;
1027 validate_session_id(attachment_id)?;
1028 let Some(dir) = self.attachments_dir(session_id).await? else {
1029 return Ok(None);
1030 };
1031 if !dir.exists() {
1032 return Ok(None);
1033 }
1034
1035 let mut rd = fs::read_dir(&dir).await?;
1036 while let Some(entry) = rd.next_entry().await? {
1037 let file_name = entry.file_name();
1038 let file_name = file_name.to_string_lossy();
1039 if !file_name.starts_with(attachment_id) {
1040 continue;
1041 }
1042 if file_name.len() <= attachment_id.len() + 1
1044 || !file_name.as_bytes()[attachment_id.len()].eq(&b'.')
1045 {
1046 continue;
1047 }
1048 let ext = file_name.split('.').next_back().unwrap_or("bin");
1049 let mime = extension_to_mime(ext).unwrap_or("application/octet-stream");
1050 let bytes = fs::read(entry.path()).await?;
1051 return Ok(Some((bytes, mime.to_string())));
1052 }
1053
1054 Ok(None)
1055 }
1056
1057 pub async fn clear_session(&self, session_id: &str) -> io::Result<bool> {
1058 let Some(mut session) = self.load_session(session_id).await? else {
1059 return Ok(false);
1060 };
1061
1062 let system_msg = session
1064 .messages
1065 .iter()
1066 .find(|m| matches!(m.role, Role::System))
1067 .cloned();
1068 session.messages.clear();
1069 if let Some(system) = system_msg {
1070 session.messages.push(system);
1071 }
1072
1073 session.token_usage = None;
1075 session.conversation_summary = None;
1076 session.updated_at = Utc::now();
1077
1078 if let Ok(Some(dir)) = self.attachments_dir(session_id).await {
1080 let _ = fs::remove_dir_all(&dir).await;
1081 let _ = fs::create_dir_all(&dir).await;
1082 }
1083
1084 self.save_session(&session).await?;
1085 Ok(true)
1086 }
1087
1088 pub async fn cleanup(&self, mode: CleanupMode, keep_pinned: bool) -> io::Result<CleanupResult> {
1089 let entries = {
1091 self.index
1092 .read()
1093 .await
1094 .sessions
1095 .values()
1096 .cloned()
1097 .collect::<Vec<_>>()
1098 };
1099
1100 let pinned_child_roots: HashSet<String> = if keep_pinned {
1101 entries
1102 .iter()
1103 .filter(|e| e.kind == SessionKind::Child && e.pinned)
1104 .filter_map(|e| e.parent_session_id.clone())
1105 .collect()
1106 } else {
1107 HashSet::new()
1108 };
1109
1110 let is_protected = |e: &SessionIndexEntry| -> bool {
1112 if !keep_pinned {
1113 return false;
1114 }
1115 if e.pinned {
1116 return true;
1117 }
1118 if e.kind == SessionKind::Root && pinned_child_roots.contains(&e.id) {
1120 return true;
1121 }
1122 false
1123 };
1124
1125 let mut delete_child_ids = HashSet::<String>::new();
1127 let mut delete_root_ids = HashSet::<String>::new();
1128
1129 match mode {
1130 CleanupMode::Children => {
1131 for e in entries.iter().filter(|e| e.kind == SessionKind::Child) {
1132 if is_protected(e) {
1133 continue;
1134 }
1135 delete_child_ids.insert(e.id.clone());
1136 }
1137 }
1138 CleanupMode::All | CleanupMode::Empty => {
1139 for root in entries.iter().filter(|e| e.kind == SessionKind::Root) {
1141 if is_protected(root) {
1142 continue;
1143 }
1144 if mode == CleanupMode::Empty && root.message_count > 1 {
1145 continue;
1146 }
1147 delete_root_ids.insert(root.id.clone());
1148 }
1149
1150 for child in entries.iter().filter(|e| e.kind == SessionKind::Child) {
1152 if delete_root_ids.contains(&child.root_session_id) {
1153 continue; }
1155 if is_protected(child) {
1156 continue;
1157 }
1158 if mode == CleanupMode::Empty && child.message_count > 1 {
1159 continue;
1160 }
1161 delete_child_ids.insert(child.id.clone());
1162 }
1163 }
1164 }
1165
1166 let mut deleted_ids = HashSet::<String>::new();
1168 for root_id in delete_root_ids.iter() {
1169 for e in entries.iter().filter(|e| e.root_session_id == *root_id) {
1170 deleted_ids.insert(e.id.clone());
1171 }
1172 }
1173 for child_id in delete_child_ids.iter() {
1174 deleted_ids.insert(child_id.clone());
1175 }
1176
1177 for root_id in delete_root_ids.iter() {
1179 let _ = self.delete_session_recursive(root_id, true).await?;
1180 }
1181 for child_id in delete_child_ids.iter() {
1182 let _ = self.delete_session_recursive(child_id, true).await?;
1183 }
1184 let mut deleted_session_ids: Vec<String> = deleted_ids.into_iter().collect();
1185 deleted_session_ids.sort();
1186 Ok(CleanupResult {
1187 deleted_count: deleted_session_ids.len(),
1188 deleted_session_ids,
1189 })
1190 }
1191
1192 pub async fn dev_reset(&self) -> io::Result<()> {
1198 let _lifecycle = self.lock_session_lifecycle_exclusive().await?;
1199 let _guard = self.write_lock.lock().await;
1200
1201 let _ = fs::remove_dir_all(&self.sessions_dir).await;
1203 fs::create_dir_all(&self.sessions_dir).await?;
1204
1205 {
1207 let mut index = self.index.write().await;
1208 *index = SessionsIndex::empty();
1209 self.persist_index_locked(&index).await?;
1210 }
1211
1212 Ok(())
1213 }
1214
1215 pub async fn delete_session_recursive(
1220 &self,
1221 session_id: &str,
1222 force: bool,
1223 ) -> io::Result<bool> {
1224 let _lifecycle = self.lock_session_lifecycle_exclusive().await?;
1225 self.delete_session_recursive_locked(session_id, force)
1226 .await
1227 }
1228
1229 async fn delete_session_recursive_locked(
1230 &self,
1231 session_id: &str,
1232 force: bool,
1233 ) -> io::Result<bool> {
1234 let entry = self.get_index_entry(session_id).await;
1235 let Some(entry) = entry else {
1236 return Ok(false);
1237 };
1238
1239 if !force && entry.pinned {
1240 return Err(other_io_error(
1241 "refusing to delete pinned session without force",
1242 ));
1243 }
1244
1245 match entry.kind {
1246 SessionKind::Child => {
1247 let abs_dir = self.abs_path_from_rel(&entry.rel_path);
1248 let _ = fs::remove_dir_all(&abs_dir).await;
1249 self.update_index(|index| {
1250 index.sessions.remove(session_id);
1251 Ok(())
1252 })
1253 .await?;
1254 if let Err(error) = self.search_index.delete_session(session_id).await {
1255 tracing::warn!(
1256 "failed to delete session search index row for {}: {}",
1257 session_id,
1258 error
1259 );
1260 }
1261 Ok(true)
1262 }
1263 SessionKind::Root => {
1264 let root_id = entry.id.clone();
1265 let abs_dir = self.abs_path_from_rel(&entry.rel_path);
1266 let _ = fs::remove_dir_all(&abs_dir).await;
1267
1268 let to_remove_ids = {
1269 let index = self.index.read().await;
1270 index
1271 .sessions
1272 .values()
1273 .filter(|e| e.root_session_id == root_id)
1274 .map(|e| e.id.clone())
1275 .collect::<Vec<_>>()
1276 };
1277
1278 self.update_index(|index| {
1279 for id in &to_remove_ids {
1280 index.sessions.remove(id);
1281 }
1282 Ok(())
1283 })
1284 .await?;
1285
1286 for id in to_remove_ids {
1287 if let Err(error) = self.search_index.delete_session(&id).await {
1288 tracing::warn!(
1289 "failed to delete session search index row for {}: {}",
1290 id,
1291 error
1292 );
1293 }
1294 }
1295 Ok(true)
1296 }
1297 }
1298 }
1299}
1300
1301#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1302pub enum CleanupMode {
1303 All,
1304 Empty,
1305 Children,
1306}
1307
1308#[derive(Debug, Clone, Serialize, Deserialize)]
1309pub struct CleanupResult {
1310 pub deleted_count: usize,
1311 pub deleted_session_ids: Vec<String>,
1312}
1313
1314pub(crate) async fn atomic_write(path: &Path, bytes: &[u8]) -> io::Result<()> {
1328 use tokio::io::AsyncWriteExt;
1329
1330 let tmp = path.with_extension(format!("tmp.{}", Uuid::new_v4()));
1331 let write_result = async {
1332 let mut file = fs::File::create(&tmp).await?;
1333 file.write_all(bytes).await?;
1334 file.sync_all().await
1336 }
1337 .await;
1338 if let Err(e) = write_result {
1339 let _ = fs::remove_file(&tmp).await;
1340 return Err(e);
1341 }
1342 atomic_rename(&tmp, path).await
1343}
1344
1345async fn atomic_rename(from: &Path, to: &Path) -> io::Result<()> {
1346 match fs::rename(from, to).await {
1348 Ok(()) => Ok(()),
1349 Err(err) => {
1350 if to.exists() {
1351 let _ = fs::remove_file(to).await;
1352 }
1353 fs::rename(from, to).await.map_err(|e| {
1354 other_io_error(format!(
1355 "failed to rename {:?} -> {:?}: {} (original: {})",
1356 from, to, e, err
1357 ))
1358 })
1359 }
1360 }
1361}
1362
1363fn parse_data_url_base64(url: &str) -> Option<(String, String)> {
1364 let trimmed = url.trim();
1366 if !trimmed.starts_with("data:") {
1367 return None;
1368 }
1369 let trimmed = trimmed.strip_prefix("data:")?;
1370 let (header, data) = trimmed.split_once(',')?;
1371 if !header.contains(";base64") {
1372 return None;
1373 }
1374 let mime = header.split(';').next()?.trim().to_string();
1375 Some((mime, data.trim().to_string()))
1376}
1377
1378fn mime_to_extension(mime: &str) -> Option<&'static str> {
1379 match mime.trim().to_ascii_lowercase().as_str() {
1380 "image/png" => Some("png"),
1381 "image/jpeg" => Some("jpg"),
1382 "image/webp" => Some("webp"),
1383 "image/gif" => Some("gif"),
1384 "image/bmp" => Some("bmp"),
1385 _ => None,
1386 }
1387}
1388
1389fn extension_to_mime(ext: &str) -> Option<&'static str> {
1390 match ext.trim().to_ascii_lowercase().as_str() {
1391 "png" => Some("image/png"),
1392 "jpg" | "jpeg" => Some("image/jpeg"),
1393 "webp" => Some("image/webp"),
1394 "gif" => Some("image/gif"),
1395 "bmp" => Some("image/bmp"),
1396 _ => None,
1397 }
1398}
1399
1400#[async_trait::async_trait]
1401impl Storage for SessionStoreV2 {
1402 async fn save_session(&self, session: &Session) -> io::Result<()> {
1403 let rel_path = self.ensure_session_dirs(session).await?;
1404 let abs_dir = self.abs_path_from_rel(&rel_path);
1405 let path = abs_dir.join("session.json");
1406
1407 self.write_runtime_sidecar(&abs_dir, session).await?;
1414
1415 let tmp = path.with_extension(format!("json.tmp.{}", Uuid::new_v4()));
1416 let bytes =
1417 serde_json::to_vec_pretty(session).map_err(|e| other_io_error(e.to_string()))?;
1418 fs::write(&tmp, bytes).await?;
1419 atomic_rename(&tmp, &path).await?;
1420
1421 self.upsert_index_from_session(session, rel_path).await?;
1422 if let Err(error) = self.search_index.upsert_session(session).await {
1423 tracing::warn!(
1424 "failed to update session search index for {}: {}",
1425 session.id,
1426 error
1427 );
1428 }
1429 Ok(())
1430 }
1431
1432 async fn load_session(&self, session_id: &str) -> io::Result<Option<Session>> {
1433 validate_session_id(session_id)?;
1434 let Some(path) = self.session_json_path(session_id).await? else {
1435 return Ok(None);
1436 };
1437 if !path.exists() {
1438 return Ok(None);
1439 }
1440 let raw = fs::read_to_string(path).await?;
1441 let session: Session = serde_json::from_str(&raw)
1442 .map_err(|e| other_io_error(format!("invalid session.json: {e}")))?;
1443 let sidecar = self.read_runtime_sidecar(session_id).await?;
1444 let mut session = overlay_runtime_sidecar(session, sidecar);
1445 session.clear_stale_root_token_budget();
1447 Ok(Some(session))
1448 }
1449
1450 async fn delete_session(&self, session_id: &str) -> io::Result<bool> {
1451 self.delete_session_recursive(session_id, true).await
1453 }
1454
1455 async fn save_runtime_state(&self, session: &Session) -> io::Result<()> {
1456 let Some(rel) = self.resolve_rel_path(&session.id).await else {
1460 return self.save_session(session).await;
1463 };
1464 let abs_dir = self.abs_path_from_rel(&rel);
1465 self.write_runtime_sidecar(&abs_dir, session).await?;
1466
1467 let workspace_path = session
1472 .workspace_path_meta()
1473 .map(|value| value.trim().to_string())
1474 .filter(|value| !value.is_empty());
1475 let project_id = normalized_project_id(session);
1476 let runtime_index_changed = self
1477 .get_index_entry(&session.id)
1478 .await
1479 .is_some_and(|entry| {
1480 entry.workspace_path != workspace_path || entry.project_id != project_id
1481 });
1482 if runtime_index_changed {
1483 self.update_index(|index| {
1484 if let Some(entry) = index.sessions.get_mut(&session.id) {
1485 entry.workspace_path = workspace_path;
1486 entry.project_id = project_id;
1487 }
1488 Ok(())
1489 })
1490 .await?;
1491 }
1492 Ok(())
1493 }
1494
1495 async fn load_runtime_control_plane(&self, session_id: &str) -> io::Result<Option<Session>> {
1496 validate_session_id(session_id)?;
1497 if let Some(side) = self.read_runtime_sidecar(session_id).await? {
1500 return Ok(Some(side));
1501 }
1502 self.load_session(session_id).await
1503 }
1504
1505 async fn list_child_run_statuses(
1506 &self,
1507 parent_session_id: &str,
1508 ) -> io::Result<Vec<(String, Option<String>)>> {
1509 let index = self.index.read().await;
1510 Ok(index
1511 .sessions
1512 .values()
1513 .filter(|entry| {
1514 entry.kind == SessionKind::Child
1515 && entry.parent_session_id.as_deref() == Some(parent_session_id)
1516 })
1517 .map(|entry| (entry.id.clone(), entry.last_run_status.clone()))
1518 .collect())
1519 }
1520
1521 async fn list_sessions_by_run_status(
1522 &self,
1523 status: &str,
1524 ) -> io::Result<Vec<(String, Option<String>)>> {
1525 let index = self.index.read().await;
1526 Ok(index
1527 .sessions
1528 .values()
1529 .filter(|entry| entry.last_run_status.as_deref() == Some(status))
1530 .map(|entry| (entry.id.clone(), entry.parent_session_id.clone()))
1531 .collect())
1532 }
1533
1534 async fn append_token_usage_record(&self, session_id: &str, json_line: &str) -> io::Result<()> {
1535 use tokio::io::AsyncWriteExt;
1536
1537 validate_session_id(session_id)?;
1538 let Some(rel) = self.resolve_rel_path(session_id).await else {
1542 return Ok(());
1543 };
1544 let path = self.abs_path_from_rel(&rel).join(TOKEN_USAGE_FILE);
1545
1546 let mut line = json_line.trim_end_matches('\n').to_string();
1548 line.push('\n');
1549
1550 let mut file = fs::OpenOptions::new()
1551 .create(true)
1552 .append(true)
1553 .open(&path)
1554 .await?;
1555 file.write_all(line.as_bytes()).await?;
1556 file.flush().await?;
1570 Ok(())
1571 }
1572}
1573
1574#[async_trait::async_trait]
1575impl AttachmentReader for SessionStoreV2 {
1576 async fn read_attachment(
1577 &self,
1578 session_id: &str,
1579 attachment_id: &str,
1580 ) -> io::Result<Option<(Vec<u8>, String)>> {
1581 SessionStoreV2::read_attachment(self, session_id, attachment_id).await
1582 }
1583}
1584
1585#[cfg(test)]
1586mod tests {
1587 use super::*;
1588 use bamboo_domain::{SessionInboxError, SessionInboxPort, SessionMessageEnvelope};
1589 use std::io;
1590 use std::sync::Arc;
1591 use tempfile::TempDir;
1592
1593 async fn create_temp_storage() -> io::Result<(SessionStoreV2, TempDir)> {
1594 let temp_dir = TempDir::new().map_err(io::Error::other)?;
1595 let bamboo_home = temp_dir.path().to_path_buf();
1596 let storage = SessionStoreV2::new(bamboo_home).await?;
1597 Ok((storage, temp_dir))
1598 }
1599
1600 #[tokio::test]
1601 async fn test_new_creates_sessions_directory() -> io::Result<()> {
1602 let temp_dir = TempDir::new().map_err(io::Error::other)?;
1603 let bamboo_home = temp_dir.path().to_path_buf();
1604 let sessions_dir = bamboo_home.join("sessions");
1605
1606 assert!(!sessions_dir.exists());
1607 let _storage = SessionStoreV2::new(bamboo_home).await?;
1608 assert!(sessions_dir.exists());
1609
1610 Ok(())
1611 }
1612
1613 #[tokio::test]
1614 async fn root_token_budget_cleared_on_load_child_preserved() -> io::Result<()> {
1615 let (storage, _dir) = create_temp_storage().await?;
1619
1620 let mut root = Session::new("root-1", "m");
1621 root.token_budget = Some(bamboo_domain::TokenBudget::for_model(1000));
1622 storage.save_session(&root).await?;
1623 let loaded = storage.load_session("root-1").await?.expect("root present");
1624 assert!(
1625 loaded.token_budget.is_none(),
1626 "stale Root token_budget must be cleared on load"
1627 );
1628
1629 let parent = Session::new("root-1", "m");
1630 let mut child = Session::new_child_of("child-1", &parent, "m", "c");
1631 child.token_budget = Some(bamboo_domain::TokenBudget::for_model(500));
1632 storage.save_session(&child).await?;
1633 let loaded_child = storage
1634 .load_session("child-1")
1635 .await?
1636 .expect("child present");
1637 assert!(
1638 loaded_child.token_budget.is_some(),
1639 "Child assigned sub-budget must be preserved on load"
1640 );
1641
1642 Ok(())
1643 }
1644
1645 #[tokio::test]
1646 async fn test_new_creates_index_file() -> io::Result<()> {
1647 let temp_dir = TempDir::new().map_err(io::Error::other)?;
1648 let bamboo_home = temp_dir.path().to_path_buf();
1649 let index_path = bamboo_home.join("sessions.json");
1650
1651 assert!(!index_path.exists());
1652 let _storage = SessionStoreV2::new(bamboo_home).await?;
1653 assert!(index_path.exists());
1654
1655 Ok(())
1656 }
1657
1658 use bamboo_domain::session::types::Message;
1661 use bamboo_domain::AgentRuntimeState;
1662
1663 fn session_with_history(id: &str, messages: usize, run_id: &str) -> Session {
1664 let mut s = Session::new(id.to_string(), "test-model".to_string());
1665 for i in 0..messages {
1666 s.add_message(Message::user(format!("msg-{i}")));
1667 }
1668 s.agent_runtime_state = Some(AgentRuntimeState::new(run_id));
1669 s
1670 }
1671
1672 async fn read_session_json_raw(storage: &SessionStoreV2, id: &str) -> String {
1673 let path = storage.session_json_path(id).await.unwrap().unwrap();
1674 tokio::fs::read_to_string(path).await.unwrap()
1675 }
1676
1677 #[tokio::test]
1678 async fn append_token_usage_record_writes_jsonl_in_session_dir() -> io::Result<()> {
1679 let (storage, _t) = create_temp_storage().await?;
1680 let s = session_with_history("tu-1", 1, "run-A");
1681 storage.save_session(&s).await?;
1682
1683 storage
1684 .append_token_usage_record("tu-1", r#"{"round":1,"cache_read_input_tokens":0}"#)
1685 .await?;
1686 storage
1688 .append_token_usage_record("tu-1", "{\"round\":2,\"cache_read_input_tokens\":9000}\n")
1689 .await?;
1690
1691 let rel = storage.resolve_rel_path("tu-1").await.unwrap();
1692 let path = storage.abs_path_from_rel(&rel).join(TOKEN_USAGE_FILE);
1693 assert!(
1694 path.exists(),
1695 "token-usage.jsonl should sit in the session dir"
1696 );
1697
1698 let contents = tokio::fs::read_to_string(&path).await?;
1699 let lines: Vec<&str> = contents.lines().collect();
1700 assert_eq!(
1708 lines.len(),
1709 2,
1710 "one line per appended record; actual token-usage.jsonl = {contents:?}"
1711 );
1712 assert!(lines[0].contains("\"round\":1"));
1713 assert!(lines[1].contains("\"round\":2"));
1714 for line in lines {
1716 serde_json::from_str::<serde_json::Value>(line).expect("each line is valid JSON");
1717 }
1718 Ok(())
1719 }
1720
1721 #[tokio::test]
1722 async fn append_token_usage_record_is_noop_for_unindexed_session() -> io::Result<()> {
1723 let (storage, _t) = create_temp_storage().await?;
1724 storage
1726 .append_token_usage_record("never-saved", r#"{"round":1}"#)
1727 .await?;
1728 assert!(storage.resolve_rel_path("never-saved").await.is_none());
1729 Ok(())
1730 }
1731
1732 #[tokio::test]
1733 async fn save_session_writes_runtime_sidecar() -> io::Result<()> {
1734 let (storage, _t) = create_temp_storage().await?;
1735 let s = session_with_history("sc-1", 2, "run-A");
1736 storage.save_session(&s).await?;
1737
1738 let sidecar_path = storage.runtime_json_path("sc-1").await?.unwrap();
1739 assert!(
1740 sidecar_path.exists(),
1741 "save_session must write runtime.json"
1742 );
1743
1744 let side = storage.read_runtime_sidecar("sc-1").await?.unwrap();
1746 assert!(side.messages.is_empty(), "sidecar messages must be cleared");
1747 assert_eq!(side.agent_runtime_state.as_ref().unwrap().run_id, "run-A");
1748 Ok(())
1749 }
1750
1751 #[tokio::test]
1752 async fn save_runtime_state_does_not_rewrite_session_json_messages() -> io::Result<()> {
1753 let (storage, _t) = create_temp_storage().await?;
1754
1755 let s = session_with_history("sc-2", 3, "run-A");
1757 storage.save_session(&s).await?;
1758 let raw_before = read_session_json_raw(&storage, "sc-2").await;
1759 assert!(raw_before.contains("msg-2"));
1760
1761 let mut s2 = s.clone();
1765 s2.agent_runtime_state = Some(AgentRuntimeState::new("run-B"));
1766 s2.add_message(Message::user("msg-3-should-not-persist"));
1767 storage.save_runtime_state(&s2).await?;
1768
1769 let raw_after = read_session_json_raw(&storage, "sc-2").await;
1770 assert_eq!(
1771 raw_before, raw_after,
1772 "save_runtime_state must not touch session.json"
1773 );
1774
1775 let loaded = storage.load_session("sc-2").await?.unwrap();
1777 assert_eq!(loaded.agent_runtime_state.as_ref().unwrap().run_id, "run-B");
1778 assert_eq!(
1779 loaded.messages.len(),
1780 3,
1781 "runtime-only save must not add a message"
1782 );
1783 Ok(())
1784 }
1785
1786 #[tokio::test]
1787 async fn save_runtime_state_falls_back_to_full_save_when_unpersisted() -> io::Result<()> {
1788 let (storage, _t) = create_temp_storage().await?;
1789 let s = session_with_history("sc-3", 1, "run-A");
1792 storage.save_runtime_state(&s).await?;
1793
1794 let loaded = storage.load_session("sc-3").await?;
1795 assert!(
1796 loaded.is_some(),
1797 "fallback full save must create the session"
1798 );
1799 assert_eq!(loaded.unwrap().messages.len(), 1);
1800 Ok(())
1801 }
1802
1803 #[tokio::test]
1804 async fn corrupt_sidecar_is_ignored_and_session_still_loads() -> io::Result<()> {
1805 let (storage, _t) = create_temp_storage().await?;
1806 let s = session_with_history("sc-4", 2, "run-A");
1807 storage.save_session(&s).await?;
1808
1809 let sidecar_path = storage.runtime_json_path("sc-4").await?.unwrap();
1811 tokio::fs::write(&sidecar_path, b"{ not valid json").await?;
1812
1813 let loaded = storage.load_session("sc-4").await?.unwrap();
1815 assert_eq!(loaded.messages.len(), 2);
1816 assert_eq!(loaded.agent_runtime_state.as_ref().unwrap().run_id, "run-A");
1817 Ok(())
1818 }
1819
1820 #[tokio::test]
1821 async fn corrupt_index_is_backed_up_and_rebuilt_from_disk() -> io::Result<()> {
1822 let temp_dir = TempDir::new().map_err(io::Error::other)?;
1827 let bamboo_home = temp_dir.path().to_path_buf();
1828
1829 {
1831 let storage = SessionStoreV2::new(bamboo_home.clone()).await?;
1832 let root = Session::new("root-1", "m");
1833 storage.save_session(&root).await?;
1834 let child = Session::new_child_of("child-1", &root, "m", "c");
1835 storage.save_session(&child).await?;
1836 }
1837
1838 let index_path = bamboo_home.join("sessions.json");
1840 tokio::fs::write(&index_path, b"{ not valid json").await?;
1841
1842 let recovered = SessionStoreV2::new(bamboo_home.clone()).await?;
1844
1845 assert_eq!(
1848 recovered.resolve_rel_path("root-1").await.as_deref(),
1849 Some("sessions/root-1"),
1850 "root must be recovered with its on-disk rel_path"
1851 );
1852 assert_eq!(
1853 recovered.resolve_rel_path("child-1").await.as_deref(),
1854 Some("sessions/root-1/children/child-1"),
1855 "child must be recovered with its on-disk rel_path"
1856 );
1857 assert!(
1858 recovered.get_index_entry("root-1").await.is_some(),
1859 "root index entry must exist after rebuild"
1860 );
1861 assert!(
1862 recovered.load_session("root-1").await?.is_some(),
1863 "recovered root must load from disk"
1864 );
1865 let loaded_child = recovered
1866 .load_session("child-1")
1867 .await?
1868 .expect("recovered child must load from disk");
1869 assert_eq!(loaded_child.parent_session_id.as_deref(), Some("root-1"));
1870 assert_eq!(loaded_child.root_session_id, "root-1");
1871
1872 assert!(
1875 bamboo_home.join("sessions.json.bak").exists(),
1876 "corrupt sessions.json must be backed up to sessions.json.bak"
1877 );
1878 assert!(
1879 index_path.exists(),
1880 "a fresh sessions.json must be written after rebuild"
1881 );
1882
1883 Ok(())
1884 }
1885
1886 #[tokio::test]
1887 async fn v2_index_migrates_workspace_paths_from_root_and_child_sessions() -> io::Result<()> {
1888 let temp_dir = TempDir::new().map_err(io::Error::other)?;
1889 let bamboo_home = temp_dir.path().to_path_buf();
1890
1891 {
1892 let storage = SessionStoreV2::new(bamboo_home.clone()).await?;
1893 let mut root = Session::new("workspace-root", "m");
1894 root.set_workspace_path_meta(" /workspaces/root ");
1895 storage.save_session(&root).await?;
1896
1897 let mut child = Session::new_child_of("workspace-child", &root, "m", "child");
1898 child.set_workspace_path_meta("/workspaces/child");
1899 storage.save_session(&child).await?;
1900
1901 let legacy_without_workspace = Session::new("workspace-missing", "m");
1902 storage.save_session(&legacy_without_workspace).await?;
1903
1904 let broken_dir = bamboo_home.join("sessions/broken");
1906 tokio::fs::create_dir_all(&broken_dir).await?;
1907 tokio::fs::write(broken_dir.join("session.json"), b"{ invalid json").await?;
1908 }
1909
1910 let index_path = bamboo_home.join("sessions.json");
1911 let mut legacy: serde_json::Value =
1912 serde_json::from_slice(&tokio::fs::read(&index_path).await?)
1913 .map_err(|error| other_io_error(error.to_string()))?;
1914 legacy["version"] = serde_json::json!(2);
1915 for entry in legacy["sessions"]
1916 .as_object_mut()
1917 .expect("sessions object")
1918 .values_mut()
1919 {
1920 entry
1921 .as_object_mut()
1922 .expect("entry")
1923 .remove("workspace_path");
1924 }
1925 tokio::fs::write(
1926 &index_path,
1927 serde_json::to_vec_pretty(&legacy)
1928 .map_err(|error| other_io_error(error.to_string()))?,
1929 )
1930 .await?;
1931
1932 let migrated = SessionStoreV2::new(bamboo_home.clone()).await?;
1933 assert_eq!(
1934 migrated
1935 .get_index_entry("workspace-root")
1936 .await
1937 .and_then(|entry| entry.workspace_path),
1938 Some("/workspaces/root".to_string())
1939 );
1940 assert_eq!(
1941 migrated
1942 .get_index_entry("workspace-child")
1943 .await
1944 .and_then(|entry| entry.workspace_path),
1945 Some("/workspaces/child".to_string())
1946 );
1947 assert!(migrated.get_index_entry("broken").await.is_none());
1948 assert_eq!(
1949 migrated
1950 .get_index_entry("workspace-missing")
1951 .await
1952 .and_then(|entry| entry.workspace_path),
1953 None
1954 );
1955
1956 let persisted: SessionsIndex = serde_json::from_slice(&tokio::fs::read(index_path).await?)
1957 .map_err(|error| other_io_error(error.to_string()))?;
1958 assert_eq!(persisted.version, SESSIONS_INDEX_VERSION);
1959 Ok(())
1960 }
1961
1962 #[tokio::test]
1963 async fn v3_index_migrates_project_ids_from_root_and_child_sessions() -> io::Result<()> {
1964 let temp_dir = TempDir::new().map_err(io::Error::other)?;
1965 let bamboo_home = temp_dir.path().to_path_buf();
1966
1967 {
1968 let storage = SessionStoreV2::new(bamboo_home.clone()).await?;
1969 let mut root = Session::new("project-root", "m");
1970 root.set_project_id_meta("01JROOTPROJECT00000000000000");
1971 storage.save_session(&root).await?;
1972
1973 let mut child = Session::new_child_of("project-child", &root, "m", "child");
1974 child.metadata.insert(
1975 "project_id".to_string(),
1976 "01JCHILDPROJECT000000000000".to_string(),
1977 );
1978 storage.save_session(&child).await?;
1979
1980 let unassigned = Session::new("project-unassigned", "m");
1981 storage.save_session(&unassigned).await?;
1982 }
1983
1984 let index_path = bamboo_home.join("sessions.json");
1985 let mut legacy: serde_json::Value =
1986 serde_json::from_slice(&tokio::fs::read(&index_path).await?)
1987 .map_err(|error| other_io_error(error.to_string()))?;
1988 legacy["version"] = serde_json::json!(3);
1989 for entry in legacy["sessions"]
1990 .as_object_mut()
1991 .expect("sessions object")
1992 .values_mut()
1993 {
1994 entry.as_object_mut().expect("entry").remove("project_id");
1995 }
1996 tokio::fs::write(
1997 &index_path,
1998 serde_json::to_vec_pretty(&legacy)
1999 .map_err(|error| other_io_error(error.to_string()))?,
2000 )
2001 .await?;
2002
2003 let migrated = SessionStoreV2::new(bamboo_home.clone()).await?;
2004 assert_eq!(
2005 migrated
2006 .get_index_entry("project-root")
2007 .await
2008 .and_then(|entry| entry.project_id),
2009 Some("01JROOTPROJECT00000000000000".to_string())
2010 );
2011 assert_eq!(
2012 migrated
2013 .get_index_entry("project-child")
2014 .await
2015 .and_then(|entry| entry.project_id),
2016 Some("01JCHILDPROJECT000000000000".to_string())
2017 );
2018 assert_eq!(
2019 migrated
2020 .get_index_entry("project-unassigned")
2021 .await
2022 .and_then(|entry| entry.project_id),
2023 None
2024 );
2025
2026 let persisted: SessionsIndex = serde_json::from_slice(&tokio::fs::read(index_path).await?)
2027 .map_err(|error| other_io_error(error.to_string()))?;
2028 assert_eq!(persisted.version, SESSIONS_INDEX_VERSION);
2029 Ok(())
2030 }
2031
2032 #[tokio::test]
2033 async fn workspace_path_updates_index_on_full_and_runtime_only_saves() -> io::Result<()> {
2034 let (storage, _temp_dir) = create_temp_storage().await?;
2035 let mut session = Session::new("workspace-update", "m");
2036 session.set_workspace_path_meta(" /workspaces/first ");
2037 storage.save_session(&session).await?;
2038 assert_eq!(
2039 storage
2040 .get_index_entry(&session.id)
2041 .await
2042 .and_then(|entry| entry.workspace_path),
2043 Some("/workspaces/first".to_string())
2044 );
2045
2046 session.set_workspace_path_meta(" /workspaces/latest ");
2047 storage.save_runtime_state(&session).await?;
2048 assert_eq!(
2049 storage
2050 .get_index_entry(&session.id)
2051 .await
2052 .and_then(|entry| entry.workspace_path),
2053 Some("/workspaces/latest".to_string())
2054 );
2055 Ok(())
2056 }
2057
2058 #[tokio::test]
2059 async fn project_id_updates_index_on_full_and_runtime_only_saves() -> io::Result<()> {
2060 let (storage, _temp_dir) = create_temp_storage().await?;
2061 let mut session = Session::new("project-update", "m");
2062 session.set_project_id_meta(" 01JPROJECTFIRST000000000000 ");
2063 storage.save_session(&session).await?;
2064 assert_eq!(
2065 storage
2066 .get_index_entry(&session.id)
2067 .await
2068 .and_then(|entry| entry.project_id),
2069 Some("01JPROJECTFIRST000000000000".to_string())
2070 );
2071
2072 session.set_project_id_meta(" 01JPROJECTLATEST00000000000 ");
2073 storage.save_runtime_state(&session).await?;
2074 assert_eq!(
2075 storage
2076 .get_index_entry(&session.id)
2077 .await
2078 .and_then(|entry| entry.project_id),
2079 Some("01JPROJECTLATEST00000000000".to_string())
2080 );
2081 Ok(())
2082 }
2083
2084 #[tokio::test]
2085 async fn malformed_legacy_project_id_isolated_from_session_index() -> io::Result<()> {
2086 let (storage, _temp_dir) = create_temp_storage().await?;
2087 let mut malformed = Session::new("project-malformed", "m");
2088 malformed
2089 .metadata
2090 .insert("project_id".to_string(), "../unsafe".to_string());
2091 storage.save_session(&malformed).await?;
2092
2093 let healthy = Session::new("project-healthy", "m");
2094 storage.save_session(&healthy).await?;
2095
2096 assert_eq!(
2097 storage
2098 .get_index_entry(&malformed.id)
2099 .await
2100 .and_then(|entry| entry.project_id),
2101 None
2102 );
2103 assert!(storage.get_index_entry(&healthy.id).await.is_some());
2104 Ok(())
2105 }
2106
2107 #[tokio::test]
2108 async fn rebuild_overlays_runtime_sidecar_control_plane() -> io::Result<()> {
2109 let temp_dir = TempDir::new().map_err(io::Error::other)?;
2115 let bamboo_home = temp_dir.path().to_path_buf();
2116
2117 {
2118 let storage = SessionStoreV2::new(bamboo_home.clone()).await?;
2119
2120 let mut root = Session::new("rb-overlay", "m");
2122 root.metadata
2123 .insert("last_run_status".into(), "running".into());
2124 storage.save_session(&root).await?;
2125
2126 let mut updated = root.clone();
2129 updated
2130 .metadata
2131 .insert("last_run_status".into(), "completed".into());
2132 storage.save_runtime_state(&updated).await?;
2133
2134 let raw = read_session_json_raw(&storage, "rb-overlay").await;
2136 assert!(
2137 raw.contains("running"),
2138 "session.json must still carry the pre-sidecar status"
2139 );
2140 }
2141
2142 tokio::fs::write(bamboo_home.join("sessions.json"), b"{ not valid json").await?;
2144 let recovered = SessionStoreV2::new(bamboo_home.clone()).await?;
2145
2146 let entry = recovered
2150 .get_index_entry("rb-overlay")
2151 .await
2152 .expect("root recovered into rebuilt index");
2153 assert_eq!(
2154 entry.last_run_status.as_deref(),
2155 Some("completed"),
2156 "rebuild must overlay runtime.json control-plane, not the stale session.json"
2157 );
2158 Ok(())
2159 }
2160
2161 #[tokio::test]
2162 async fn missing_index_starts_empty_and_does_not_back_up() -> io::Result<()> {
2163 let temp_dir = TempDir::new().map_err(io::Error::other)?;
2166 let bamboo_home = temp_dir.path().to_path_buf();
2167
2168 let storage = SessionStoreV2::new(bamboo_home.clone()).await?;
2169 assert!(storage.list_index_entries().await.is_empty());
2170 assert!(
2171 !bamboo_home.join("sessions.json.bak").exists(),
2172 "a missing index must not produce a .bak backup"
2173 );
2174 Ok(())
2175 }
2176
2177 #[tokio::test]
2180 async fn migration_backfills_sidecars_for_legacy_sessions() -> io::Result<()> {
2181 let temp_dir = TempDir::new().map_err(io::Error::other)?;
2182 let bamboo_home = temp_dir.path().to_path_buf();
2183 let storage = SessionStoreV2::new(bamboo_home.clone()).await?;
2184
2185 let a = session_with_history("mig-a", 3, "run-A");
2188 let b = session_with_history("mig-b", 1, "run-B");
2189 storage.save_session(&a).await?;
2190 storage.save_session(&b).await?;
2191 for id in ["mig-a", "mig-b"] {
2192 let sidecar = storage.runtime_json_path(id).await?.unwrap();
2193 tokio::fs::remove_file(&sidecar).await?;
2194 assert!(!sidecar.exists());
2195 }
2196
2197 let migrated = storage.migrate_runtime_sidecars().await?;
2198 assert_eq!(migrated, 2, "both legacy sessions get a sidecar");
2199
2200 for (id, run) in [("mig-a", "run-A"), ("mig-b", "run-B")] {
2202 let side = storage.read_runtime_sidecar(id).await?.unwrap();
2203 assert!(side.messages.is_empty());
2204 assert_eq!(side.agent_runtime_state.as_ref().unwrap().run_id, run);
2205 }
2206 assert_eq!(
2208 storage.load_session("mig-a").await?.unwrap().messages.len(),
2209 3
2210 );
2211
2212 let marker = bamboo_home.join(RUNTIME_SIDECAR_MIGRATION_MARKER);
2214 assert!(marker.exists());
2215 assert_eq!(storage.migrate_runtime_sidecars().await?, 0);
2216 Ok(())
2217 }
2218
2219 #[tokio::test]
2220 async fn migration_is_idempotent_and_skips_existing_sidecars() -> io::Result<()> {
2221 let (storage, _t) = create_temp_storage().await?;
2222 storage
2224 .save_session(&session_with_history("mig-c", 2, "run-C"))
2225 .await?;
2226 let first = storage.migrate_runtime_sidecars().await?;
2227 assert_eq!(first, 0, "session saved in new format needs no migration");
2228 assert_eq!(storage.migrate_runtime_sidecars().await?, 0);
2230 Ok(())
2231 }
2232
2233 #[tokio::test]
2234 async fn migration_drops_legacy_denormalized_children_from_sidecar() -> io::Result<()> {
2235 let (storage, _t) = create_temp_storage().await?;
2239 let mut s = session_with_history("mig-legacy", 1, "run-L");
2240 storage.save_session(&s).await?;
2241
2242 let dir = storage.abs_path_from_rel(&storage.resolve_rel_path("mig-legacy").await.unwrap());
2245 s.agent_runtime_state = Some(AgentRuntimeState::new("run-L"));
2246 let mut value = serde_json::to_value(&s).unwrap();
2247 value["agent_runtime_state"]["children"]["active_ids"] = serde_json::json!(["ghost-child"]);
2248 tokio::fs::write(
2249 dir.join("session.json"),
2250 serde_json::to_vec_pretty(&value).unwrap(),
2251 )
2252 .await?;
2253 tokio::fs::remove_file(storage.runtime_json_path("mig-legacy").await?.unwrap()).await?;
2254
2255 assert_eq!(storage.migrate_runtime_sidecars().await?, 1);
2256
2257 let raw_sidecar =
2258 tokio::fs::read_to_string(storage.runtime_json_path("mig-legacy").await?.unwrap())
2259 .await?;
2260 assert!(
2261 !raw_sidecar.contains("ghost-child") && !raw_sidecar.contains("active_ids"),
2262 "legacy denormalized children must not survive migration: {raw_sidecar}"
2263 );
2264 Ok(())
2265 }
2266
2267 #[tokio::test]
2268 async fn list_child_run_statuses_filters_by_parent_and_reports_status() -> io::Result<()> {
2269 let (storage, _t) = create_temp_storage().await?;
2270
2271 let parent = Session::new("p-root".to_string(), "m".to_string());
2274 storage.save_session(&parent).await?;
2275 let other = Session::new("p-other".to_string(), "m".to_string());
2276 storage.save_session(&other).await?;
2277
2278 let mut c1 = Session::new_child("ch-done", "p-root", "m", "c1");
2279 c1.metadata
2280 .insert("last_run_status".to_string(), "completed".to_string());
2281 storage.save_session(&c1).await?;
2282
2283 let c2 = Session::new_child("ch-pending", "p-root", "m", "c2");
2284 storage.save_session(&c2).await?;
2285
2286 let foreign = Session::new_child("ch-foreign", "p-other", "m", "x");
2287 storage.save_session(&foreign).await?;
2288
2289 let mut got = storage.list_child_run_statuses("p-root").await?;
2290 got.sort_by(|a, b| a.0.cmp(&b.0));
2291 assert_eq!(got.len(), 2, "only p-root's children: {got:?}");
2292 assert_eq!(got[0].0, "ch-done");
2293 assert_eq!(got[0].1.as_deref(), Some("completed"));
2294 assert_eq!(got[1].0, "ch-pending");
2295 assert!(got[1].1.as_deref() != Some("completed"));
2297 Ok(())
2298 }
2299
2300 #[tokio::test]
2301 async fn list_sessions_by_run_status_matches_index_and_reports_parent() -> io::Result<()> {
2302 let (storage, _t) = create_temp_storage().await?;
2303
2304 let mut root = Session::new("r-susp".to_string(), "m".to_string());
2305 root.metadata
2306 .insert("last_run_status".to_string(), "suspended".to_string());
2307 storage.save_session(&root).await?;
2308
2309 let mut child = Session::new_child("ch-run", "r-susp", "m", "c");
2310 child
2311 .metadata
2312 .insert("last_run_status".to_string(), "running".to_string());
2313 storage.save_session(&child).await?;
2314
2315 let mut done = Session::new("r-done".to_string(), "m".to_string());
2316 done.metadata
2317 .insert("last_run_status".to_string(), "completed".to_string());
2318 storage.save_session(&done).await?;
2319
2320 let suspended = storage.list_sessions_by_run_status("suspended").await?;
2321 assert_eq!(suspended, vec![("r-susp".to_string(), None)]);
2322
2323 let running = storage.list_sessions_by_run_status("running").await?;
2324 assert_eq!(
2325 running,
2326 vec![("ch-run".to_string(), Some("r-susp".to_string()))]
2327 );
2328
2329 assert!(storage
2330 .list_sessions_by_run_status("timeout")
2331 .await?
2332 .is_empty());
2333 Ok(())
2334 }
2335
2336 #[tokio::test]
2337 async fn load_runtime_control_plane_reads_sidecar_without_messages() -> io::Result<()> {
2338 let (storage, _t) = create_temp_storage().await?;
2339 let s = session_with_history("sc-5", 5, "run-A");
2340 storage.save_session(&s).await?;
2341
2342 let cp = storage.load_runtime_control_plane("sc-5").await?.unwrap();
2343 assert!(
2344 cp.messages.is_empty(),
2345 "control-plane load must skip the message history"
2346 );
2347 assert_eq!(cp.agent_runtime_state.as_ref().unwrap().run_id, "run-A");
2348 Ok(())
2349 }
2350
2351 #[tokio::test]
2352 async fn test_save_and_load_session() -> io::Result<()> {
2353 let (storage, _temp_dir) = create_temp_storage().await?;
2354 let session = Session::new("session-1", "test-model");
2355
2356 storage.save_session(&session).await?;
2357 let loaded = storage.load_session(&session.id).await?;
2358
2359 assert!(loaded.is_some());
2360 let loaded = loaded.unwrap();
2361 assert_eq!(loaded.id, session.id);
2362 assert_eq!(loaded.model, session.model);
2363
2364 Ok(())
2365 }
2366
2367 #[tokio::test]
2368 async fn test_load_session_returns_none_when_not_found() -> io::Result<()> {
2369 let (storage, _temp_dir) = create_temp_storage().await?;
2370 let loaded = storage.load_session("nonexistent").await?;
2371 assert!(loaded.is_none());
2372 Ok(())
2373 }
2374
2375 #[tokio::test]
2376 async fn nested_grandchild_persists_under_root() -> io::Result<()> {
2377 let (storage, _t) = create_temp_storage().await?;
2382 let root = Session::new("root-1", "m");
2383 storage.save_session(&root).await?;
2384 let child = Session::new_child_of("child-1", &root, "m", "c");
2385 storage.save_session(&child).await?;
2386 let grandchild = Session::new_child_of("gc-1", &child, "m", "g");
2387 storage.save_session(&grandchild).await?;
2388
2389 let loaded = storage.load_session("gc-1").await?.expect("grandchild");
2390 assert_eq!(loaded.parent_session_id.as_deref(), Some("child-1"));
2391 assert_eq!(loaded.root_session_id, "root-1");
2392 assert_eq!(loaded.spawn_depth, 2);
2393
2394 let entry = storage.get_index_entry("gc-1").await.expect("indexed");
2396 assert_eq!(entry.parent_session_id.as_deref(), Some("child-1"));
2397 assert_eq!(entry.root_session_id, "root-1");
2398 Ok(())
2399 }
2400
2401 #[tokio::test]
2402 async fn test_list_index_entries_empty() -> io::Result<()> {
2403 let (storage, _temp_dir) = create_temp_storage().await?;
2404 let entries = storage.list_index_entries().await;
2405 assert!(entries.is_empty());
2406 Ok(())
2407 }
2408
2409 #[tokio::test]
2410 async fn test_list_index_entries_with_sessions() -> io::Result<()> {
2411 let (storage, _temp_dir) = create_temp_storage().await?;
2412
2413 let session1 = Session::new("session-1", "model-1");
2414 let session2 = Session::new("session-2", "model-2");
2415
2416 storage.save_session(&session1).await?;
2417 storage.save_session(&session2).await?;
2418
2419 let entries = storage.list_index_entries().await;
2420 assert_eq!(entries.len(), 2);
2421
2422 Ok(())
2423 }
2424
2425 #[tokio::test]
2426 async fn test_get_index_entry() -> io::Result<()> {
2427 let (storage, _temp_dir) = create_temp_storage().await?;
2428 let session = Session::new("session-1", "test-model");
2429
2430 storage.save_session(&session).await?;
2431
2432 let entry = storage.get_index_entry(&session.id).await;
2433 assert!(entry.is_some());
2434 let entry = entry.unwrap();
2435 assert_eq!(entry.id, session.id);
2436
2437 Ok(())
2438 }
2439
2440 #[tokio::test]
2441 async fn test_get_index_entry_returns_none_when_not_found() -> io::Result<()> {
2442 let (storage, _temp_dir) = create_temp_storage().await?;
2443 let entry = storage.get_index_entry("nonexistent").await;
2444 assert!(entry.is_none());
2445 Ok(())
2446 }
2447
2448 #[tokio::test]
2449 async fn test_delete_session() -> io::Result<()> {
2450 let (storage, _temp_dir) = create_temp_storage().await?;
2451 let session = Session::new("session-1", "test-model");
2452
2453 storage.save_session(&session).await?;
2454 assert!(storage.load_session(&session.id).await?.is_some());
2455
2456 let deleted = storage.delete_session(&session.id).await?;
2457 assert!(deleted);
2458 assert!(storage.load_session(&session.id).await?.is_none());
2459
2460 Ok(())
2461 }
2462
2463 #[tokio::test]
2464 async fn test_delete_session_returns_false_when_not_found() -> io::Result<()> {
2465 let (storage, _temp_dir) = create_temp_storage().await?;
2466 let deleted = storage.delete_session("nonexistent").await?;
2467 assert!(!deleted);
2468 Ok(())
2469 }
2470
2471 #[tokio::test]
2472 async fn delete_wins_before_delivery_without_recreating_an_orphan_inbox() -> io::Result<()> {
2473 let temp_dir = TempDir::new().map_err(io::Error::other)?;
2474 let storage = Arc::new(SessionStoreV2::new(temp_dir.path().to_path_buf()).await?);
2475 let session = Session::new("delete-inbox-race", "model");
2476 storage.save_session(&session).await?;
2477 let rel_path = storage
2478 .resolve_rel_path(&session.id)
2479 .await
2480 .expect("saved session has an indexed path");
2481 let session_dir = storage.abs_path_from_rel(&rel_path);
2482 let delivery_storage = Arc::new(SessionStoreV2::new(temp_dir.path().to_path_buf()).await?);
2487 let inbox = crate::FileSessionInbox::new(
2488 delivery_storage,
2489 bamboo_domain::SessionInboxLimits::default(),
2490 );
2491
2492 let lifecycle = storage.lock_session_lifecycle_exclusive().await?;
2496 let target = session.id.clone();
2497 let delivery = tokio::spawn(async move {
2498 inbox
2499 .deliver(&SessionMessageEnvelope::user_input(target, "late"))
2500 .await
2501 });
2502 tokio::task::yield_now().await;
2503 assert!(!delivery.is_finished());
2504
2505 assert!(
2506 storage
2507 .delete_session_recursive_locked(&session.id, true)
2508 .await?
2509 );
2510 assert!(!session_dir.exists());
2511 drop(lifecycle);
2512
2513 let error = tokio::time::timeout(std::time::Duration::from_secs(2), delivery)
2514 .await
2515 .expect("blocked delivery must resume after deletion")
2516 .expect("delivery task must not panic")
2517 .expect_err("a deleted target cannot acknowledge a new inbox");
2518 assert!(matches!(
2519 error,
2520 SessionInboxError::TargetNotFound(ref id) if id == &session.id
2521 ));
2522 assert!(
2523 !session_dir.exists(),
2524 "failed delivery must not recreate an orphan session/inbox tree"
2525 );
2526 Ok(())
2527 }
2528
2529 #[test]
2530 fn test_validate_session_id_empty() {
2531 assert!(validate_session_id("").is_err());
2532 }
2533
2534 #[test]
2535 fn test_validate_session_id_with_slash() {
2536 assert!(validate_session_id("session/1").is_err());
2537 }
2538
2539 #[test]
2540 fn test_validate_session_id_with_backslash() {
2541 assert!(validate_session_id("session\\1").is_err());
2542 }
2543
2544 #[test]
2545 fn test_validate_session_id_with_double_dot() {
2546 assert!(validate_session_id("session..1").is_err());
2547 }
2548
2549 #[test]
2550 fn test_validate_session_id_valid() {
2551 assert!(validate_session_id("session-123").is_ok());
2552 }
2553
2554 #[test]
2555 fn test_root_rel_path() {
2556 let path = SessionStoreV2::root_rel_path("session-123");
2557 assert_eq!(path, "sessions/session-123");
2558 }
2559
2560 #[test]
2561 fn test_child_rel_path() {
2562 let path = SessionStoreV2::child_rel_path("root-1", "child-2");
2563 assert_eq!(path, "sessions/root-1/children/child-2");
2564 }
2565
2566 #[test]
2567 fn test_mime_to_extension() {
2568 assert_eq!(mime_to_extension("image/png"), Some("png"));
2569 assert_eq!(mime_to_extension("image/jpeg"), Some("jpg"));
2570 assert_eq!(mime_to_extension("image/webp"), Some("webp"));
2571 assert_eq!(mime_to_extension("image/gif"), Some("gif"));
2572 assert_eq!(mime_to_extension("image/bmp"), Some("bmp"));
2573 assert_eq!(mime_to_extension("unknown/type"), None);
2574 }
2575
2576 #[test]
2577 fn test_extension_to_mime() {
2578 assert_eq!(extension_to_mime("png"), Some("image/png"));
2579 assert_eq!(extension_to_mime("jpg"), Some("image/jpeg"));
2580 assert_eq!(extension_to_mime("jpeg"), Some("image/jpeg"));
2581 assert_eq!(extension_to_mime("webp"), Some("image/webp"));
2582 assert_eq!(extension_to_mime("gif"), Some("image/gif"));
2583 assert_eq!(extension_to_mime("bmp"), Some("image/bmp"));
2584 assert_eq!(extension_to_mime("unknown"), None);
2585 }
2586
2587 #[test]
2588 fn test_extension_to_mime_case_insensitive() {
2589 assert_eq!(extension_to_mime("PNG"), Some("image/png"));
2590 assert_eq!(extension_to_mime("JPG"), Some("image/jpeg"));
2591 assert_eq!(extension_to_mime("JPEG"), Some("image/jpeg"));
2592 }
2593
2594 #[test]
2595 fn test_extension_to_mime_with_whitespace() {
2596 assert_eq!(extension_to_mime(" png "), Some("image/png"));
2597 assert_eq!(extension_to_mime("\tjpg\t"), Some("image/jpeg"));
2598 }
2599}