vtcode_core/scheduler/
store.rs

1use anyhow::{Context, Result, bail};
2use chrono::{DateTime, Local, Utc};
3use std::fs;
4use std::io::Write;
5use std::path::{Path, PathBuf};
6use std::time::Duration;
7
8use super::session::ScheduleCreateInput;
9use super::spec::{ScheduledTaskDefinition, ScheduledTaskRecord, ScheduledTaskRuntimeState, ScheduledTaskSummary};
10use super::{MAX_SCHEDULED_TASKS, atomic_write, initialize_runtime_state};
11use vtcode_commons::VtCodePaths;
12
13const CLAIM_STALE_SECS: u64 = 15 * 60;
14const DURABLE_STORE_DIR: &str = "scheduled_tasks";
15
16#[derive(Debug, Clone)]
17pub struct SchedulerPaths {
18    pub config_root: PathBuf,
19    pub data_root: PathBuf,
20}
21
22impl SchedulerPaths {
23    pub fn new_default() -> Result<Self> {
24        let paths = VtCodePaths::resolve().context("Failed to resolve VT Code paths")?;
25        let config_root = paths.config_path(DURABLE_STORE_DIR)?;
26        let data_root = paths.state_path(DURABLE_STORE_DIR)?;
27        Ok(Self { config_root, data_root })
28    }
29
30    #[must_use]
31    pub fn tasks_dir(&self) -> PathBuf {
32        self.config_root.join("tasks")
33    }
34
35    #[must_use]
36    pub fn state_dir(&self) -> PathBuf {
37        self.data_root.join("state")
38    }
39
40    #[must_use]
41    pub fn claims_dir(&self) -> PathBuf {
42        self.data_root.join("claims")
43    }
44
45    #[must_use]
46    pub fn runs_dir(&self) -> PathBuf {
47        self.data_root.join("runs")
48    }
49
50    pub fn ensure_dirs(&self) -> Result<()> {
51        for dir in [
52            &self.config_root,
53            &self.data_root,
54            &self.tasks_dir(),
55            &self.state_dir(),
56            &self.claims_dir(),
57            &self.runs_dir(),
58        ] {
59            VtCodePaths::ensure_user_dir(dir)
60                .with_context(|| format!("Failed to create scheduler directory {}", dir.display()))?;
61        }
62        Ok(())
63    }
64
65    #[must_use]
66    pub fn definition_path(&self, id: &str) -> PathBuf {
67        self.tasks_dir().join(format!("{id}.toml"))
68    }
69
70    #[must_use]
71    pub fn runtime_path(&self, id: &str) -> PathBuf {
72        self.state_dir().join(format!("{id}.json"))
73    }
74
75    #[must_use]
76    pub fn claim_path(&self, id: &str) -> PathBuf {
77        self.claims_dir().join(format!("{id}.claim"))
78    }
79
80    #[must_use]
81    pub fn artifact_dir(&self, id: &str, run_at: DateTime<Utc>) -> PathBuf {
82        self.runs_dir().join(id).join(run_at.format("%Y%m%dT%H%M%SZ").to_string())
83    }
84}
85
86#[derive(Debug, Clone)]
87pub struct DurableTaskStore {
88    paths: SchedulerPaths,
89}
90
91impl DurableTaskStore {
92    #[must_use]
93    pub fn with_paths(paths: SchedulerPaths) -> Self {
94        Self { paths }
95    }
96
97    pub fn new_default() -> Result<Self> {
98        Ok(Self { paths: SchedulerPaths::new_default()? })
99    }
100
101    #[must_use]
102    pub fn paths(&self) -> &SchedulerPaths {
103        &self.paths
104    }
105
106    pub fn create(&self, definition: ScheduledTaskDefinition) -> Result<ScheduledTaskSummary> {
107        self.paths.ensure_dirs()?;
108        let current_count = self.definition_paths()?.len();
109        if current_count >= MAX_SCHEDULED_TASKS {
110            bail!("VT Code supports at most {MAX_SCHEDULED_TASKS} durable scheduled tasks");
111        }
112
113        let runtime = initialize_runtime_state(&definition)?;
114        self.write_definition(&definition)?;
115        self.write_runtime(&definition.id, &runtime)?;
116        Ok(ScheduledTaskRecord { definition, runtime }.summary())
117    }
118
119    pub fn create_from_input(
120        &self,
121        input: ScheduleCreateInput,
122        now: DateTime<Local>,
123        default_workspace: Option<PathBuf>,
124    ) -> Result<ScheduledTaskSummary> {
125        let definition = input.build_definition(now, default_workspace)?;
126        self.create(definition)
127    }
128
129    pub fn list(&self) -> Result<Vec<ScheduledTaskSummary>> {
130        let mut records = self.load_records()?;
131        records.sort_by_key(|record| record.runtime.next_run_at);
132        Ok(records.into_iter().map(|record| record.summary()).collect())
133    }
134
135    pub fn delete(&self, id: &str) -> Result<Option<ScheduledTaskSummary>> {
136        let Some(record) = self.load_record(id)? else {
137            return Ok(None);
138        };
139        let _ = fs::remove_file(self.paths.definition_path(id));
140        let _ = fs::remove_file(self.paths.runtime_path(id));
141        let _ = fs::remove_file(self.paths.claim_path(id));
142        Ok(Some(record.summary()))
143    }
144
145    pub fn load_record(&self, id: &str) -> Result<Option<ScheduledTaskRecord>> {
146        let definition_path = self.paths.definition_path(id);
147        if !definition_path.exists() {
148            return Ok(None);
149        }
150        let definition = read_definition(&definition_path)?;
151        let runtime = match self.read_runtime(&definition.id)? {
152            Some(runtime) => runtime,
153            None => initialize_runtime_state(&definition)?,
154        };
155        Ok(Some(ScheduledTaskRecord { definition, runtime }))
156    }
157
158    pub fn update_runtime(&self, record: &ScheduledTaskRecord) -> Result<()> {
159        self.write_runtime(&record.definition.id, &record.runtime)
160    }
161
162    pub(crate) async fn load_records_async(&self) -> Result<Vec<ScheduledTaskRecord>> {
163        let store = self.clone();
164        tokio::task::spawn_blocking(move || store.load_records())
165            .await
166            .context("Scheduled task loading task panicked")?
167    }
168
169    pub(crate) async fn update_runtime_async(&self, record: ScheduledTaskRecord) -> Result<()> {
170        let store = self.clone();
171        tokio::task::spawn_blocking(move || store.update_runtime(&record))
172            .await
173            .context("Scheduled task runtime update task panicked")?
174    }
175
176    fn load_records(&self) -> Result<Vec<ScheduledTaskRecord>> {
177        self.paths.ensure_dirs()?;
178        let mut records = Vec::new();
179        for definition_path in self.definition_paths()? {
180            let definition = read_definition(&definition_path)?;
181            let runtime = self
182                .read_runtime(&definition.id)?
183                .unwrap_or(initialize_runtime_state(&definition)?);
184            records.push(ScheduledTaskRecord { definition, runtime });
185        }
186        Ok(records)
187    }
188
189    fn definition_paths(&self) -> Result<Vec<PathBuf>> {
190        self.paths.ensure_dirs()?;
191        let mut paths = Vec::new();
192        for entry in fs::read_dir(self.paths.tasks_dir())
193            .with_context(|| format!("Failed to read {}", self.paths.tasks_dir().display()))?
194        {
195            let entry = entry?;
196            let path = entry.path();
197            if path.extension().and_then(|value| value.to_str()) == Some("toml") {
198                paths.push(path);
199            }
200        }
201        paths.sort();
202        Ok(paths)
203    }
204
205    fn write_definition(&self, definition: &ScheduledTaskDefinition) -> Result<()> {
206        let serialized = toml::to_string_pretty(definition).context("Failed to serialize task definition")?;
207        atomic_write(&self.paths.definition_path(&definition.id), serialized.as_bytes())
208    }
209
210    fn read_runtime(&self, id: &str) -> Result<Option<ScheduledTaskRuntimeState>> {
211        let path = self.paths.runtime_path(id);
212        if !path.exists() {
213            return Ok(None);
214        }
215        let raw = fs::read_to_string(&path).with_context(|| format!("Failed to read {}", path.display()))?;
216        let runtime = serde_json::from_str(&raw).with_context(|| format!("Failed to parse {}", path.display()))?;
217        Ok(Some(runtime))
218    }
219
220    fn write_runtime(&self, id: &str, runtime: &ScheduledTaskRuntimeState) -> Result<()> {
221        let serialized = serde_json::to_vec_pretty(runtime).context("Failed to serialize runtime state")?;
222        atomic_write(&self.paths.runtime_path(id), &serialized)
223    }
224}
225
226fn read_definition(path: &Path) -> Result<ScheduledTaskDefinition> {
227    let raw = fs::read_to_string(path).with_context(|| format!("Failed to read {}", path.display()))?;
228    toml::from_str(&raw).with_context(|| format!("Failed to parse {}", path.display()))
229}
230
231fn try_acquire_claim(paths: &SchedulerPaths, id: &str) -> Result<bool> {
232    paths.ensure_dirs()?;
233    let path = paths.claim_path(id);
234    match VtCodePaths::create_private_file(&path) {
235        Ok(mut file) => {
236            let timestamp = Utc::now().to_rfc3339();
237            file.write_all(timestamp.as_bytes())
238                .with_context(|| format!("Failed to write {}", path.display()))?;
239            Ok(true)
240        }
241        Err(error)
242            if error.chain().any(|cause| {
243                cause
244                    .downcast_ref::<std::io::Error>()
245                    .is_some_and(|io_error| io_error.kind() == std::io::ErrorKind::AlreadyExists)
246            }) =>
247        {
248            if claim_is_stale(&path)? {
249                let _ = fs::remove_file(&path);
250                return try_acquire_claim(paths, id);
251            }
252            Ok(false)
253        }
254        Err(error) => Err(error).with_context(|| format!("Failed to create {}", path.display())),
255    }
256}
257
258pub(crate) async fn try_acquire_claim_async(paths: SchedulerPaths, id: String) -> Result<bool> {
259    tokio::task::spawn_blocking(move || try_acquire_claim(&paths, &id))
260        .await
261        .context("Scheduled task claim acquisition task panicked")?
262}
263
264fn claim_is_stale(path: &Path) -> Result<bool> {
265    let metadata = fs::metadata(path).with_context(|| format!("Failed to stat {}", path.display()))?;
266    let modified = metadata
267        .modified()
268        .with_context(|| format!("Failed to read modification time for {}", path.display()))?;
269    let elapsed = modified.elapsed().unwrap_or_default();
270    Ok(elapsed >= Duration::from_secs(CLAIM_STALE_SECS))
271}
272
273fn release_claim(paths: &SchedulerPaths, id: &str) -> Result<()> {
274    let path = paths.claim_path(id);
275    if path.exists() {
276        fs::remove_file(&path).with_context(|| format!("Failed to remove {}", path.display()))?;
277    }
278    Ok(())
279}
280
281pub(crate) async fn release_claim_async(paths: SchedulerPaths, id: String) -> Result<()> {
282    tokio::task::spawn_blocking(move || release_claim(&paths, &id))
283        .await
284        .context("Scheduled task claim release task panicked")?
285}
vtcode_core/scheduler/store.rs

vtcode_core/scheduler/
store.rs
