1use std::path::Path;
11
12use tokio::fs;
13
14use crate::error::SessionError;
15use crate::event::{SessionEvent, SessionEventEnvelope};
16use crate::log::SessionEventLog;
17use crate::replay::ReplayEngine;
18use crate::store::SessionStore;
19
20const BLOBS_DIR_NAME: &str = "blobs";
22
23#[derive(Debug, Clone)]
25pub struct ForkResult {
26 pub new_session_id: String,
28 pub events_copied: usize,
31}
32
33pub struct ForkEngine;
35
36impl ForkEngine {
37 #[tracing::instrument(name = "session.fork.run", skip_all, level = "info", fields(at_seq))]
57 pub async fn fork(
58 data_dir: &Path,
59 src_id: &str,
60 new_id: &str,
61 at_seq: Option<u64>,
62 store: &SessionStore,
63 owner: Option<&str>,
64 ) -> Result<ForkResult, SessionError> {
65 if store.get(src_id).await?.is_none() {
66 return Err(SessionError::NotFound(src_id.to_owned()));
67 }
68
69 let src_dir = crate::session_dir(data_dir, src_id);
70 let src_log = SessionEventLog::open(&src_dir).await?;
71 let all_events = src_log.read_all().await?;
72
73 let total = u64::try_from(all_events.len()).unwrap_or(u64::MAX);
74 let at_seq = at_seq.unwrap_or(total);
75 if at_seq > total {
76 return Err(SessionError::InvalidForkPoint(format!(
77 "at_seq={at_seq} exceeds source session's event count={total}"
78 )));
79 }
80
81 ReplayEngine::replay(&src_dir, Some(at_seq)).await?;
84
85 let take_n = usize::try_from(at_seq).unwrap_or(usize::MAX);
86 let to_copy: Vec<_> = all_events.iter().take(take_n).cloned().collect();
87 let (cwd, provider_name, model) = to_copy
88 .iter()
89 .find_map(|e| match &e.kind {
90 SessionEvent::SessionStarted {
91 cwd,
92 provider_name,
93 model,
94 ..
95 } => Some((cwd.clone(), provider_name.clone(), model.clone())),
96 _ => None,
97 })
98 .unwrap_or_default();
99
100 let child_dir = crate::session_dir(data_dir, new_id);
101 let child_log = SessionEventLog::open(&child_dir).await?;
102
103 child_log
104 .append(
105 None,
106 None,
107 SessionEvent::SessionStarted {
108 session_id: new_id.to_owned(),
109 cwd,
110 provider_name,
111 model,
112 forked_from: Some((src_id.to_owned(), at_seq)),
113 },
114 )
115 .await?;
116 for envelope in &to_copy {
117 child_log
118 .append(envelope.turn_id, envelope.parent_seq, envelope.kind.clone())
119 .await?;
120 }
121
122 copy_referenced_blobs(&src_dir, &child_dir, &to_copy).await?;
123
124 store.record_fork(new_id, src_id, at_seq, owner).await?;
125 store
126 .update_seq(
127 new_id,
128 child_log.last_seq().unwrap_or(0),
129 to_copy.len() as u64 + 1,
130 )
131 .await?;
132
133 src_log
135 .append(
136 None,
137 None,
138 SessionEvent::ForkPoint {
139 new_session_id: new_id.to_owned(),
140 },
141 )
142 .await?;
143
144 Ok(ForkResult {
145 new_session_id: new_id.to_owned(),
146 events_copied: to_copy.len(),
147 })
148 }
149}
150
151async fn copy_referenced_blobs(
183 src_dir: &Path,
184 child_dir: &Path,
185 events: &[SessionEventEnvelope],
186) -> Result<(), SessionError> {
187 let mut hashes: Vec<&str> = Vec::new();
188 for envelope in events {
189 let SessionEvent::UserMessage { image_refs, .. } = &envelope.kind else {
190 continue;
191 };
192 for hash in image_refs {
193 validate_blob_hash(hash)?;
194 hashes.push(hash.as_str());
195 }
196 }
197
198 if hashes.is_empty() {
199 return Ok(());
200 }
201
202 hashes.sort_unstable();
208 hashes.dedup();
209
210 let src_blobs = src_dir.join(BLOBS_DIR_NAME);
211 let child_blobs = child_dir.join(BLOBS_DIR_NAME);
212 fs::create_dir_all(&child_blobs).await?;
213 crate::log::set_permissions(&child_blobs, 0o700).await?;
214
215 for hash in hashes {
216 let src_blob = src_blobs.join(hash);
217 let child_blob = child_blobs.join(hash);
218
219 match fs::hard_link(&src_blob, &child_blob).await {
220 Ok(()) => {}
221 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
222 tracing::warn!(
223 blob = hash,
224 path = %src_blob.display(),
225 "fork: referenced blob missing on parent's disk, skipping"
226 );
227 }
228 Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
229 tracing::debug!(
244 blob = hash,
245 path = %child_blob.display(),
246 "fork: blob already linked in child, skipping"
247 );
248 }
249 Err(_) => {
250 fs::copy(&src_blob, &child_blob).await?;
255 }
256 }
257 }
258
259 Ok(())
260}
261
262fn validate_blob_hash(hash: &str) -> Result<(), SessionError> {
268 if hash.is_empty() || !hash.chars().all(|c| c.is_ascii_hexdigit()) {
269 return Err(SessionError::InvalidBlobHash(hash.to_owned()));
270 }
271 Ok(())
272}
273
274#[cfg(test)]
275mod tests {
276 use super::*;
277 use crate::store::SessionStore;
278
279 async fn make_pool() -> zeph_db::DbPool {
280 let config = zeph_db::DbConfig {
281 url: ":memory:".to_owned(),
282 ..Default::default()
283 };
284 let pool = config
285 .connect()
286 .await
287 .expect("connect in-memory sqlite pool");
288 zeph_db::run_migrations(&pool)
289 .await
290 .expect("run migrations");
291 pool
292 }
293
294 async fn seed_parent(data_dir: &Path, store: &SessionStore, id: &str) {
295 store.create(id).await.unwrap();
296 let dir = crate::session_dir(data_dir, id);
297 let log = SessionEventLog::open(&dir).await.unwrap();
298 log.append(
299 None,
300 None,
301 SessionEvent::SessionStarted {
302 session_id: id.to_owned(),
303 cwd: "/repo".to_owned(),
304 provider_name: "claude".to_owned(),
305 model: "opus".to_owned(),
306 forked_from: None,
307 },
308 )
309 .await
310 .unwrap();
311 log.append(
312 None,
313 None,
314 SessionEvent::UserMessage {
315 text: "hello".to_owned(),
316 image_refs: vec![],
317 },
318 )
319 .await
320 .unwrap();
321 log.append(
322 None,
323 None,
324 SessionEvent::AssistantMessage {
325 parts: vec![zeph_llm::provider::MessagePart::Text {
326 text: "hi".to_owned(),
327 }],
328 },
329 )
330 .await
331 .unwrap();
332 log.append(
333 None,
334 None,
335 SessionEvent::UserMessage {
336 text: "second turn".to_owned(),
337 image_refs: vec![],
338 },
339 )
340 .await
341 .unwrap();
342 store
343 .update_seq(id, log.last_seq().unwrap(), 4)
344 .await
345 .unwrap();
346 }
347
348 #[tokio::test]
349 async fn test_fork_copies_events() {
350 let store = SessionStore::new(make_pool().await);
351 let data_dir = tempfile::tempdir().unwrap();
352 seed_parent(data_dir.path(), &store, "parent").await;
353
354 let result = ForkEngine::fork(data_dir.path(), "parent", "child", Some(3), &store, None)
355 .await
356 .unwrap();
357 assert_eq!(result.events_copied, 3);
358 assert_eq!(result.new_session_id, "child");
359
360 let child_dir = crate::session_dir(data_dir.path(), &result.new_session_id);
361 let child_log = SessionEventLog::open(&child_dir).await.unwrap();
362 let events = child_log.read_all().await.unwrap();
363 assert_eq!(events.len(), 4);
365 }
366
367 #[tokio::test]
368 async fn test_fork_provenance_metadata() {
369 let store = SessionStore::new(make_pool().await);
370 let data_dir = tempfile::tempdir().unwrap();
371 seed_parent(data_dir.path(), &store, "parent").await;
372
373 ForkEngine::fork(data_dir.path(), "parent", "child", Some(2), &store, None)
374 .await
375 .unwrap();
376
377 let meta = store.get("child").await.unwrap().unwrap();
378 assert_eq!(meta.forked_from.as_deref(), Some("parent"));
379 assert_eq!(meta.forked_at_seq, Some(2));
380 }
381
382 #[tokio::test]
386 async fn fork_propagates_owner_to_child_row() {
387 let pool = make_pool().await;
388 let store = SessionStore::new(pool.clone());
389 let data_dir = tempfile::tempdir().unwrap();
390 seed_parent(data_dir.path(), &store, "parent").await;
391
392 ForkEngine::fork(
393 data_dir.path(),
394 "parent",
395 "child",
396 Some(2),
397 &store,
398 Some("alice"),
399 )
400 .await
401 .unwrap();
402
403 let owner_key: Option<String> = zeph_db::query_scalar(zeph_db::sql!(
404 "SELECT owner_key FROM acp_sessions WHERE id = ?"
405 ))
406 .bind("child")
407 .fetch_one(&pool)
408 .await
409 .unwrap();
410 assert_eq!(owner_key.as_deref(), Some("alice"));
411 }
412
413 #[tokio::test]
414 async fn test_fork_appends_forkpoint_to_parent() {
415 let store = SessionStore::new(make_pool().await);
416 let data_dir = tempfile::tempdir().unwrap();
417 seed_parent(data_dir.path(), &store, "parent").await;
418
419 ForkEngine::fork(data_dir.path(), "parent", "child", Some(2), &store, None)
420 .await
421 .unwrap();
422
423 let parent_dir = crate::session_dir(data_dir.path(), "parent");
424 let parent_log = SessionEventLog::open(&parent_dir).await.unwrap();
425 let events = parent_log.read_all().await.unwrap();
426 assert!(matches!(
427 events.last().unwrap().kind,
428 SessionEvent::ForkPoint { .. }
429 ));
430 }
431
432 #[tokio::test]
433 async fn test_fork_rejects_seq_beyond_source() {
434 let store = SessionStore::new(make_pool().await);
435 let data_dir = tempfile::tempdir().unwrap();
436 seed_parent(data_dir.path(), &store, "parent").await;
437
438 let err = ForkEngine::fork(data_dir.path(), "parent", "child", Some(100), &store, None)
439 .await
440 .unwrap_err();
441 assert!(matches!(err, SessionError::InvalidForkPoint(_)));
442 }
443
444 #[tokio::test]
445 async fn test_fork_rejects_unknown_source() {
446 let store = SessionStore::new(make_pool().await);
447 let data_dir = tempfile::tempdir().unwrap();
448
449 let err = ForkEngine::fork(data_dir.path(), "no-such", "child", Some(0), &store, None)
450 .await
451 .unwrap_err();
452 assert!(matches!(err, SessionError::NotFound(_)));
453 }
454
455 #[tokio::test]
456 async fn test_fork_none_copies_everything() {
457 let store = SessionStore::new(make_pool().await);
458 let data_dir = tempfile::tempdir().unwrap();
459 seed_parent(data_dir.path(), &store, "parent").await;
460
461 let result = ForkEngine::fork(data_dir.path(), "parent", "child", None, &store, None)
462 .await
463 .unwrap();
464 assert_eq!(result.events_copied, 4);
466 }
467
468 #[tokio::test]
471 async fn test_fork_copies_referenced_blobs() {
472 let store = SessionStore::new(make_pool().await);
473 let data_dir = tempfile::tempdir().unwrap();
474 seed_parent(data_dir.path(), &store, "parent").await;
475
476 let parent_dir = crate::session_dir(data_dir.path(), "parent");
477 let parent_blobs = parent_dir.join("blobs");
478 tokio::fs::create_dir_all(&parent_blobs).await.unwrap();
479 tokio::fs::write(parent_blobs.join("a1b2c3"), b"image-bytes")
480 .await
481 .unwrap();
482
483 let parent_log = SessionEventLog::open(&parent_dir).await.unwrap();
484 parent_log
485 .append(
486 None,
487 None,
488 SessionEvent::UserMessage {
489 text: "with image".to_owned(),
490 image_refs: vec!["a1b2c3".to_owned()],
491 },
492 )
493 .await
494 .unwrap();
495 store.update_seq("parent", 4, 5).await.unwrap();
496
497 let result = ForkEngine::fork(data_dir.path(), "parent", "child", Some(5), &store, None)
498 .await
499 .unwrap();
500 assert_eq!(result.events_copied, 5);
501
502 let child_dir = crate::session_dir(data_dir.path(), "child");
503 let child_blob = child_dir.join("blobs").join("a1b2c3");
504 let copied = tokio::fs::read(&child_blob).await.unwrap();
505 assert_eq!(copied, b"image-bytes");
506 }
507
508 #[tokio::test]
512 async fn test_fork_skips_missing_blob_without_failing() {
513 let store = SessionStore::new(make_pool().await);
514 let data_dir = tempfile::tempdir().unwrap();
515 seed_parent(data_dir.path(), &store, "parent").await;
516
517 let parent_dir = crate::session_dir(data_dir.path(), "parent");
518 let parent_log = SessionEventLog::open(&parent_dir).await.unwrap();
519 parent_log
520 .append(
521 None,
522 None,
523 SessionEvent::UserMessage {
524 text: "with missing image".to_owned(),
525 image_refs: vec!["deadbeef".to_owned()],
526 },
527 )
528 .await
529 .unwrap();
530 store.update_seq("parent", 4, 5).await.unwrap();
531
532 let result = ForkEngine::fork(data_dir.path(), "parent", "child", Some(5), &store, None)
533 .await
534 .unwrap();
535 assert_eq!(result.events_copied, 5);
536
537 let child_dir = crate::session_dir(data_dir.path(), "child");
538 assert!(!child_dir.join("blobs").join("deadbeef").exists());
539 }
540
541 #[tokio::test]
545 async fn test_fork_without_image_refs_creates_no_blobs_dir() {
546 let store = SessionStore::new(make_pool().await);
547 let data_dir = tempfile::tempdir().unwrap();
548 seed_parent(data_dir.path(), &store, "parent").await;
549
550 ForkEngine::fork(data_dir.path(), "parent", "child", Some(2), &store, None)
551 .await
552 .unwrap();
553
554 let child_dir = crate::session_dir(data_dir.path(), "child");
555 assert!(!child_dir.join("blobs").exists());
556 }
557
558 #[tokio::test]
563 async fn test_fork_rejects_path_traversal_in_image_refs() {
564 let store = SessionStore::new(make_pool().await);
565 let data_dir = tempfile::tempdir().unwrap();
566 seed_parent(data_dir.path(), &store, "parent").await;
567
568 let parent_dir = crate::session_dir(data_dir.path(), "parent");
569 let parent_log = SessionEventLog::open(&parent_dir).await.unwrap();
570 parent_log
571 .append(
572 None,
573 None,
574 SessionEvent::UserMessage {
575 text: "malicious ref".to_owned(),
576 image_refs: vec!["../../../etc/passwd".to_owned()],
577 },
578 )
579 .await
580 .unwrap();
581 store.update_seq("parent", 4, 5).await.unwrap();
582
583 let err = ForkEngine::fork(data_dir.path(), "parent", "child", Some(5), &store, None)
584 .await
585 .unwrap_err();
586 assert!(matches!(err, SessionError::InvalidBlobHash(_)));
587
588 let child_dir = crate::session_dir(data_dir.path(), "child");
590 assert!(!child_dir.join("blobs").exists());
591 }
592
593 #[tokio::test]
597 async fn test_fork_rejects_absolute_path_in_image_refs() {
598 let store = SessionStore::new(make_pool().await);
599 let data_dir = tempfile::tempdir().unwrap();
600 seed_parent(data_dir.path(), &store, "parent").await;
601
602 let parent_dir = crate::session_dir(data_dir.path(), "parent");
603 let parent_log = SessionEventLog::open(&parent_dir).await.unwrap();
604 parent_log
605 .append(
606 None,
607 None,
608 SessionEvent::UserMessage {
609 text: "malicious absolute ref".to_owned(),
610 image_refs: vec!["/etc/passwd".to_owned()],
611 },
612 )
613 .await
614 .unwrap();
615 store.update_seq("parent", 4, 5).await.unwrap();
616
617 let err = ForkEngine::fork(data_dir.path(), "parent", "child", Some(5), &store, None)
618 .await
619 .unwrap_err();
620 assert!(matches!(err, SessionError::InvalidBlobHash(_)));
621 }
622
623 #[tokio::test]
627 async fn test_fork_dedups_duplicate_blob_hash() {
628 let store = SessionStore::new(make_pool().await);
629 let data_dir = tempfile::tempdir().unwrap();
630 seed_parent(data_dir.path(), &store, "parent").await;
631
632 let parent_dir = crate::session_dir(data_dir.path(), "parent");
633 let parent_blobs = parent_dir.join("blobs");
634 tokio::fs::create_dir_all(&parent_blobs).await.unwrap();
635 tokio::fs::write(parent_blobs.join("cafe01"), b"shared-bytes")
636 .await
637 .unwrap();
638
639 let parent_log = SessionEventLog::open(&parent_dir).await.unwrap();
640 parent_log
641 .append(
642 None,
643 None,
644 SessionEvent::UserMessage {
645 text: "first ref".to_owned(),
646 image_refs: vec!["cafe01".to_owned()],
647 },
648 )
649 .await
650 .unwrap();
651 parent_log
652 .append(
653 None,
654 None,
655 SessionEvent::UserMessage {
656 text: "second ref, same hash".to_owned(),
657 image_refs: vec!["cafe01".to_owned()],
658 },
659 )
660 .await
661 .unwrap();
662 store.update_seq("parent", 4, 6).await.unwrap();
663
664 let result = ForkEngine::fork(data_dir.path(), "parent", "child", Some(6), &store, None)
665 .await
666 .unwrap();
667 assert_eq!(result.events_copied, 6);
668
669 let child_dir = crate::session_dir(data_dir.path(), "child");
670 let child_blob = child_dir.join("blobs").join("cafe01");
671 assert_eq!(tokio::fs::read(&child_blob).await.unwrap(), b"shared-bytes");
672 }
673
674 #[tokio::test]
681 async fn test_copy_referenced_blobs_retry_does_not_truncate_shared_blob() {
682 let data_dir = tempfile::tempdir().unwrap();
683 let src_dir = data_dir.path().join("parent");
684 let child_dir = data_dir.path().join("child");
685
686 let src_blobs = src_dir.join("blobs");
687 tokio::fs::create_dir_all(&src_blobs).await.unwrap();
688 let original_content = b"image-bytes-not-empty";
689 tokio::fs::write(src_blobs.join("a1b2c3"), original_content)
690 .await
691 .unwrap();
692
693 let events = vec![SessionEventEnvelope {
694 seq: 0,
695 ts_ms: 0,
696 turn_id: None,
697 parent_seq: None,
698 kind: SessionEvent::UserMessage {
699 text: "with image".to_owned(),
700 image_refs: vec!["a1b2c3".to_owned()],
701 },
702 }];
703
704 copy_referenced_blobs(&src_dir, &child_dir, &events)
706 .await
707 .unwrap();
708
709 let child_blob = child_dir.join("blobs").join("a1b2c3");
710 assert_eq!(
711 tokio::fs::read(&child_blob).await.unwrap(),
712 original_content
713 );
714
715 copy_referenced_blobs(&src_dir, &child_dir, &events)
718 .await
719 .unwrap();
720
721 assert_eq!(
722 tokio::fs::read(&child_blob).await.unwrap(),
723 original_content,
724 "child blob must not be truncated by a retried fork against the same child_dir"
725 );
726 assert_eq!(
727 tokio::fs::read(src_blobs.join("a1b2c3")).await.unwrap(),
728 original_content,
729 "parent's original blob must not be truncated by a retried fork against the same child_dir"
730 );
731 }
732
733 #[cfg(unix)]
736 #[tokio::test]
737 async fn test_fork_sets_0700_on_child_blobs_dir() {
738 use std::os::unix::fs::PermissionsExt;
739
740 let store = SessionStore::new(make_pool().await);
741 let data_dir = tempfile::tempdir().unwrap();
742 seed_parent(data_dir.path(), &store, "parent").await;
743
744 let parent_dir = crate::session_dir(data_dir.path(), "parent");
745 let parent_blobs = parent_dir.join("blobs");
746 tokio::fs::create_dir_all(&parent_blobs).await.unwrap();
747 tokio::fs::write(parent_blobs.join("a1b2c3"), b"image-bytes")
748 .await
749 .unwrap();
750
751 let parent_log = SessionEventLog::open(&parent_dir).await.unwrap();
752 parent_log
753 .append(
754 None,
755 None,
756 SessionEvent::UserMessage {
757 text: "with image".to_owned(),
758 image_refs: vec!["a1b2c3".to_owned()],
759 },
760 )
761 .await
762 .unwrap();
763 store.update_seq("parent", 4, 5).await.unwrap();
764
765 ForkEngine::fork(data_dir.path(), "parent", "child", Some(5), &store, None)
766 .await
767 .unwrap();
768
769 let child_dir = crate::session_dir(data_dir.path(), "child");
770 let meta = tokio::fs::metadata(child_dir.join("blobs")).await.unwrap();
771 assert_eq!(meta.permissions().mode() & 0o777, 0o700);
772 }
773}