objects/store/
writer_lease.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Exclusive writer leases for agent-controlled thread mutations.
3
4use std::path::{Path, PathBuf};
5
6use chrono::{DateTime, Utc};
7use serde::{Deserialize, Serialize};
8
9use crate::{
10    fs_atomic::write_file_atomic,
11    lock::RepoLock,
12    store::{HeddleError, Liveness, Result, reservation_liveness_at},
13};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16#[serde(rename_all = "snake_case")]
17pub enum WriterLeaseStatus {
18    Active,
19    Complete,
20    Abandoned,
21}
22
23impl std::fmt::Display for WriterLeaseStatus {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        match self {
26            Self::Active => write!(f, "active"),
27            Self::Complete => write!(f, "complete"),
28            Self::Abandoned => write!(f, "abandoned"),
29        }
30    }
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct WriterLease {
35    pub lease_id: String,
36    pub thread: String,
37    #[serde(default)]
38    pub actor_session_id: Option<String>,
39    #[serde(default)]
40    pub task_assignment_id: Option<String>,
41    #[serde(default)]
42    pub anchor_state: Option<String>,
43    #[serde(default)]
44    pub anchor_root: Option<String>,
45    #[serde(default)]
46    pub path: Option<PathBuf>,
47    pub token_hash: String,
48    #[serde(default)]
49    pub pid: Option<u32>,
50    #[serde(default)]
51    pub boot_id: Option<String>,
52    pub heartbeat_at: DateTime<Utc>,
53    pub started_at: DateTime<Utc>,
54    pub status: WriterLeaseStatus,
55    #[serde(default)]
56    pub completed_at: Option<DateTime<Utc>>,
57}
58
59impl WriterLease {
60    pub fn lease_expires_at(&self) -> DateTime<Utc> {
61        self.heartbeat_at + crate::store::AGENT_LEASE_DURATION
62    }
63
64    pub fn liveness_at(&self, now: DateTime<Utc>) -> Liveness {
65        if self.status != WriterLeaseStatus::Active {
66            return Liveness::Dead;
67        }
68        reservation_liveness_at(
69            self.pid,
70            self.boot_id.as_deref(),
71            Some(self.heartbeat_at),
72            now,
73        )
74    }
75}
76
77#[derive(Debug, Clone)]
78pub struct WriterLeaseDraft {
79    pub thread: String,
80    pub actor_session_id: Option<String>,
81    pub task_assignment_id: Option<String>,
82    pub anchor_state: Option<String>,
83    pub anchor_root: Option<String>,
84    pub path: Option<PathBuf>,
85    pub pid: Option<u32>,
86    pub boot_id: Option<String>,
87}
88
89#[derive(Debug)]
90pub struct WriterLeaseGrant {
91    pub lease: WriterLease,
92    pub token: String,
93}
94
95#[derive(Debug)]
96pub enum WriterLeaseReserveOutcome {
97    Reserved(WriterLeaseGrant),
98    LiveOwner(WriterLease),
99}
100
101#[derive(Debug)]
102pub enum WriterLeaseAuthOutcome {
103    Authorized(WriterLease),
104    Missing,
105    TokenMismatch,
106    Inactive(WriterLease),
107}
108
109pub struct WriterLeaseStore {
110    leases_dir: PathBuf,
111}
112
113impl WriterLeaseStore {
114    pub fn new(heddle_dir: &Path) -> Self {
115        Self {
116            leases_dir: heddle_dir.join("writer-leases"),
117        }
118    }
119
120    fn lock_path(&self) -> PathBuf {
121        self.leases_dir.join(".lock")
122    }
123
124    fn write_lock(&self) -> Result<crate::lock::WriteLockGuard> {
125        RepoLock::at(self.lock_path()).write().map_err(|err| {
126            HeddleError::Config(format!("failed to acquire writer lease lock: {err}"))
127        })
128    }
129
130    fn lease_path(&self, lease_id: &str) -> Result<PathBuf> {
131        validate_lease_id(lease_id)?;
132        Ok(self.leases_dir.join(format!("{lease_id}.toml")))
133    }
134
135    fn load_path(&self, path: &Path) -> Result<Option<WriterLease>> {
136        if !path.exists() {
137            return Ok(None);
138        }
139        let content = std::fs::read_to_string(path)?;
140        toml::from_str(&content)
141            .map(Some)
142            .map_err(|err| HeddleError::Config(err.to_string()))
143    }
144
145    fn write_lease(&self, lease: &WriterLease) -> Result<()> {
146        crate::fs_atomic::create_dir_all_durable(&self.leases_dir)?;
147        let content =
148            toml::to_string_pretty(lease).map_err(|err| HeddleError::Config(err.to_string()))?;
149        Ok(write_file_atomic(
150            &self.lease_path(&lease.lease_id)?,
151            content.as_bytes(),
152        )?)
153    }
154
155    fn list_locked(&self) -> Result<Vec<WriterLease>> {
156        if !self.leases_dir.exists() {
157            return Ok(Vec::new());
158        }
159        let mut leases = Vec::new();
160        for entry in std::fs::read_dir(&self.leases_dir)? {
161            let path = entry?.path();
162            if path
163                .extension()
164                .is_some_and(|extension| extension == "toml")
165                && let Some(lease) = self.load_path(&path)?
166            {
167                leases.push(lease);
168            }
169        }
170        leases.sort_by_key(|lease| std::cmp::Reverse(lease.started_at));
171        Ok(leases)
172    }
173
174    fn reap_expired_locked(&self, now: DateTime<Utc>) -> Result<()> {
175        for mut lease in self.list_locked()? {
176            if lease.liveness_at(now) == Liveness::Dead && lease.status == WriterLeaseStatus::Active
177            {
178                lease.status = WriterLeaseStatus::Abandoned;
179                lease.completed_at = Some(now);
180                self.write_lease(&lease)?;
181            }
182        }
183        Ok(())
184    }
185
186    pub fn reserve(
187        &self,
188        draft: WriterLeaseDraft,
189        now: DateTime<Utc>,
190    ) -> Result<WriterLeaseReserveOutcome> {
191        let _lock = self.write_lock()?;
192        self.reap_expired_locked(now)?;
193        if let Some(owner) = self
194            .list_locked()?
195            .into_iter()
196            .find(|lease| lease.thread == draft.thread && lease.status == WriterLeaseStatus::Active)
197        {
198            return Ok(WriterLeaseReserveOutcome::LiveOwner(owner));
199        }
200
201        let lease_id = generate_writer_lease_id();
202        let token = generate_writer_lease_token();
203        let lease = WriterLease {
204            lease_id,
205            thread: draft.thread,
206            actor_session_id: draft.actor_session_id,
207            task_assignment_id: draft.task_assignment_id,
208            anchor_state: draft.anchor_state,
209            anchor_root: draft.anchor_root,
210            path: draft.path,
211            token_hash: token_hash(&token),
212            pid: draft.pid,
213            boot_id: draft.boot_id,
214            heartbeat_at: now,
215            started_at: now,
216            status: WriterLeaseStatus::Active,
217            completed_at: None,
218        };
219        self.write_lease(&lease)?;
220        Ok(WriterLeaseReserveOutcome::Reserved(WriterLeaseGrant {
221            lease,
222            token,
223        }))
224    }
225
226    pub fn authenticate_and_renew(
227        &self,
228        lease_id: &str,
229        token: &str,
230        now: DateTime<Utc>,
231    ) -> Result<WriterLeaseAuthOutcome> {
232        let _lock = self.write_lock()?;
233        let path = self.lease_path(lease_id)?;
234        let Some(mut lease) = self.load_path(&path)? else {
235            return Ok(WriterLeaseAuthOutcome::Missing);
236        };
237        if lease.status != WriterLeaseStatus::Active || lease.liveness_at(now) == Liveness::Dead {
238            if lease.status == WriterLeaseStatus::Active {
239                lease.status = WriterLeaseStatus::Abandoned;
240                lease.completed_at = Some(now);
241                self.write_lease(&lease)?;
242            }
243            return Ok(WriterLeaseAuthOutcome::Inactive(lease));
244        }
245        if token_hash(token) != lease.token_hash {
246            return Ok(WriterLeaseAuthOutcome::TokenMismatch);
247        }
248        lease.heartbeat_at = now;
249        self.write_lease(&lease)?;
250        Ok(WriterLeaseAuthOutcome::Authorized(lease))
251    }
252
253    pub fn release(
254        &self,
255        lease_id: &str,
256        token: &str,
257        status: WriterLeaseStatus,
258        now: DateTime<Utc>,
259    ) -> Result<WriterLeaseAuthOutcome> {
260        let _lock = self.write_lock()?;
261        let path = self.lease_path(lease_id)?;
262        let Some(mut lease) = self.load_path(&path)? else {
263            return Ok(WriterLeaseAuthOutcome::Missing);
264        };
265        if lease.status != WriterLeaseStatus::Active {
266            return Ok(WriterLeaseAuthOutcome::Inactive(lease));
267        }
268        if token_hash(token) != lease.token_hash {
269            return Ok(WriterLeaseAuthOutcome::TokenMismatch);
270        }
271        lease.status = status;
272        lease.completed_at = Some(now);
273        self.write_lease(&lease)?;
274        Ok(WriterLeaseAuthOutcome::Authorized(lease))
275    }
276
277    pub fn list(&self) -> Result<Vec<WriterLease>> {
278        let _lock = self.write_lock()?;
279        self.reap_expired_locked(Utc::now())?;
280        self.list_locked()
281    }
282
283    /// List persisted leases without reaping expired active records.
284    ///
285    /// Cleanup previews use this read-only view so residue detection cannot
286    /// mutate lease storage as a side effect.
287    pub fn list_without_reaping(&self) -> Result<Vec<WriterLease>> {
288        self.list_locked()
289    }
290
291    pub fn abandon_thread(&self, thread: &str, now: DateTime<Utc>) -> Result<()> {
292        let _lock = self.write_lock()?;
293        for mut lease in self.list_locked()? {
294            if lease.thread == thread && lease.status == WriterLeaseStatus::Active {
295                lease.status = WriterLeaseStatus::Abandoned;
296                lease.completed_at = Some(now);
297                self.write_lease(&lease)?;
298            }
299        }
300        Ok(())
301    }
302
303    pub fn load(&self, lease_id: &str) -> Result<Option<WriterLease>> {
304        self.load_path(&self.lease_path(lease_id)?)
305    }
306}
307
308pub fn generate_writer_lease_id() -> String {
309    format!("lease-{}", random_base32())
310}
311
312pub fn generate_writer_lease_token() -> String {
313    format!("hwl_{}", random_base32())
314}
315
316fn random_base32() -> String {
317    let random_bytes: [u8; 24] = rand::random();
318    base32::encode(base32::Alphabet::Rfc4648 { padding: false }, &random_bytes).to_lowercase()
319}
320
321fn token_hash(token: &str) -> String {
322    blake3::hash(token.as_bytes()).to_hex().to_string()
323}
324
325fn validate_lease_id(lease_id: &str) -> Result<()> {
326    if lease_id.starts_with("lease-")
327        && lease_id
328            .bytes()
329            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-')
330    {
331        return Ok(());
332    }
333    Err(HeddleError::Config(format!(
334        "invalid writer lease id '{lease_id}'"
335    )))
336}
337
338#[cfg(test)]
339mod tests {
340    use tempfile::TempDir;
341
342    use super::*;
343
344    fn draft(thread: &str) -> WriterLeaseDraft {
345        WriterLeaseDraft {
346            thread: thread.to_string(),
347            actor_session_id: Some("agent-one".to_string()),
348            task_assignment_id: None,
349            anchor_state: Some("hd-state".to_string()),
350            anchor_root: Some("root".to_string()),
351            path: None,
352            pid: None,
353            boot_id: None,
354        }
355    }
356
357    #[test]
358    fn token_is_required_to_renew_or_release() {
359        let temp = TempDir::new().unwrap();
360        let store = WriterLeaseStore::new(temp.path());
361        let now = Utc::now();
362        let WriterLeaseReserveOutcome::Reserved(grant) =
363            store.reserve(draft("feature/a"), now).unwrap()
364        else {
365            panic!("first lease should reserve");
366        };
367        assert!(matches!(
368            store
369                .authenticate_and_renew(&grant.lease.lease_id, "wrong", now)
370                .unwrap(),
371            WriterLeaseAuthOutcome::TokenMismatch
372        ));
373        assert!(matches!(
374            store
375                .authenticate_and_renew(&grant.lease.lease_id, &grant.token, now)
376                .unwrap(),
377            WriterLeaseAuthOutcome::Authorized(_)
378        ));
379    }
380
381    #[test]
382    fn expired_lease_does_not_block_a_new_owner() {
383        let temp = TempDir::new().unwrap();
384        let store = WriterLeaseStore::new(temp.path());
385        let now = Utc::now();
386        let WriterLeaseReserveOutcome::Reserved(first) =
387            store.reserve(draft("feature/a"), now).unwrap()
388        else {
389            panic!("first lease should reserve");
390        };
391        let later = now + crate::store::AGENT_LEASE_DURATION + chrono::Duration::seconds(1);
392        let WriterLeaseReserveOutcome::Reserved(second) =
393            store.reserve(draft("feature/a"), later).unwrap()
394        else {
395            panic!("expired lease should not block");
396        };
397        assert_ne!(first.lease.lease_id, second.lease.lease_id);
398    }
399
400    #[test]
401    fn stored_lease_does_not_contain_bearer_token() {
402        let temp = TempDir::new().unwrap();
403        let store = WriterLeaseStore::new(temp.path());
404        let WriterLeaseReserveOutcome::Reserved(grant) =
405            store.reserve(draft("feature/a"), Utc::now()).unwrap()
406        else {
407            panic!("first lease should reserve");
408        };
409        let persisted = std::fs::read_to_string(
410            temp.path()
411                .join("writer-leases")
412                .join(format!("{}.toml", grant.lease.lease_id)),
413        )
414        .unwrap();
415        assert!(!persisted.contains(&grant.token));
416        assert!(persisted.contains("token_hash"));
417    }
418}
objects/store/writer_lease.rs

objects/store/
writer_lease.rs
