1use chrono::{DateTime, Utc};
5use regex::Regex;
6use std::collections::{BTreeMap, BTreeSet};
7use std::fs;
8use std::path::Path;
9
10use crate::cas::AttachmentFileType;
11use crate::error::{Result, StorageError};
12
13use super::ocr::{ensure_empty_state_asset, retain_supported_ocr_annotations_ids};
14use super::paths::{
15 attachment_manifest_path, context_window_path, messages_path, ocr_annotations_path,
16};
17use super::{
18 default_ocr_annotations, AttachmentManifest, AttachmentManifestEntry, ContextWindow,
19 OcrAnnotations, SideChatData, SideChatMetadata, ThreadData, ThreadMessage, ThreadMetadata,
20 ThreadStorage, WorkspaceMetadata,
21};
22
23fn copy_dir_all(source: &Path, destination: &Path) -> Result<()> {
24 fs::create_dir_all(destination)?;
25
26 for entry in fs::read_dir(source)? {
27 let entry = entry?;
28 let file_type = entry.file_type()?;
29 let destination_path = destination.join(entry.file_name());
30
31 if file_type.is_dir() {
32 copy_dir_all(&entry.path(), &destination_path)?;
33 } else {
34 fs::copy(entry.path(), destination_path)?;
35 }
36 }
37
38 Ok(())
39}
40
41fn normalize_hash(value: &str) -> Option<String> {
42 let trimmed = value.trim();
43 if trimmed.len() == 64 && trimmed.bytes().all(|byte| byte.is_ascii_hexdigit()) {
44 return Some(trimmed.to_ascii_lowercase());
45 }
46
47 Path::new(trimmed)
48 .file_stem()
49 .and_then(|value| value.to_str())
50 .and_then(normalize_hash)
51}
52
53fn attachment_display_names(content: &str) -> BTreeMap<String, String> {
54 let re = Regex::new(r"\[([^\]\n]+)\]\((<[^>\n]+>|[^)\n]+)\)")
55 .expect("attachment markdown regex must compile");
56 let mut names = BTreeMap::new();
57
58 for capture in re.captures_iter(content) {
59 let Some(label) = capture.get(1).map(|value| value.as_str().trim()) else {
60 continue;
61 };
62 let Some(raw_path) = capture.get(2).map(|value| value.as_str().trim()) else {
63 continue;
64 };
65 let unwrapped = raw_path
66 .strip_prefix('<')
67 .and_then(|value| value.strip_suffix('>'))
68 .unwrap_or(raw_path);
69 let Some(path) = unwrapped.strip_prefix("file://") else {
70 continue;
71 };
72 if let Some(hash) = normalize_hash(path) {
73 names.entry(hash).or_insert_with(|| label.to_string());
74 }
75 }
76
77 names
78}
79
80fn sort_attachment_manifest(manifest: &mut AttachmentManifest, initial_hash: &str) {
81 manifest.sort_by(|left, right| {
82 let left_initial = left.attachment_hash == initial_hash;
83 let right_initial = right.attachment_hash == initial_hash;
84 match (left_initial, right_initial) {
85 (true, false) => std::cmp::Ordering::Less,
86 (false, true) => std::cmp::Ordering::Greater,
87 _ => right
88 .last_mention_at
89 .cmp(&left.last_mention_at)
90 .then_with(|| left.attachment_hash.cmp(&right.attachment_hash)),
91 }
92 });
93}
94
95fn validate_message_ids(messages: &[ThreadMessage]) -> Result<()> {
96 let mut seen = BTreeSet::new();
97 for message in messages {
98 let id = message.id();
99 if !ThreadMessage::is_valid_id(id) {
100 return Err(StorageError::InvalidThreadMessage(format!(
101 "message ID `{id}` must use the msg-<UUID> format"
102 )));
103 }
104 if !seen.insert(id) {
105 return Err(StorageError::InvalidThreadMessage(format!(
106 "duplicate message ID `{id}`"
107 )));
108 }
109 let mut attachment_hashes = BTreeSet::new();
110 for attachment in message.attachments() {
111 let hash = attachment.attachment_hash.as_str();
112 if hash.len() != 64
113 || !hash
114 .bytes()
115 .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
116 {
117 return Err(StorageError::InvalidThreadMessage(format!(
118 "attachment hash `{hash}` is not a canonical BLAKE3 hash"
119 )));
120 }
121 if !attachment_hashes.insert(hash) {
122 return Err(StorageError::InvalidThreadMessage(format!(
123 "message `{id}` repeats attachment hash `{hash}`"
124 )));
125 }
126 }
127 }
128 Ok(())
129}
130
131impl ThreadStorage {
132 pub fn attachment_manifest_entry(
133 &self,
134 attachment_hash: &str,
135 display_name: &str,
136 last_mention_at: DateTime<Utc>,
137 ) -> Result<AttachmentManifestEntry> {
138 let hash = normalize_hash(attachment_hash).ok_or(StorageError::InvalidHash)?;
139 self.find_object_blob(&hash)?;
140 let mut object_manifest = self.load_object_manifest(&hash)?;
141
142 if object_manifest.file_context.file_type == AttachmentFileType::TextLocal
143 && object_manifest.file_context.file_brief.is_none()
144 {
145 let bytes = fs::read(self.find_object_blob(&hash)?)?;
146 object_manifest.file_context.file_brief =
147 Some(std::str::from_utf8(&bytes)?.to_string());
148 self.save_object_manifest(&hash, &object_manifest)?;
149 }
150
151 Ok(AttachmentManifestEntry {
152 attachment_hash: hash,
153 display_name: display_name.trim().to_string(),
154 file_type: object_manifest.file_context.file_type,
155 file_brief: object_manifest.file_context.file_brief,
156 last_mention_at,
157 })
158 }
159
160 fn apply_user_message_attachments_to_manifest(
161 &self,
162 manifest: &mut AttachmentManifest,
163 initial_hash: &str,
164 message: &ThreadMessage,
165 ) -> Result<()> {
166 let ThreadMessage::User {
167 content,
168 timestamp,
169 attachments,
170 ..
171 } = message
172 else {
173 return Ok(());
174 };
175 let names = attachment_display_names(content);
176
177 for attachment in attachments {
178 let hash =
179 normalize_hash(&attachment.attachment_hash).ok_or(StorageError::InvalidHash)?;
180 let fallback_name = self
181 .find_object_blob(&hash)?
182 .file_name()
183 .and_then(|value| value.to_str())
184 .unwrap_or("attachment")
185 .to_string();
186 let display_name = names.get(&hash).cloned().unwrap_or(fallback_name);
187 let fresh = self.attachment_manifest_entry(&hash, &display_name, *timestamp)?;
188
189 if let Some(existing) = manifest
190 .iter_mut()
191 .find(|entry| entry.attachment_hash == hash)
192 {
193 existing.display_name = fresh.display_name;
194 existing.file_type = fresh.file_type;
195 existing.file_brief = fresh.file_brief;
196 if existing.last_mention_at < *timestamp {
197 existing.last_mention_at = *timestamp;
198 }
199 } else {
200 manifest.push(fresh);
201 }
202 }
203
204 sort_attachment_manifest(manifest, initial_hash);
205 Ok(())
206 }
207
208 fn apply_user_message_attachments(
209 &self,
210 thread: &mut ThreadData,
211 message: &ThreadMessage,
212 ) -> Result<()> {
213 self.apply_user_message_attachments_to_manifest(
214 &mut thread.attachment_manifest,
215 &thread.metadata.image_hash,
216 message,
217 )?;
218 Ok(())
219 }
220
221 fn save_thread_files(&self, thread: &ThreadData) -> Result<()> {
222 let thread_dir = self.thread_dir(&thread.metadata.id);
223 fs::create_dir_all(&thread_dir)?;
224
225 let ocr_path = ocr_annotations_path(&thread_dir);
226 let mut ocr_data = thread.ocr_data.clone();
227 ensure_empty_state_asset(&mut ocr_data);
228 retain_supported_ocr_annotations_ids(&mut ocr_data);
229 super::atomic_write(
230 &ocr_path,
231 serde_json::to_string_pretty(&ocr_data)?.as_bytes(),
232 )?;
233
234 let context_path = context_window_path(&thread_dir);
235 if !context_path.exists() {
236 super::atomic_write(
237 &context_path,
238 serde_json::to_string_pretty(&thread.context_window)?.as_bytes(),
239 )?;
240 }
241
242 super::atomic_write(
243 &messages_path(&thread_dir),
244 serde_json::to_string_pretty(&thread.messages)?.as_bytes(),
245 )?;
246 super::atomic_write(
247 &attachment_manifest_path(&thread_dir),
248 serde_json::to_string_pretty(&thread.attachment_manifest)?.as_bytes(),
249 )?;
250 Ok(())
251 }
252
253 pub fn save_thread(&self, thread: &ThreadData) -> Result<()> {
254 self.save_thread_files(thread)?;
255 self.update_index(&thread.metadata)
256 }
257
258 pub fn save_thread_in_workspace(&self, thread: &ThreadData, workspace_id: &str) -> Result<()> {
259 self.save_thread_files(thread)?;
260 self.update_index_in_workspace(&thread.metadata, Some(workspace_id))
261 }
262
263 fn save_sidechat_files(&self, sidechat: &SideChatData) -> Result<()> {
264 let thread_dir = self.thread_dir(&sidechat.metadata.id);
265 fs::create_dir_all(&thread_dir)?;
266 super::atomic_write(
267 &context_window_path(&thread_dir),
268 serde_json::to_string_pretty(&sidechat.context_window)?.as_bytes(),
269 )?;
270 super::atomic_write(
271 &messages_path(&thread_dir),
272 serde_json::to_string_pretty(&sidechat.messages)?.as_bytes(),
273 )?;
274 super::atomic_write(
275 &attachment_manifest_path(&thread_dir),
276 serde_json::to_string_pretty(&sidechat.attachment_manifest)?.as_bytes(),
277 )?;
278 Ok(())
279 }
280
281 pub fn save_sidechat(&self, sidechat: &SideChatData) -> Result<()> {
282 validate_message_ids(&sidechat.messages)?;
283 let mut persisted = sidechat.clone();
284 let messages = persisted.messages.clone();
285 for message in &messages {
286 self.apply_user_message_attachments_to_manifest(
287 &mut persisted.attachment_manifest,
288 "",
289 message,
290 )?;
291 }
292 self.save_sidechat_files(&persisted)?;
293 self.update_sidechat_index(&persisted.metadata)
294 }
295
296 pub fn load_sidechat(&self, sidechat_id: &str) -> Result<SideChatData> {
297 let thread_dir = self.thread_dir(sidechat_id);
298 if !thread_dir.exists() {
299 return Err(StorageError::ThreadNotFound(sidechat_id.to_string()));
300 }
301 let metadata = self.get_sidechat_metadata(sidechat_id)?;
302 let messages = self.load_messages(sidechat_id)?;
303 let context_window = serde_json::from_str::<ContextWindow>(&fs::read_to_string(
304 context_window_path(&thread_dir),
305 )?)?;
306 let attachment_manifest = serde_json::from_str::<AttachmentManifest>(&fs::read_to_string(
307 attachment_manifest_path(&thread_dir),
308 )?)?;
309 Ok(SideChatData {
310 metadata,
311 messages,
312 context_window,
313 attachment_manifest,
314 })
315 }
316
317 pub fn update_sidechat_metadata(&self, metadata: &SideChatMetadata) -> Result<()> {
318 if !self.thread_dir(&metadata.id).exists() {
319 return Err(StorageError::ThreadNotFound(metadata.id.clone()));
320 }
321 self.update_sidechat_index(metadata)
322 }
323
324 pub fn delete_sidechat(&self, sidechat_id: &str) -> Result<()> {
325 let thread_dir = self.thread_dir(sidechat_id);
326 if thread_dir.exists() {
327 fs::remove_dir_all(thread_dir)?;
328 }
329 self.remove_sidechat_from_index(sidechat_id)
330 }
331
332 pub fn fork_sidechat_latest(&self, sidechat_id: &str) -> Result<SideChatMetadata> {
333 let source_dir = self.thread_dir(sidechat_id);
334 if !source_dir.exists() {
335 return Err(StorageError::ThreadNotFound(sidechat_id.to_string()));
336 }
337
338 let source = self.load_sidechat(sidechat_id)?;
339 let metadata = SideChatMetadata::new(format!("forked {}", source.metadata.title));
340 let destination_dir = self.thread_dir(&metadata.id);
341 copy_dir_all(&source_dir, &destination_dir)?;
342
343 let mut forked = source;
344 forked.metadata = metadata.clone();
345 forked.attachment_manifest.clear();
346 self.save_sidechat(&forked)?;
347 Ok(metadata)
348 }
349
350 pub fn set_thread_workspace(&self, thread_id: &str, workspace_id: Option<&str>) -> Result<()> {
351 let metadata = self.get_index_metadata(thread_id)?;
352 self.update_index_in_workspace(&metadata, workspace_id)
353 }
354
355 pub fn load_thread(&self, thread_id: &str) -> Result<ThreadData> {
356 let thread_dir = self.thread_dir(thread_id);
357 if !thread_dir.exists() {
358 return Err(StorageError::ThreadNotFound(thread_id.to_string()));
359 }
360
361 let metadata = self.get_index_metadata(thread_id)?;
362 let ocr_path = ocr_annotations_path(&thread_dir);
363 let mut annotations_changed = false;
364 let mut ocr_data: OcrAnnotations = if ocr_path.exists() {
365 serde_json::from_str(&fs::read_to_string(&ocr_path)?)?
366 } else {
367 default_ocr_annotations()
368 };
369 annotations_changed |= ensure_empty_state_asset(&mut ocr_data);
370 annotations_changed |= retain_supported_ocr_annotations_ids(&mut ocr_data);
371 if annotations_changed {
372 super::atomic_write(
373 &ocr_path,
374 serde_json::to_string_pretty(&ocr_data)?.as_bytes(),
375 )?;
376 }
377
378 let messages = self.load_messages(thread_id)?;
379 let context_window = if context_window_path(&thread_dir).exists() {
380 serde_json::from_str::<ContextWindow>(&fs::read_to_string(context_window_path(
381 &thread_dir,
382 ))?)?
383 } else {
384 ContextWindow::default()
385 };
386 let attachment_manifest = serde_json::from_str::<AttachmentManifest>(&fs::read_to_string(
387 attachment_manifest_path(&thread_dir),
388 )?)?;
389 let image_tone = self.get_image_tone(&metadata.image_hash);
390 let reverse_image_search = self.get_reverse_image_search_cache(&metadata.image_hash)?;
391
392 Ok(ThreadData {
393 metadata,
394 messages,
395 ocr_data,
396 context_window,
397 reverse_image_search,
398 attachment_manifest,
399 image_tone,
400 })
401 }
402
403 pub fn load_messages(&self, thread_id: &str) -> Result<Vec<ThreadMessage>> {
404 let thread_dir = self.thread_dir(thread_id);
405 if !thread_dir.exists() {
406 return Err(StorageError::ThreadNotFound(thread_id.to_string()));
407 }
408
409 let path = messages_path(&thread_dir);
410 if !path.exists() {
411 return Ok(Vec::new());
412 }
413
414 Ok(serde_json::from_str::<Vec<ThreadMessage>>(
415 &fs::read_to_string(path)?,
416 )?)
417 }
418
419 pub fn list_threads(&self) -> Result<Vec<ThreadMetadata>> {
420 let index = self.read_index()?;
421 let mut threads = index
422 .workspaces
423 .into_iter()
424 .flat_map(|workspace| workspace.threads.into_values())
425 .chain(index.unassigned_threads.into_values())
426 .collect::<Vec<_>>();
427 threads.sort_by_key(|thread| std::cmp::Reverse(thread.updated_at));
428 Ok(threads)
429 }
430
431 pub fn list_workspaces(&self) -> Result<Vec<WorkspaceMetadata>> {
432 Ok(self.read_index()?.workspaces)
433 }
434
435 pub fn delete_thread(&self, thread_id: &str) -> Result<()> {
436 self.delete_threads(&[thread_id.to_string()])
437 }
438
439 pub fn delete_threads(&self, thread_ids: &[String]) -> Result<()> {
440 for thread_id in thread_ids {
441 let thread_dir = self.thread_dir(thread_id);
442 if thread_dir.exists() {
443 fs::remove_dir_all(&thread_dir)?;
444 }
445 }
446 self.remove_many_from_index(thread_ids)
447 }
448
449 pub fn fork_thread_latest(&self, thread_id: &str) -> Result<ThreadMetadata> {
450 let source_dir = self.thread_dir(thread_id);
451 if !source_dir.exists() {
452 return Err(StorageError::ThreadNotFound(thread_id.to_string()));
453 }
454
455 let source_workspace_id = self.get_thread_workspace_id(thread_id)?;
456 let source_thread = self.load_thread(thread_id)?;
457 let mut metadata = ThreadMetadata::new(
458 format!("forked {}", source_thread.metadata.title),
459 source_thread.metadata.image_hash.clone(),
460 );
461 metadata.pinned_at = None;
462 let destination_dir = self.thread_dir(&metadata.id);
463 copy_dir_all(&source_dir, &destination_dir)?;
464
465 let mut forked_thread = source_thread;
466 forked_thread.metadata = metadata.clone();
467 let initial_hash = forked_thread.metadata.image_hash.clone();
468 let initial = forked_thread
469 .attachment_manifest
470 .iter()
471 .find(|entry| entry.attachment_hash == initial_hash)
472 .cloned()
473 .ok_or(StorageError::ImageNotFound(initial_hash))?;
474 forked_thread.attachment_manifest = vec![initial];
475 let retained_messages = forked_thread.messages.clone();
476 for message in &retained_messages {
477 self.apply_user_message_attachments(&mut forked_thread, message)?;
478 }
479 self.save_thread_files(&forked_thread)?;
480 self.update_index_in_workspace(&metadata, source_workspace_id.as_deref())?;
481 Ok(metadata)
482 }
483
484 pub fn update_thread_metadata(&self, metadata: &ThreadMetadata) -> Result<()> {
485 if !self.thread_dir(&metadata.id).exists() {
486 return Err(StorageError::ThreadNotFound(metadata.id.clone()));
487 }
488 self.update_index(metadata)?;
489 Ok(())
490 }
491}