Add SessionIndex and SessionMaterializer · jauntywk.bsky.social/opencode-session-rs@6cf7e87

+258

src/index.rs

··· 1 + use crate::id::{MessageId, PartId, SessionId}; 2 + use crate::storage::{FileReader, MappedFile, StoragePaths}; 3 + use crate::types::Part; 4 + use crate::Result; 5 + use std::collections::HashMap; 6 + use std::path::PathBuf; 7 + use std::sync::Arc; 8 + 9 + #[derive(Debug, Clone)] 10 + pub struct SessionMeta { 11 + pub id: SessionId, 12 + pub title: String, 13 + pub created: i64, 14 + pub updated: i64, 15 + pub project_id: String, 16 + pub message_count: usize, 17 + } 18 + 19 + #[derive(Debug, Clone)] 20 + pub struct MessageMeta { 21 + pub id: MessageId, 22 + pub session_id: SessionId, 23 + pub role: MessageRole, 24 + pub part_count: usize, 25 + } 26 + 27 + #[derive(Debug, Clone, Copy, PartialEq, Eq)] 28 + pub enum MessageRole { 29 + User, 30 + Assistant, 31 + } 32 + 33 + #[derive(Debug, Clone)] 34 + pub struct PartRef { 35 + pub id: PartId, 36 + pub message_id: MessageId, 37 + pub session_id: SessionId, 38 + pub part_type: PartType, 39 + pub path: PathBuf, 40 + mmap: Option<Arc<MappedFile>>, 41 + } 42 + 43 + impl PartRef { 44 + pub fn new( 45 + id: PartId, 46 + message_id: MessageId, 47 + session_id: SessionId, 48 + part_type: PartType, 49 + path: PathBuf, 50 + ) -> Self { 51 + Self { 52 + id, 53 + message_id, 54 + session_id, 55 + part_type, 56 + path, 57 + mmap: None, 58 + } 59 + } 60 + 61 + pub fn load(&mut self, reader: &FileReader) -> Result<Part> { 62 + reader.read_part(&self.message_id, &self.id) 63 + } 64 + 65 + pub fn load_mapped( 66 + &mut self, 67 + cache: &crate::storage::MappedFileCache, 68 + ) -> Result<Arc<MappedFile>> { 69 + if self.mmap.is_none() { 70 + self.mmap = Some(cache.get(&self.path)?); 71 + } 72 + Ok(self.mmap.as_ref().unwrap().clone()) 73 + } 74 + } 75 + 76 + #[derive(Debug, Clone, Copy, PartialEq, Eq)] 77 + pub enum PartType { 78 + Text, 79 + Reasoning, 80 + Tool, 81 + File, 82 + StepStart, 83 + StepFinish, 84 + Snapshot, 85 + Patch, 86 + Agent, 87 + Subtask, 88 + Compaction, 89 + Retry, 90 + } 91 + 92 + impl PartType { 93 + pub fn from_str(s: &str) -> Option<Self> { 94 + match s { 95 + "text" => Some(Self::Text), 96 + "reasoning" => Some(Self::Reasoning), 97 + "tool" => Some(Self::Tool), 98 + "file" => Some(Self::File), 99 + "step-start" => Some(Self::StepStart), 100 + "step-finish" => Some(Self::StepFinish), 101 + "snapshot" => Some(Self::Snapshot), 102 + "patch" => Some(Self::Patch), 103 + "agent" => Some(Self::Agent), 104 + "subtask" => Some(Self::Subtask), 105 + "compaction" => Some(Self::Compaction), 106 + "retry" => Some(Self::Retry), 107 + _ => None, 108 + } 109 + } 110 + } 111 + 112 + #[derive(Debug, Default)] 113 + pub struct SessionIndex { 114 + pub sessions: HashMap<SessionId, SessionMeta>, 115 + pub messages: HashMap<MessageId, MessageMeta>, 116 + pub parts: HashMap<PartId, PartRef>, 117 + pub by_session: HashMap<SessionId, Vec<MessageId>>, 118 + pub by_message: HashMap<MessageId, Vec<PartId>>, 119 + pub by_project: HashMap<String, Vec<SessionId>>, 120 + } 121 + 122 + impl SessionIndex { 123 + pub fn new() -> Self { 124 + Self::default() 125 + } 126 + 127 + pub fn build(paths: &StoragePaths) -> Result<Self> { 128 + let mut index = Self::new(); 129 + let reader = FileReader::with_paths(paths.clone()); 130 + 131 + for project_id in paths.project_dirs()? { 132 + let session_ids = reader.list_sessions(&project_id)?; 133 + 134 + for session_id in &session_ids { 135 + if let Ok(session) = reader.read_session(&project_id, session_id) { 136 + let message_ids = reader.list_messages(session_id).unwrap_or_default(); 137 + let message_count = message_ids.len(); 138 + 139 + index.sessions.insert( 140 + session_id.clone(), 141 + SessionMeta { 142 + id: session_id.clone(), 143 + title: session.title, 144 + created: session.time.created, 145 + updated: session.time.updated, 146 + project_id: project_id.clone(), 147 + message_count, 148 + }, 149 + ); 150 + 151 + index 152 + .by_session 153 + .insert(session_id.clone(), message_ids.clone()); 154 + 155 + for msg_id in &message_ids { 156 + if let Ok(msg) = reader.read_message(session_id, msg_id) { 157 + let role = match &msg { 158 + crate::types::Message::User(_) => MessageRole::User, 159 + crate::types::Message::Assistant(_) => MessageRole::Assistant, 160 + }; 161 + 162 + let part_ids = reader.list_parts(msg_id).unwrap_or_default(); 163 + let part_count = part_ids.len(); 164 + 165 + index.messages.insert( 166 + msg_id.clone(), 167 + MessageMeta { 168 + id: msg_id.clone(), 169 + session_id: session_id.clone(), 170 + role, 171 + part_count, 172 + }, 173 + ); 174 + 175 + index.by_message.insert(msg_id.clone(), part_ids.clone()); 176 + 177 + for part_id in &part_ids { 178 + let path = paths.part_file(msg_id, part_id); 179 + index.parts.insert( 180 + part_id.clone(), 181 + PartRef::new( 182 + part_id.clone(), 183 + msg_id.clone(), 184 + session_id.clone(), 185 + PartType::Text, 186 + path, 187 + ), 188 + ); 189 + } 190 + } 191 + } 192 + } 193 + } 194 + 195 + index.by_project.insert(project_id.clone(), session_ids); 196 + } 197 + 198 + Ok(index) 199 + } 200 + 201 + pub fn session(&self, id: &SessionId) -> Option<&SessionMeta> { 202 + self.sessions.get(id) 203 + } 204 + 205 + pub fn message(&self, id: &MessageId) -> Option<&MessageMeta> { 206 + self.messages.get(id) 207 + } 208 + 209 + pub fn part(&self, id: &PartId) -> Option<&PartRef> { 210 + self.parts.get(id) 211 + } 212 + 213 + pub fn sessions_for_project(&self, project_id: &str) -> Vec<&SessionMeta> { 214 + self.by_project 215 + .get(project_id) 216 + .map(|ids| ids.iter().filter_map(|id| self.sessions.get(id)).collect()) 217 + .unwrap_or_default() 218 + } 219 + 220 + pub fn messages_for_session(&self, session_id: &SessionId) -> Vec<&MessageMeta> { 221 + self.by_session 222 + .get(session_id) 223 + .map(|ids| ids.iter().filter_map(|id| self.messages.get(id)).collect()) 224 + .unwrap_or_default() 225 + } 226 + 227 + pub fn parts_for_message(&self, message_id: &MessageId) -> Vec<&PartRef> { 228 + self.by_message 229 + .get(message_id) 230 + .map(|ids| ids.iter().filter_map(|id| self.parts.get(id)).collect()) 231 + .unwrap_or_default() 232 + } 233 + 234 + pub fn session_count(&self) -> usize { 235 + self.sessions.len() 236 + } 237 + 238 + pub fn message_count(&self) -> usize { 239 + self.messages.len() 240 + } 241 + 242 + pub fn part_count(&self) -> usize { 243 + self.parts.len() 244 + } 245 + } 246 + 247 + #[cfg(test)] 248 + mod tests { 249 + use super::*; 250 + 251 + #[test] 252 + fn test_index_new() { 253 + let index = SessionIndex::new(); 254 + assert_eq!(index.session_count(), 0); 255 + assert_eq!(index.message_count(), 0); 256 + assert_eq!(index.part_count(), 0); 257 + } 258 + }

