1use 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#[derive(Debug, Clone)]
20pub struct ForkResult {
21 pub new_session_id: String,
23 pub events_copied: usize,
26}
27
28pub struct ForkEngine;
30
31impl ForkEngine {
32 #[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 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 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 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 assert_eq!(result.events_copied, 4);
302 }
303}