weavatrix_rust_refactor/
token.rs

1//! Gate three: a confirmation bound to one plan, one repository, and one use.
2//!
3//! The token is not a password — it is proof that the exact plan being applied is the one a
4//! preview already checked against the working tree. So it carries a fingerprint of the plan
5//! rather than a random value alone, and presenting it for a different plan fails as loudly as
6//! presenting no token at all.
7
8use blazingly_json::{Value, json};
9use std::collections::HashMap;
10use std::path::Path;
11use std::sync::Mutex;
12use std::time::{Duration, SystemTime, UNIX_EPOCH};
13use weavatrix_refactor_plan::EditPlan;
14
15/// How long a confirmation stays valid. Long enough to read a preview, short enough that a stale
16/// tree cannot hide behind it.
17const TOKEN_TTL: Duration = Duration::from_secs(5 * 60);
18
19/// A confirmation handed to the caller after a successful preview.
20pub struct ConfirmToken {
21    pub value: String,
22    pub expires_at: u64,
23}
24
25struct Issued {
26    fingerprint: String,
27    repository: String,
28    expires_at: u64,
29    /// The previewed plan itself, so an apply can present the token alone.
30    ///
31    /// Without this the agent echoes the whole plan back — bytes it just received, paid again
32    /// as completion tokens. The token already proves which plan was previewed; storing the
33    /// plan makes that proof executable.
34    plan: EditPlan,
35}
36
37/// The issued confirmations of one server process.
38#[derive(Default)]
39pub struct TokenStore(Mutex<HashMap<String, Issued>>);
40
41fn now() -> u64 {
42    SystemTime::now()
43        .duration_since(UNIX_EPOCH)
44        .map_or(0, |elapsed| elapsed.as_secs())
45}
46
47/// A stable fingerprint of everything in the plan that decides what gets written.
48///
49/// Ordering, paths, hashes, ranges and both texts all take part: change any of them and the
50/// token stops matching, which is exactly the property that makes it a proof rather than a
51/// formality.
52fn fingerprint(plan: &EditPlan) -> String {
53    let mut material = String::from(&plan.operation);
54    for file in &plan.files {
55        material.push('\u{1}');
56        material.push_str(&file.path);
57        material.push('\u{2}');
58        material.push_str(&file.sha256);
59        for edit in &file.edits {
60            use std::fmt::Write as _;
61            material.push('\u{3}');
62            let _ = write!(
63                material,
64                "{}:{}:{}:{}:{}:{}:{}",
65                edit.start_line,
66                edit.start_char,
67                edit.end_line,
68                edit.end_char,
69                edit.before,
70                edit.after,
71                edit.provenance.as_str()
72            );
73        }
74    }
75    // A non-cryptographic digest is enough: this value never leaves the process and an attacker
76    // who can call this server can call the planner too.
77    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
78    for byte in material.as_bytes() {
79        hash ^= u64::from(*byte);
80        hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
81    }
82    format!("{hash:016x}")
83}
84
85impl TokenStore {
86    /// Issues a confirmation for a previewed plan.
87    pub fn issue(&self, plan: &EditPlan, repository: &Path) -> ConfirmToken {
88        let expires_at = now() + TOKEN_TTL.as_secs();
89        let fingerprint = fingerprint(plan);
90        let value = format!(
91            "{fingerprint}{:016x}",
92            now().wrapping_mul(0x9e37_79b9_7f4a_7c15)
93        );
94        if let Ok(mut issued) = self.0.lock() {
95            issued.retain(|_, token| token.expires_at > now());
96            issued.insert(
97                value.clone(),
98                Issued {
99                    fingerprint,
100                    repository: repository.display().to_string(),
101                    expires_at,
102                    plan: plan.clone(),
103                },
104            );
105        }
106        ConfirmToken { value, expires_at }
107    }
108
109    /// Consumes a confirmation presented alone and returns the plan it previewed.
110    ///
111    /// The checks are the ones `consume` makes, minus the fingerprint comparison — there is no
112    /// second copy of the plan to compare, which is the point: the stored plan *is* the one the
113    /// fingerprint was computed from.
114    ///
115    /// # Errors
116    ///
117    /// Returns the contract refusal — `TOKEN_UNKNOWN`, `TOKEN_EXPIRED` or
118    /// `TOKEN_REPOSITORY_MISMATCH` — when the confirmation cannot be honoured. The token is
119    /// spent either way, so a caller cannot probe with the same value twice.
120    pub fn consume_for_plan(
121        &self,
122        presented: Option<&str>,
123        repository: &Path,
124    ) -> Result<EditPlan, Value> {
125        let Some(presented) = presented else {
126            return Err(json!({
127                "status": "TOKEN_UNKNOWN",
128                "reason": "mode=\"apply\" requires the confirm_token issued by a preview. \
129                           Nothing was written.",
130            }));
131        };
132        let Ok(mut issued) = self.0.lock() else {
133            return Err(json!({
134                "status": "TOKEN_UNKNOWN",
135                "reason": "the token store is unavailable. Nothing was written.",
136            }));
137        };
138        let Some(token) = issued.remove(presented) else {
139            return Err(json!({
140                "status": "TOKEN_UNKNOWN",
141                "reason": "the confirmation is not one this server issued, or it was already \
142                           used. Nothing was written.",
143            }));
144        };
145        if token.expires_at <= now() {
146            return Err(json!({
147                "status": "TOKEN_EXPIRED",
148                "reason": "the confirmation expired; preview again to get a fresh one. Nothing \
149                           was written.",
150            }));
151        }
152        if token.repository != repository.display().to_string() {
153            return Err(json!({
154                "status": "TOKEN_REPOSITORY_MISMATCH",
155                "reason": "the confirmation belongs to a different repository. Nothing was written.",
156            }));
157        }
158        Ok(token.plan)
159    }
160
161    /// Consumes a confirmation, or returns the refusal that says why it could not be.
162    ///
163    /// Consuming happens whether or not the checks pass: a token that was presented is spent,
164    /// so a caller cannot probe with the same value twice.
165    pub fn consume(
166        &self,
167        presented: Option<&str>,
168        plan: &EditPlan,
169        repository: &Path,
170    ) -> Option<Value> {
171        let Some(presented) = presented else {
172            return Some(json!({
173                "status": "TOKEN_UNKNOWN",
174                "reason": "mode=\"apply\" requires the confirm_token issued by a preview of this \
175                           exact plan. Nothing was written.",
176            }));
177        };
178        let mut issued = self.0.lock().ok()?;
179        let Some(token) = issued.remove(presented) else {
180            return Some(json!({
181                "status": "TOKEN_UNKNOWN",
182                "reason": "the confirmation is not one this server issued, or it was already used. \
183                           Nothing was written.",
184            }));
185        };
186        if token.expires_at <= now() {
187            return Some(json!({
188                "status": "TOKEN_EXPIRED",
189                "reason": "the confirmation expired; preview again to get a fresh one. Nothing was written.",
190            }));
191        }
192        if token.repository != repository.display().to_string() {
193            return Some(json!({
194                "status": "TOKEN_REPOSITORY_MISMATCH",
195                "reason": "the confirmation belongs to a different repository. Nothing was written.",
196            }));
197        }
198        if token.fingerprint != fingerprint(plan) {
199            return Some(json!({
200                "status": "TOKEN_PLAN_MISMATCH",
201                "reason": "the plan changed after it was previewed; the confirmation proves a \
202                           different plan. Nothing was written.",
203            }));
204        }
205        None
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use super::TokenStore;
212    use blazingly_json::Value;
213    use std::path::Path;
214    use weavatrix_refactor_plan::{EditPlan, FileEdit, Provenance, TextEdit};
215
216    fn plan(after: &str) -> EditPlan {
217        EditPlan::new(
218            "rename_symbol",
219            vec![FileEdit::new(
220                "src/a.rs",
221                "0".repeat(64),
222                vec![TextEdit {
223                    start_line: 1,
224                    start_char: 0,
225                    end_line: 1,
226                    end_char: 3,
227                    before: "one".to_owned(),
228                    after: after.to_owned(),
229                    provenance: Provenance::new(Provenance::EXACT_LSP),
230                    extensions: std::collections::BTreeMap::new(),
231                }],
232            )],
233        )
234    }
235
236    fn status(value: Option<&Value>) -> Option<&str> {
237        value?.get("status")?.as_str()
238    }
239
240    #[test]
241    fn a_previewed_plan_applies_once() {
242        let store = TokenStore::default();
243        let repository = Path::new("/repo");
244        let token = store.issue(&plan("two"), repository);
245        assert!(
246            store
247                .consume(Some(&token.value), &plan("two"), repository)
248                .is_none()
249        );
250        // The same value a second time is spent, whatever it proved the first time.
251        let replay = store.consume(Some(&token.value), &plan("two"), repository);
252        assert_eq!(status(replay.as_ref()), Some("TOKEN_UNKNOWN"));
253    }
254
255    #[test]
256    fn a_token_does_not_travel_to_another_plan() {
257        let store = TokenStore::default();
258        let repository = Path::new("/repo");
259        let token = store.issue(&plan("two"), repository);
260        let refusal = store.consume(Some(&token.value), &plan("three"), repository);
261        assert_eq!(status(refusal.as_ref()), Some("TOKEN_PLAN_MISMATCH"));
262    }
263
264    #[test]
265    fn a_token_does_not_travel_to_another_repository() {
266        let store = TokenStore::default();
267        let token = store.issue(&plan("two"), Path::new("/repo"));
268        let refusal = store.consume(Some(&token.value), &plan("two"), Path::new("/other"));
269        assert_eq!(status(refusal.as_ref()), Some("TOKEN_REPOSITORY_MISMATCH"));
270    }
271
272    #[test]
273    fn applying_without_a_confirmation_is_refused() {
274        let store = TokenStore::default();
275        let refusal = store.consume(None, &plan("two"), Path::new("/repo"));
276        assert_eq!(status(refusal.as_ref()), Some("TOKEN_UNKNOWN"));
277    }
278
279    #[test]
280    fn an_invented_confirmation_is_refused() {
281        let store = TokenStore::default();
282        let refusal = store.consume(Some("deadbeef"), &plan("two"), Path::new("/repo"));
283        assert_eq!(status(refusal.as_ref()), Some("TOKEN_UNKNOWN"));
284    }
285}
weavatrix_rust_refactor/token.rs

weavatrix_rust_refactor/
token.rs