+4 -6

src/lib.rs

··· 1 1 pub mod error; 2 2 pub mod id; 3 + pub mod index; 3 4 pub mod loader; 5 + pub mod materializer; 4 6 pub mod storage; 5 7 pub mod types; 6 8 7 - #[cfg(feature = "async")] 8 - pub mod index; 9 - 10 - #[cfg(feature = "async")] 11 - pub mod materializer; 12 - 13 9 #[cfg(feature = "watch")] 14 10 pub mod watch; 15 11 16 12 pub use error::{Error, Result}; 17 13 pub use id::{MessageId, PartId, SessionId}; 14 + pub use index::{SessionIndex, MessageMeta, PartRef, SessionMeta}; 18 15 pub use loader::SessionLoader; 16 + pub use materializer::SessionMaterializer; 19 17 pub use types::{ 20 18 message::{AssistantMessage, Message, UserMessage}, 21 19 part::Part,

+171

src/materializer.rs

··· 1 + use crate::id::{MessageId, PartId, SessionId}; 2 + use crate::index::{MessageMeta, PartRef, SessionIndex, SessionMeta}; 3 + use crate::loader::{LoadedSession, MessageWithParts, SessionTree}; 4 + use crate::storage::{FileReader, MappedFileCache, StoragePaths}; 5 + use crate::types::{Message, Part, SessionInfo}; 6 + use crate::Result; 7 + use std::sync::Arc; 8 + 9 + pub struct SessionMaterializer { 10 + reader: FileReader, 11 + index: SessionIndex, 12 + cache: MappedFileCache, 13 + } 14 + 15 + impl SessionMaterializer { 16 + pub fn new() -> Result<Self> { 17 + let paths = StoragePaths::detect()?; 18 + Self::with_paths(paths) 19 + } 20 + 21 + pub fn with_paths(paths: StoragePaths) -> Result<Self> { 22 + let reader = FileReader::with_paths(paths); 23 + let index = SessionIndex::build(reader.paths())?; 24 + let cache = MappedFileCache::new(); 25 + Ok(Self { 26 + reader, 27 + index, 28 + cache, 29 + }) 30 + } 31 + 32 + pub fn index(&self) -> &SessionIndex { 33 + &self.index 34 + } 35 + 36 + pub fn reader(&self) -> &FileReader { 37 + &self.reader 38 + } 39 + 40 + pub fn sessions(&self) -> &std::collections::HashMap<SessionId, SessionMeta> { 41 + &self.index.sessions 42 + } 43 + 44 + pub fn session(&self, id: &SessionId) -> Option<&SessionMeta> { 45 + self.index.session(id) 46 + } 47 + 48 + pub fn projects(&self) -> &std::collections::HashMap<String, Vec<SessionId>> { 49 + &self.index.by_project 50 + } 51 + 52 + pub fn sessions_for_project(&self, project_id: &str) -> Vec<&SessionMeta> { 53 + self.index.sessions_for_project(project_id) 54 + } 55 + 56 + pub fn messages_for_session(&self, session_id: &SessionId) -> Vec<&MessageMeta> { 57 + self.index.messages_for_session(session_id) 58 + } 59 + 60 + pub fn parts_for_message(&self, message_id: &MessageId) -> Vec<&PartRef> { 61 + self.index.parts_for_message(message_id) 62 + } 63 + 64 + pub fn load_session(&self, project_id: &str, id: &SessionId) -> Result<SessionInfo> { 65 + self.reader.read_session(project_id, id) 66 + } 67 + 68 + pub fn load_session_with_diff( 69 + &self, 70 + project_id: &str, 71 + id: &SessionId, 72 + ) -> Result<LoadedSession> { 73 + let info = self.reader.read_session(project_id, id)?; 74 + let diff = self.reader.read_diff(id).ok(); 75 + Ok(LoadedSession { info, diff }) 76 + } 77 + 78 + pub fn load_message(&self, session_id: &SessionId, id: &MessageId) -> Result<Message> { 79 + self.reader.read_message(session_id, id) 80 + } 81 + 82 + pub fn load_part(&self, message_id: &MessageId, id: &PartId) -> Result<Part> { 83 + self.reader.read_part(message_id, id) 84 + } 85 + 86 + pub fn load_message_with_parts( 87 + &self, 88 + session_id: &SessionId, 89 + message_id: &MessageId, 90 + ) -> Result<MessageWithParts> { 91 + let message = self.reader.read_message(session_id, message_id)?; 92 + let part_ids = self.reader.list_parts(message_id)?; 93 + let mut parts = Vec::with_capacity(part_ids.len()); 94 + for part_id in &part_ids { 95 + parts.push(self.reader.read_part(message_id, part_id)?); 96 + } 97 + Ok(MessageWithParts { message, parts }) 98 + } 99 + 100 + pub fn load_session_tree( 101 + &self, 102 + project_id: &str, 103 + session_id: &SessionId, 104 + ) -> Result<SessionTree> { 105 + let session = self.load_session_with_diff(project_id, session_id)?; 106 + let message_metas = self.index.messages_for_session(session_id); 107 + 108 + let mut messages = Vec::with_capacity(message_metas.len()); 109 + for msg_meta in message_metas { 110 + messages.push(self.load_message_with_parts(session_id, &msg_meta.id)?); 111 + } 112 + 113 + Ok(SessionTree { session, messages }) 114 + } 115 + 116 + pub fn sessions_by_time(&self, since: i64) -> Vec<&SessionMeta> { 117 + self.index 118 + .sessions 119 + .values() 120 + .filter(|s| s.created >= since) 121 + .collect() 122 + } 123 + 124 + pub fn sessions_updated_since(&self, since: i64) -> Vec<&SessionMeta> { 125 + self.index 126 + .sessions 127 + .values() 128 + .filter(|s| s.updated >= since) 129 + .collect() 130 + } 131 + 132 + pub fn mapped_file(&self, path: &std::path::Path) -> Result<Arc<crate::storage::MappedFile>> { 133 + self.cache.get(path) 134 + } 135 + 136 + pub fn stats(&self) -> Stats { 137 + Stats { 138 + sessions: self.index.sessions.len(), 139 + messages: self.index.messages.len(), 140 + parts: self.index.parts.len(), 141 + projects: self.index.by_project.len(), 142 + cached_files: self.cache.len(), 143 + } 144 + } 145 + } 146 + 147 + #[derive(Debug)] 148 + pub struct Stats { 149 + pub sessions: usize, 150 + pub messages: usize, 151 + pub parts: usize, 152 + pub projects: usize, 153 + pub cached_files: usize, 154 + } 155 + 156 + #[cfg(test)] 157 + mod tests { 158 + use super::*; 159 + 160 + #[test] 161 + fn test_stats() { 162 + let stats = Stats { 163 + sessions: 10, 164 + messages: 100, 165 + parts: 500, 166 + projects: 2, 167 + cached_files: 3, 168 + }; 169 + assert_eq!(stats.sessions, 10); 170 + } 171 + }

+2 -2

src/storage/mod.rs

··· 1 - pub mod paths; 2 1 pub mod mmap; 2 + pub mod paths; 3 3 pub mod reader; 4 4 5 + pub use mmap::{MappedFile, MappedFileCache}; 5 6 pub use paths::StoragePaths; 6 - pub use mmap::MappedFile; 7 7 pub use reader::FileReader;

+1

src/storage/paths.rs

··· 1 1 use crate::{Error, Result}; 2 2 use std::path::PathBuf; 3 3 4 + #[derive(Clone)] 4 5 pub struct StoragePaths { 5 6 pub root: PathBuf, 6 7 pub session: PathBuf,

Configure Feed

Configure Feed