Skip to main content

vtcode_core/scheduler/
session.rs

1use anyhow::{Context, Result, anyhow, bail};
2use chrono::{DateTime, Duration as ChronoDuration, Local, LocalResult, NaiveDateTime, NaiveTime, TimeZone, Utc};
3use humantime::parse_duration as parse_human_duration;
4use regex::Regex;
5use std::collections::BTreeMap;
6use std::path::{Path, PathBuf};
7use std::sync::LazyLock;
8
9use super::spec::{
10    DueSessionPrompt, ScheduleSpec, ScheduledTaskAction, ScheduledTaskDefinition, ScheduledTaskRecord,
11    ScheduledTaskSummary, TaskRunStatus,
12};
13use super::{
14    MAX_SCHEDULED_TASKS, SESSION_TASK_EXPIRY_HOURS, advance_record_runtime, due_session_prompt,
15    initialize_runtime_state,
16};
17use crate::utils::path::normalize_path;
18
19static REMIND_AT_RE: LazyLock<Regex> = LazyLock::new(|| {
20    Regex::new(r"(?ix)^\s*remind\s+me\s+at\s+(?P<when>.+?)\s+to\s+(?P<prompt>.+)\s*$").expect("remind-at regex")
21});
22
23static REMIND_IN_RE: LazyLock<Regex> = LazyLock::new(|| {
24    Regex::new(
25        r"(?ix)
26        ^\s*in\s+
27        (?P<count>\d+)
28        \s*
29        (?P<unit>minutes|minute|hours|hour|days|day)
30        \s*,?\s*
31        (?P<prompt>.+)
32        \s*$
33    ",
34    )
35    .expect("remind-in regex")
36});
37
38static TIME_ONLY_RE: LazyLock<Regex> = LazyLock::new(|| {
39    Regex::new(r"(?ix)^\s*(?P<hour>\d{1,2})(?::(?P<minute>\d{2}))?\s*(?P<ampm>am|pm)?\s*$").expect("time-only regex")
40});
41
42#[derive(Debug, Default, Clone)]
43pub struct SessionScheduler {
44    pub(crate) tasks: BTreeMap<String, ScheduledTaskRecord>,
45}
46
47impl SessionScheduler {
48    #[must_use]
49    pub fn new() -> Self {
50        Self { tasks: BTreeMap::new() }
51    }
52
53    #[must_use]
54    pub fn len(&self) -> usize {
55        self.tasks.len()
56    }
57
58    #[must_use]
59    pub fn is_empty(&self) -> bool {
60        self.tasks.is_empty()
61    }
62
63    pub fn create_prompt_task(
64        &mut self,
65        name: Option<String>,
66        prompt: String,
67        schedule: ScheduleSpec,
68        created_at: DateTime<Utc>,
69    ) -> Result<ScheduledTaskSummary> {
70        self.ensure_capacity()?;
71        let expires_at = schedule
72            .is_recurring()
73            .then(|| created_at + ChronoDuration::hours(SESSION_TASK_EXPIRY_HOURS));
74        let definition = ScheduledTaskDefinition::new(
75            name,
76            schedule,
77            ScheduledTaskAction::Prompt { prompt },
78            None,
79            created_at,
80            expires_at,
81        )?;
82        let runtime = initialize_runtime_state(&definition)?;
83        let record = ScheduledTaskRecord { definition: definition.clone(), runtime };
84        let summary = record.summary();
85        self.tasks.insert(definition.id.clone(), record);
86        Ok(summary)
87    }
88
89    pub fn list(&self) -> Vec<ScheduledTaskSummary> {
90        self.tasks.values().map(ScheduledTaskRecord::summary).collect()
91    }
92
93    pub fn delete(&mut self, query: &str) -> Option<ScheduledTaskSummary> {
94        let query = query.trim();
95        if query.is_empty() {
96            return None;
97        }
98        if let Some(record) = self.tasks.remove(query) {
99            return Some(record.summary());
100        }
101        let key = self
102            .tasks
103            .iter()
104            .find_map(|(id, record)| record.definition.name.eq_ignore_ascii_case(query).then(|| id.clone()))?;
105        self.tasks.remove(&key).map(|record| record.summary())
106    }
107
108    pub fn collect_due_prompts(&mut self, now: DateTime<Utc>) -> Result<Vec<DueSessionPrompt>> {
109        let mut due = Vec::new();
110        let mut completed = Vec::new();
111        for record in self.tasks.values_mut() {
112            let Some(next_run_at) = record.runtime.next_run_at else {
113                continue;
114            };
115            if now < next_run_at {
116                continue;
117            }
118
119            if let Some(prompt) = due_session_prompt(record) {
120                due.push(prompt);
121            }
122
123            if advance_record_runtime(record, now, TaskRunStatus::Triggered)? {
124                completed.push(record.definition.id.clone());
125            }
126        }
127
128        for id in completed {
129            self.tasks.remove(&id);
130        }
131
132        Ok(due)
133    }
134
135    fn ensure_capacity(&self) -> Result<()> {
136        if self.tasks.len() >= MAX_SCHEDULED_TASKS {
137            bail!("A session can hold at most {MAX_SCHEDULED_TASKS} scheduled tasks");
138        }
139        Ok(())
140    }
141}
142
143#[derive(Debug, Clone, PartialEq, Eq)]
144pub enum SessionLanguageCommand {
145    CreateOneShotPrompt { prompt: String, run_at: DateTime<Utc> },
146    ListTasks,
147    CancelTask { query: String },
148}
149
150#[derive(Debug, Clone, PartialEq, Eq)]
151pub struct ScheduleCreateInput {
152    pub name: Option<String>,
153    pub prompt: Option<String>,
154    pub reminder: Option<String>,
155    pub every: Option<String>,
156    pub cron: Option<String>,
157    pub at: Option<String>,
158    pub workspace: Option<PathBuf>,
159}
160
161impl ScheduleCreateInput {
162    pub fn build_definition(
163        self,
164        now: DateTime<Local>,
165        default_workspace: Option<PathBuf>,
166    ) -> Result<ScheduledTaskDefinition> {
167        let action = match (self.prompt, self.reminder) {
168            (Some(prompt), None) => ScheduledTaskAction::Prompt { prompt },
169            (None, Some(reminder)) => ScheduledTaskAction::Reminder { message: reminder },
170            (Some(_), Some(_)) => bail!("Choose either --prompt or --reminder"),
171            (None, None) => bail!("One of --prompt or --reminder is required"),
172        };
173
174        let schedule = match (self.every, self.cron, self.at) {
175            (Some(raw), None, None) => {
176                let duration = parse_human_duration(raw.trim())
177                    .with_context(|| format!("Invalid --every duration: {}", raw.trim()))?;
178                ScheduleSpec::fixed_interval(duration)?
179            }
180            (None, Some(expression), None) => ScheduleSpec::cron5(expression)?,
181            (None, None, Some(raw)) => ScheduleSpec::one_shot(parse_local_datetime(raw.trim(), now)?),
182            _ => bail!("Choose exactly one of --every, --cron, or --at"),
183        };
184
185        let workspace = match (&action, self.workspace.or(default_workspace)) {
186            (ScheduledTaskAction::Prompt { .. }, Some(path)) => Some(path),
187            (ScheduledTaskAction::Prompt { .. }, None) => {
188                bail!("Prompt tasks require a workspace (pass --workspace or create from chat)")
189            }
190            (ScheduledTaskAction::Reminder { .. }, path) => path,
191        }
192        .map(|path| resolve_scheduled_workspace_path(&path))
193        .transpose()?;
194
195        if matches!(action, ScheduledTaskAction::Prompt { .. }) && workspace.as_ref().is_some_and(|path| !path.is_dir())
196        {
197            let workspace = workspace.as_ref().context("prompt workspace should exist")?;
198            bail!("Prompt task workspace does not exist or is not a directory: {}", workspace.display());
199        }
200
201        ScheduledTaskDefinition::new(self.name, schedule, action, workspace, now.with_timezone(&Utc), None)
202    }
203}
204
205pub(crate) fn resolve_scheduled_workspace_path(path: &Path) -> Result<PathBuf> {
206    resolve_scheduled_workspace_path_with_home(path, dirs::home_dir().as_deref())
207}
208
209pub(crate) fn resolve_scheduled_workspace_path_with_home(path: &Path, home_dir: Option<&Path>) -> Result<PathBuf> {
210    let expanded = expand_scheduled_workspace_home(path, home_dir);
211    let absolute = if expanded.is_absolute() {
212        expanded
213    } else {
214        std::env::current_dir()
215            .context("Failed to resolve current directory for scheduled task workspace")?
216            .join(expanded)
217    };
218    Ok(normalize_path(&absolute))
219}
220
221fn expand_scheduled_workspace_home(path: &Path, home_dir: Option<&Path>) -> PathBuf {
222    let Some(raw) = path.to_str() else {
223        return path.to_path_buf();
224    };
225
226    if raw == "~" {
227        return home_dir.map(Path::to_path_buf).unwrap_or_else(|| path.to_path_buf());
228    }
229
230    if let Some(rest) = raw.strip_prefix("~/")
231        && let Some(home_dir) = home_dir
232    {
233        return home_dir.join(rest);
234    }
235
236    path.to_path_buf()
237}
238
239pub fn parse_session_language_command(input: &str, now: DateTime<Local>) -> Option<Result<SessionLanguageCommand>> {
240    let trimmed = input.trim();
241    if trimmed.is_empty() {
242        return None;
243    }
244
245    let normalized = trimmed.trim_end_matches(['?', '.', '!']).to_ascii_lowercase();
246    if normalized == "what scheduled tasks do i have" {
247        return Some(Ok(SessionLanguageCommand::ListTasks));
248    }
249
250    if let Some(captures) = REMIND_AT_RE.captures(trimmed) {
251        let when = captures.name("when").map(|value| value.as_str()).unwrap_or_default();
252        let prompt = captures
253            .name("prompt")
254            .map(|value| value.as_str().trim().to_string())
255            .unwrap_or_default();
256        return Some(
257            parse_local_datetime(when, now)
258                .map(|run_at| SessionLanguageCommand::CreateOneShotPrompt { prompt, run_at }),
259        );
260    }
261
262    if let Some(captures) = REMIND_IN_RE.captures(trimmed) {
263        let count = captures["count"].parse::<i64>().ok()?;
264        let prompt = captures["prompt"].trim().to_string();
265        let delta = match captures["unit"].to_ascii_lowercase().as_str() {
266            "minute" | "minutes" => ChronoDuration::minutes(count),
267            "hour" | "hours" => ChronoDuration::hours(count),
268            "day" | "days" => ChronoDuration::days(count),
269            _ => return None,
270        };
271        return Some(Ok(SessionLanguageCommand::CreateOneShotPrompt {
272            prompt,
273            run_at: (now + delta).with_timezone(&Utc),
274        }));
275    }
276
277    if let Some(query) = trimmed.strip_prefix("cancel ") {
278        return Some(Ok(SessionLanguageCommand::CancelTask { query: query.trim().to_string() }));
279    }
280
281    None
282}
283
284pub fn parse_schedule_create_args(args: &str) -> Result<ScheduleCreateInput> {
285    let tokens = shell_words::split(args).with_context(|| format!("Failed to parse arguments: {args}"))?;
286    parse_schedule_create_tokens(&tokens)
287}
288
289pub fn parse_schedule_create_tokens(tokens: &[String]) -> Result<ScheduleCreateInput> {
290    let mut name = None;
291    let mut prompt = None;
292    let mut reminder = None;
293    let mut every = None;
294    let mut cron = None;
295    let mut at = None;
296    let mut workspace = None;
297    let mut index = 0usize;
298
299    while index < tokens.len() {
300        let token = &tokens[index];
301        let (flag, inline_value) = if let Some((left, right)) = token.split_once('=') {
302            (left, Some(right.to_string()))
303        } else {
304            (token.as_str(), None)
305        };
306
307        let take_value = |idx: &mut usize| -> Result<String> {
308            if let Some(value) = inline_value.clone() {
309                return Ok(value);
310            }
311            let Some(value) = tokens.get(*idx + 1) else {
312                bail!("Missing value for {flag}");
313            };
314            *idx += 1;
315            Ok(value.clone())
316        };
317
318        match flag {
319            "--name" => name = Some(take_value(&mut index)?),
320            "--prompt" => prompt = Some(take_value(&mut index)?),
321            "--reminder" => reminder = Some(take_value(&mut index)?),
322            "--every" => every = Some(take_value(&mut index)?),
323            "--cron" => cron = Some(take_value(&mut index)?),
324            "--at" => at = Some(take_value(&mut index)?),
325            "--workspace" => workspace = Some(PathBuf::from(take_value(&mut index)?)),
326            "--help" | "help" => {
327                bail!(
328                    "Usage: /schedule create --prompt <text>|--reminder <text> --every <dur>|--cron <expr>|--at <time> [--name <label>] [--workspace <path>]"
329                );
330            }
331            _ => bail!("Unknown option: {token}"),
332        }
333        index += 1;
334    }
335
336    Ok(ScheduleCreateInput { name, prompt, reminder, every, cron, at, workspace })
337}
338
339pub fn parse_local_datetime(raw: &str, now: DateTime<Local>) -> Result<DateTime<Utc>> {
340    let trimmed = raw.trim();
341    if trimmed.is_empty() {
342        bail!("Time value cannot be empty");
343    }
344
345    if let Ok(parsed) = DateTime::parse_from_rfc3339(trimmed) {
346        return Ok(parsed.with_timezone(&Utc));
347    }
348
349    for format in ["%Y-%m-%d %H:%M", "%Y-%m-%dT%H:%M", "%Y-%m-%d %H:%M:%S"] {
350        if let Ok(naive) = NaiveDateTime::parse_from_str(trimmed, format) {
351            return localize_naive_datetime(naive).map(|value| value.with_timezone(&Utc));
352        }
353    }
354
355    if let Some(captures) = TIME_ONLY_RE.captures(trimmed) {
356        let mut hour = captures["hour"].parse::<u32>()?;
357        let minute = captures
358            .name("minute")
359            .map(|value| value.as_str().parse::<u32>())
360            .transpose()?
361            .unwrap_or(0);
362        if minute >= 60 {
363            bail!("Invalid minute in time value");
364        }
365        if let Some(ampm) = captures.name("ampm").map(|value| value.as_str().to_ascii_lowercase()) {
366            if hour == 0 || hour > 12 {
367                bail!("Invalid 12-hour clock value");
368            }
369            if ampm == "pm" && hour != 12 {
370                hour += 12;
371            }
372            if ampm == "am" && hour == 12 {
373                hour = 0;
374            }
375        } else if hour >= 24 {
376            bail!("Invalid 24-hour clock value");
377        }
378
379        let time = NaiveTime::from_hms_opt(hour, minute, 0).ok_or_else(|| anyhow!("Invalid time value"))?;
380        let today = now.date_naive();
381        let naive = today.and_time(time);
382        let mut localized = localize_naive_datetime(naive)?;
383        if localized <= now {
384            localized = localize_naive_datetime((today + ChronoDuration::days(1)).and_time(time))?;
385        }
386        return Ok(localized.with_timezone(&Utc));
387    }
388
389    bail!("Unsupported time format. Use RFC3339, YYYY-MM-DD HH:MM, or a local time like 3pm")
390}
391
392fn localize_naive_datetime(naive: NaiveDateTime) -> Result<DateTime<Local>> {
393    match Local.from_local_datetime(&naive) {
394        LocalResult::Single(value) => Ok(value),
395        LocalResult::Ambiguous(first, _) => Ok(first),
396        LocalResult::None => bail!("Local time does not exist due to timezone transition"),
397    }
398}