codex_rollout/
rollout_reference_index.rs

1//! Indexes direct fork references found in local rollout files.
2
3use std::collections::HashMap;
4use std::collections::HashSet;
5use std::io;
6use std::path::Path;
7use std::time::Duration;
8use std::time::Instant;
9
10use codex_protocol::ThreadId;
11use codex_protocol::protocol::HistoryPosition;
12
13use crate::ARCHIVED_SESSIONS_SUBDIR;
14use crate::SESSIONS_SUBDIR;
15use crate::compression::RolloutFile;
16
17/// Direct history-base edges discovered from local rollout metadata.
18///
19/// This is a physical-rollout index, not a logical lineage resolver. Callers use it to answer
20/// cheap inverse-reference questions without each reimplementing rollout discovery.
21#[derive(Debug, Default)]
22pub struct RolloutReferenceIndex {
23    history_base_by_thread: HashMap<ThreadId, HistoryPosition>,
24    reference_counts_by_thread: HashMap<ThreadId, usize>,
25}
26
27impl RolloutReferenceIndex {
28    /// Scans active and archived local rollout metadata without a deadline.
29    pub async fn scan(codex_home: &Path) -> io::Result<Self> {
30        let Some(index) = Self::scan_with_deadline(codex_home, ScanDeadline::Unlimited).await?
31        else {
32            return Err(io::Error::other(
33                "unlimited rollout reference scan exceeded a deadline",
34            ));
35        };
36        Ok(index)
37    }
38
39    /// Scans active and archived local rollout metadata until the worker deadline expires.
40    ///
41    /// Returns None instead of a partial index when the deadline expires.
42    pub(crate) async fn scan_until(
43        codex_home: &Path,
44        started_at: Instant,
45        max_runtime: Duration,
46    ) -> io::Result<Option<Self>> {
47        Self::scan_with_deadline(
48            codex_home,
49            ScanDeadline::Until {
50                started_at,
51                max_runtime,
52            },
53        )
54        .await
55    }
56
57    /// Returns how many other discovered rollouts directly reference thread_id.
58    pub fn reference_count(&self, thread_id: ThreadId) -> usize {
59        self.reference_counts_by_thread
60            .get(&thread_id)
61            .copied()
62            .unwrap_or_default()
63    }
64
65    /// Returns the direct history-base edge for thread_id, if one was discovered.
66    pub fn history_base(&self, thread_id: ThreadId) -> Option<&HistoryPosition> {
67        self.history_base_by_thread.get(&thread_id)
68    }
69
70    async fn scan_with_deadline(
71        codex_home: &Path,
72        deadline: ScanDeadline,
73    ) -> io::Result<Option<Self>> {
74        let mut history_base_by_thread = HashMap::new();
75        let mut seen_thread_ids = HashSet::new();
76        let mut stack = vec![
77            codex_home.join(ARCHIVED_SESSIONS_SUBDIR),
78            codex_home.join(SESSIONS_SUBDIR),
79        ];
80        while let Some(directory) = stack.pop() {
81            if deadline.expired() {
82                return Ok(None);
83            }
84            let mut entries = match tokio::fs::read_dir(directory.as_path()).await {
85                Ok(entries) => entries,
86                Err(err) if err.kind() == io::ErrorKind::NotFound => continue,
87                Err(err) => return Err(err),
88            };
89            loop {
90                if deadline.expired() {
91                    return Ok(None);
92                }
93                let Some(entry) = entries.next_entry().await? else {
94                    break;
95                };
96                let path = entry.path();
97                let file_type = entry.file_type().await?;
98                if file_type.is_dir() {
99                    stack.push(path);
100                    continue;
101                }
102                if !file_type.is_file() {
103                    continue;
104                }
105                let Some(rollout_file) = RolloutFile::from_path(path) else {
106                    continue;
107                };
108                let Ok(meta) = crate::read_session_meta_line(rollout_file.path()).await else {
109                    continue;
110                };
111                let thread_id = meta.meta.id;
112                if !seen_thread_ids.insert(thread_id) {
113                    continue;
114                }
115                if let Some(history_base) = meta.meta.history_base {
116                    history_base_by_thread.insert(thread_id, history_base);
117                }
118            }
119        }
120
121        let mut reference_counts_by_thread = HashMap::new();
122        for (thread_id, history_base) in &history_base_by_thread {
123            if history_base.thread_id == *thread_id {
124                continue;
125            }
126            *reference_counts_by_thread
127                .entry(history_base.thread_id)
128                .or_default() += 1;
129        }
130        Ok(Some(Self {
131            history_base_by_thread,
132            reference_counts_by_thread,
133        }))
134    }
135}
136
137#[derive(Clone, Copy)]
138enum ScanDeadline {
139    Unlimited,
140    Until {
141        started_at: Instant,
142        max_runtime: Duration,
143    },
144}
145
146impl ScanDeadline {
147    fn expired(self) -> bool {
148        match self {
149            Self::Unlimited => false,
150            Self::Until {
151                started_at,
152                max_runtime,
153            } => started_at.elapsed() >= max_runtime,
154        }
155    }
156}
157
158#[cfg(test)]
159#[path = "rollout_reference_index_tests.rs"]
160mod tests;
codex_rollout/rollout_reference_index.rs

codex_rollout/
rollout_reference_index.rs
