Skip to main content

sac/goal/
mod.rs

1use std::path::Path;
2use std::time::SystemTime;
3
4use anyhow::Result;
5use rusqlite::OptionalExtension;
6use serde::{Deserialize, Serialize};
7use uuid::Uuid;
8
9use crate::store::open_connection;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
12#[serde(rename_all = "snake_case")]
13pub enum GoalStatus {
14    Active,
15    Paused,
16    Complete,
17    Blocked,
18    /// Set when session-level usage limit is exceeded during an active goal.
19    UsageLimited,
20    /// Set when the goal's token budget is exhausted. Terminal — only the user
21    /// can raise the budget.
22    BudgetLimited,
23}
24
25impl GoalStatus {
26    pub fn label(self) -> &'static str {
27        match self {
28            Self::Active => "active",
29            Self::Paused => "paused",
30            Self::Complete => "complete",
31            Self::Blocked => "blocked",
32            Self::UsageLimited => "usage_limited",
33            Self::BudgetLimited => "budget_limited",
34        }
35    }
36
37    pub fn is_continuable(self) -> bool {
38        matches!(self, Self::Active)
39    }
40
41    /// Returns `true` for statuses that represent a truly finished state
42    /// from which the goal will never auto-resume.  Matches Codex semantics:
43    /// only `Complete` and `BudgetLimited` are terminal.
44    ///
45    /// `Blocked` and `UsageLimited` are NOT terminal — they are resumable
46    /// states where the user (or system) can clear the condition and
47    /// continue the goal.
48    pub fn is_terminal(self) -> bool {
49        matches!(self, Self::Complete | Self::BudgetLimited)
50    }
51
52    pub fn from_str(s: &str) -> Option<Self> {
53        match s {
54            "active" => Some(Self::Active),
55            "paused" => Some(Self::Paused),
56            "complete" => Some(Self::Complete),
57            "blocked" => Some(Self::Blocked),
58            "usage_limited" => Some(Self::UsageLimited),
59            "budget_limited" => Some(Self::BudgetLimited),
60            _ => None,
61        }
62    }
63}
64
65impl std::fmt::Display for GoalStatus {
66    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        f.write_str(self.label())
68    }
69}
70
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct GoalState {
73    /// Unique identifier for this goal instance.  Used for optimistic
74    /// concurrency control: accounting operations carry the expected
75    /// `goal_id` and skip the write when the stored id differs (the goal
76    /// was replaced between read and write).
77    pub goal_id: String,
78    pub objective: String,
79    pub status: GoalStatus,
80    pub tokens_used: i64,
81    pub time_used_seconds: i64,
82    pub token_budget: Option<i64>,
83    pub created_at: String,
84    pub updated_at: String,
85}
86
87/// Generate a new random goal id (UUID v4).
88pub fn new_goal_id() -> String {
89    Uuid::new_v4().to_string()
90}
91
92pub fn save_goal(store_path: &Path, session_id: &str, goal: &GoalState) -> Result<()> {
93    let conn = open_connection(store_path)?;
94    conn.execute(
95        "INSERT INTO goals (session_id, goal_id, objective, status, tokens_used, time_used_seconds, token_budget, created_at, updated_at)
96         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)
97         ON CONFLICT(session_id) DO UPDATE SET
98             goal_id = excluded.goal_id,
99             objective = excluded.objective,
100             status = excluded.status,
101             tokens_used = excluded.tokens_used,
102             time_used_seconds = excluded.time_used_seconds,
103             token_budget = excluded.token_budget,
104             updated_at = excluded.updated_at",
105        rusqlite::params![
106            session_id,
107            goal.goal_id,
108            goal.objective,
109            goal.status.label(),
110            goal.tokens_used,
111            goal.time_used_seconds,
112            goal.token_budget,
113            goal.created_at,
114            goal.updated_at,
115        ],
116    )?;
117    Ok(())
118}
119
120pub fn load_goal(store_path: &Path, session_id: &str) -> Result<Option<GoalState>> {
121    let conn = open_connection(store_path)?;
122    let result = conn
123        .query_row(
124            "SELECT goal_id, objective, status, tokens_used, time_used_seconds, token_budget, created_at, updated_at
125             FROM goals WHERE session_id = ?1",
126            [session_id],
127            |row| {
128                Ok(GoalState {
129                    goal_id: row.get::<_, String>(0).unwrap_or_default(),
130                    objective: row.get(1)?,
131                    status: GoalStatus::from_str(&row.get::<_, String>(2)?)
132                        .unwrap_or(GoalStatus::Active),
133                    tokens_used: row.get(3)?,
134                    time_used_seconds: row.get(4)?,
135                    token_budget: row.get(5)?,
136                    created_at: row.get(6)?,
137                    updated_at: row.get(7)?,
138                })
139            },
140        )
141        .optional()?;
142    Ok(result)
143}
144
145/// Result of accounting goal usage.
146#[derive(Debug, Clone, Copy, PartialEq, Eq)]
147pub enum AccountingOutcome {
148    /// Usage was recorded; goal remains within budget (or has no budget).
149    WithinBudget,
150    /// Usage was recorded; goal now exceeds its token budget.
151    BudgetExceeded,
152    /// The accounting was skipped because the stored `goal_id` did not
153    /// match the expected value (the goal was replaced between the
154    /// snapshot and the write).
155    Skipped,
156}
157
158/// Atomically increment `tokens_used` and `time_used_seconds` for the
159/// session's goal, returning whether the budget is now exceeded.
160///
161/// When `expected_goal_id` is `Some`, the update is conditional: the row
162/// is only modified when the stored `goal_id` matches.  If it does not
163/// match the goal was replaced between the snapshot and the write, and
164/// `AccountingOutcome::Skipped` is returned to avoid charging the wrong
165/// goal instance.  This mirrors Codex's optimistic concurrency check in
166/// `account_thread_goal_usage`.
167pub fn account_goal_usage(
168    store_path: &Path,
169    session_id: &str,
170    token_delta: i64,
171    time_delta_seconds: i64,
172    expected_goal_id: Option<&str>,
173) -> Result<AccountingOutcome> {
174    let conn = open_connection(store_path)?;
175    let now = now_utc();
176
177    let rows_affected = if let Some(expected_id) = expected_goal_id {
178        conn.execute(
179            "UPDATE goals
180             SET tokens_used = tokens_used + ?1,
181                 time_used_seconds = time_used_seconds + ?2,
182                 updated_at = ?3
183             WHERE session_id = ?4 AND goal_id = ?5",
184            rusqlite::params![token_delta, time_delta_seconds, now, session_id, expected_id],
185        )?
186    } else {
187        conn.execute(
188            "UPDATE goals
189             SET tokens_used = tokens_used + ?1,
190                 time_used_seconds = time_used_seconds + ?2,
191                 updated_at = ?3
192             WHERE session_id = ?4",
193            rusqlite::params![token_delta, time_delta_seconds, now, session_id],
194        )?
195    };
196
197    if rows_affected == 0 {
198        // Either the goal does not exist, or the goal_id didn't match
199        // (the goal was replaced). In both cases, skip accounting.
200        return Ok(AccountingOutcome::Skipped);
201    }
202
203    // Read back the current state to check budget
204    let (tokens_used, token_budget): (i64, Option<i64>) = conn.query_row(
205        "SELECT tokens_used, token_budget FROM goals WHERE session_id = ?1",
206        [session_id],
207        |row| Ok((row.get(0)?, row.get(1)?)),
208    )?;
209    match token_budget {
210        Some(budget) if tokens_used >= budget => Ok(AccountingOutcome::BudgetExceeded),
211        _ => Ok(AccountingOutcome::WithinBudget),
212    }
213}
214
215pub fn delete_goal(store_path: &Path, session_id: &str) -> Result<bool> {
216    let conn = open_connection(store_path)?;
217    let rows_deleted = conn.execute("DELETE FROM goals WHERE session_id = ?1", [session_id])?;
218    Ok(rows_deleted > 0)
219}
220
221pub fn now_utc() -> String {
222    let now = SystemTime::now()
223        .duration_since(SystemTime::UNIX_EPOCH)
224        .unwrap_or_default();
225    let secs = now.as_secs();
226    let time_of_day = secs % 86400;
227    let hours = time_of_day / 3600;
228    let minutes = (time_of_day % 3600) / 60;
229    let seconds = time_of_day % 60;
230
231    // days since epoch to Y/M/D (simplified leap year calculation)
232    let days = secs / 86400;
233    let (year, month, day) = days_to_ymd(days);
234    format!(
235        "{:04}-{:02}-{:02} {:02}:{:02}:{:02}",
236        year, month, day, hours, minutes, seconds
237    )
238}
239
240fn days_to_ymd(days: u64) -> (u64, u64, u64) {
241    // Algorithm from Howard Hinnant's date algorithms
242    let z = days + 719468;
243    let era = z / 146097;
244    let doe = z - era * 146097;
245    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
246    let y = yoe + era * 400;
247    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
248    let mp = (5 * doy + 2) / 153;
249    let d = doy - (153 * mp + 2) / 5 + 1;
250    let m = if mp < 10 { mp + 3 } else { mp - 9 };
251    let y = if m <= 2 { y + 1 } else { y };
252    (y, m, d)
253}