Skip to main content

zeph_session/
fork.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! [`ForkEngine`]: eager-copy session forking (spec §7).
5//!
6//! Copy-on-write forking is explicitly deferred (spec §7.2, §15 NEVER) — eager copy is simple and
7//! self-contained for MVP, and robust to either side independently condensing the shared prefix
8//! afterward (the child log is fully self-contained; `forked_at_seq` is historical metadata only).
9
10use std::path::Path;
11
12use crate::error::SessionError;
13use crate::event::SessionEvent;
14use crate::log::SessionEventLog;
15use crate::replay::ReplayEngine;
16use crate::store::SessionStore;
17
18/// The result of a successful fork.
19#[derive(Debug, Clone)]
20pub struct ForkResult {
21    /// The newly allocated child session id.
22    pub new_session_id: String,
23    /// Number of events copied from the parent's log (excludes the child's own `SessionStarted`
24    /// header, which is synthesized fresh).
25    pub events_copied: usize,
26}
27
28/// Forks a session at a given `seq`, producing a new, fully self-contained child session.
29pub struct ForkEngine;
30
31impl ForkEngine {
32    /// Fork `src_id` at `at_seq` into a caller-allocated `new_id` (`at_seq` is an exclusive upper
33    /// bound — matches [`ReplayEngine::replay`]'s `up_to` semantics: the child receives events
34    /// `[0, at_seq)` from the parent, plus a synthetic `SessionStarted` header recording
35    /// `forked_from`). `at_seq = None` forks at the current end of the log (copies everything) —
36    /// the default for callers with no explicit cut point (ACP's `fork_session`, which has no
37    /// `seq` parameter, and the CLI's optional `--at`).
38    ///
39    /// `new_id` is caller-supplied rather than minted internally: callers such as ACP's
40    /// `do_fork_session` need the id before the fork call completes (to construct the session's
41    /// `LoopbackChannel`/entry), and the CLI mints a fresh `SessionId::generate()` before calling
42    /// in.
43    ///
44    /// # Errors
45    ///
46    /// Returns [`SessionError::NotFound`] if `src_id` has no session-store row,
47    /// [`SessionError::InvalidForkPoint`] if `at_seq` exceeds the parent log's event count, or
48    /// [`SessionError::Io`]/[`SessionError::Db`] if the copy or store update fails.
49    #[tracing::instrument(name = "session.fork.run", skip_all, level = "info", fields(at_seq))]
50    pub async fn fork(
51        data_dir: &Path,
52        src_id: &str,
53        new_id: &str,
54        at_seq: Option<u64>,
55        store: &SessionStore,
56    ) -> Result<ForkResult, SessionError> {
57        if store.get(src_id).await?.is_none() {
58            return Err(SessionError::NotFound(src_id.to_owned()));
59        }
60
61        let src_dir = crate::session_dir(data_dir, src_id);
62        let src_log = SessionEventLog::open(&src_dir).await?;
63        let all_events = src_log.read_all().await?;
64
65        let total = u64::try_from(all_events.len()).unwrap_or(u64::MAX);
66        let at_seq = at_seq.unwrap_or(total);
67        if at_seq > total {
68            return Err(SessionError::InvalidForkPoint(format!(
69                "at_seq={at_seq} exceeds source session's event count={total}"
70            )));
71        }
72
73        // Validate the cut point is internally consistent (spec §7.2 step 2) — replay must not
74        // error. The reconstructed state itself is not needed further here.
75        ReplayEngine::replay(&src_dir, Some(at_seq)).await?;
76
77        let take_n = usize::try_from(at_seq).unwrap_or(usize::MAX);
78        let to_copy: Vec<_> = all_events.iter().take(take_n).cloned().collect();
79        let (cwd, provider_name, model) = to_copy
80            .iter()
81            .find_map(|e| match &e.kind {
82                SessionEvent::SessionStarted {
83                    cwd,
84                    provider_name,
85                    model,
86                    ..
87                } => Some((cwd.clone(), provider_name.clone(), model.clone())),
88                _ => None,
89            })
90            .unwrap_or_default();
91
92        let child_dir = crate::session_dir(data_dir, new_id);
93        let child_log = SessionEventLog::open(&child_dir).await?;
94
95        child_log
96            .append(
97                None,
98                None,
99                SessionEvent::SessionStarted {
100                    session_id: new_id.to_owned(),
101                    cwd,
102                    provider_name,
103                    model,
104                    forked_from: Some((src_id.to_owned(), at_seq)),
105                },
106            )
107            .await?;
108        for envelope in &to_copy {
109            child_log
110                .append(envelope.turn_id, envelope.parent_seq, envelope.kind.clone())
111                .await?;
112        }
113
114        store.record_fork(new_id, src_id, at_seq).await?;
115        store
116            .update_seq(
117                new_id,
118                child_log.last_seq().unwrap_or(0),
119                to_copy.len() as u64 + 1,
120            )
121            .await?;
122
123        // Non-destructive provenance record on the parent (spec §7.2 step 8).
124        src_log
125            .append(
126                None,
127                None,
128                SessionEvent::ForkPoint {
129                    new_session_id: new_id.to_owned(),
130                },
131            )
132            .await?;
133
134        Ok(ForkResult {
135            new_session_id: new_id.to_owned(),
136            events_copied: to_copy.len(),
137        })
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::store::SessionStore;
145
146    async fn make_pool() -> zeph_db::DbPool {
147        let config = zeph_db::DbConfig {
148            url: ":memory:".to_owned(),
149            ..Default::default()
150        };
151        let pool = config
152            .connect()
153            .await
154            .expect("connect in-memory sqlite pool");
155        zeph_db::run_migrations(&pool)
156            .await
157            .expect("run migrations");
158        pool
159    }
160
161    async fn seed_parent(data_dir: &Path, store: &SessionStore, id: &str) {
162        store.create(id).await.unwrap();
163        let dir = crate::session_dir(data_dir, id);
164        let log = SessionEventLog::open(&dir).await.unwrap();
165        log.append(
166            None,
167            None,
168            SessionEvent::SessionStarted {
169                session_id: id.to_owned(),
170                cwd: "/repo".to_owned(),
171                provider_name: "claude".to_owned(),
172                model: "opus".to_owned(),
173                forked_from: None,
174            },
175        )
176        .await
177        .unwrap();
178        log.append(
179            None,
180            None,
181            SessionEvent::UserMessage {
182                text: "hello".to_owned(),
183                image_refs: vec![],
184            },
185        )
186        .await
187        .unwrap();
188        log.append(
189            None,
190            None,
191            SessionEvent::AssistantMessage {
192                parts: vec![zeph_llm::provider::MessagePart::Text {
193                    text: "hi".to_owned(),
194                }],
195            },
196        )
197        .await
198        .unwrap();
199        log.append(
200            None,
201            None,
202            SessionEvent::UserMessage {
203                text: "second turn".to_owned(),
204                image_refs: vec![],
205            },
206        )
207        .await
208        .unwrap();
209        store
210            .update_seq(id, log.last_seq().unwrap(), 4)
211            .await
212            .unwrap();
213    }
214
215    #[tokio::test]
216    async fn test_fork_copies_events() {
217        let store = SessionStore::new(make_pool().await);
218        let data_dir = tempfile::tempdir().unwrap();
219        seed_parent(data_dir.path(), &store, "parent").await;
220
221        let result = ForkEngine::fork(data_dir.path(), "parent", "child", Some(3), &store)
222            .await
223            .unwrap();
224        assert_eq!(result.events_copied, 3);
225        assert_eq!(result.new_session_id, "child");
226
227        let child_dir = crate::session_dir(data_dir.path(), &result.new_session_id);
228        let child_log = SessionEventLog::open(&child_dir).await.unwrap();
229        let events = child_log.read_all().await.unwrap();
230        // 1 synthesized SessionStarted header + 3 copied events.
231        assert_eq!(events.len(), 4);
232    }
233
234    #[tokio::test]
235    async fn test_fork_provenance_metadata() {
236        let store = SessionStore::new(make_pool().await);
237        let data_dir = tempfile::tempdir().unwrap();
238        seed_parent(data_dir.path(), &store, "parent").await;
239
240        ForkEngine::fork(data_dir.path(), "parent", "child", Some(2), &store)
241            .await
242            .unwrap();
243
244        let meta = store.get("child").await.unwrap().unwrap();
245        assert_eq!(meta.forked_from.as_deref(), Some("parent"));
246        assert_eq!(meta.forked_at_seq, Some(2));
247    }
248
249    #[tokio::test]
250    async fn test_fork_appends_forkpoint_to_parent() {
251        let store = SessionStore::new(make_pool().await);
252        let data_dir = tempfile::tempdir().unwrap();
253        seed_parent(data_dir.path(), &store, "parent").await;
254
255        ForkEngine::fork(data_dir.path(), "parent", "child", Some(2), &store)
256            .await
257            .unwrap();
258
259        let parent_dir = crate::session_dir(data_dir.path(), "parent");
260        let parent_log = SessionEventLog::open(&parent_dir).await.unwrap();
261        let events = parent_log.read_all().await.unwrap();
262        assert!(matches!(
263            events.last().unwrap().kind,
264            SessionEvent::ForkPoint { .. }
265        ));
266    }
267
268    #[tokio::test]
269    async fn test_fork_rejects_seq_beyond_source() {
270        let store = SessionStore::new(make_pool().await);
271        let data_dir = tempfile::tempdir().unwrap();
272        seed_parent(data_dir.path(), &store, "parent").await;
273
274        let err = ForkEngine::fork(data_dir.path(), "parent", "child", Some(100), &store)
275            .await
276            .unwrap_err();
277        assert!(matches!(err, SessionError::InvalidForkPoint(_)));
278    }
279
280    #[tokio::test]
281    async fn test_fork_rejects_unknown_source() {
282        let store = SessionStore::new(make_pool().await);
283        let data_dir = tempfile::tempdir().unwrap();
284
285        let err = ForkEngine::fork(data_dir.path(), "no-such", "child", Some(0), &store)
286            .await
287            .unwrap_err();
288        assert!(matches!(err, SessionError::NotFound(_)));
289    }
290
291    #[tokio::test]
292    async fn test_fork_none_copies_everything() {
293        let store = SessionStore::new(make_pool().await);
294        let data_dir = tempfile::tempdir().unwrap();
295        seed_parent(data_dir.path(), &store, "parent").await;
296
297        let result = ForkEngine::fork(data_dir.path(), "parent", "child", None, &store)
298            .await
299            .unwrap();
300        // seed_parent appends 4 events total.
301        assert_eq!(result.events_copied, 4);
302    }
303}