Skip to main content

aion_core/
assistant_document.rs

1//! Structured edits to an assistant session's shared document.
2//!
3//! The assistant's `assistant_document_edit` tool submits edits as
4//! **replace-exactly-once** operations — an `old_string` that must occur
5//! exactly once in the document as it now stands, and the `new_string` that
6//! replaces it. This is the one edit idiom a language model reliably produces:
7//! it names the bytes it read, so a stale read fails loudly instead of landing
8//! somewhere else, and it needs no offset arithmetic that drifts the moment the
9//! operator types.
10//!
11//! # One applier, two callers
12//!
13//! [`apply_document_edits`] is the ONLY definition of what a batch of edits
14//! does to a text. The server's append path calls it to VALIDATE a batch before
15//! recording it, and the transcript projection calls it again to fold a
16//! recorded batch into the shared document. Because both read the same
17//! function, an edit the append path accepted is an edit the projection applies
18//! identically — there is no second matcher to disagree with the first.
19//!
20//! # Batches are atomic
21//!
22//! Edits in one batch apply in order, each against the text the previous one
23//! produced. The first edit that cannot apply fails the WHOLE batch — nothing
24//! before it is kept — so a recorded batch always applied cleanly in full, and
25//! a refused batch changed nothing at all.
26
27use serde::{Deserialize, Serialize};
28
29/// One replace-exactly-once edit.
30#[derive(Serialize, Deserialize, ts_rs::TS, Clone, Debug, PartialEq, Eq)]
31pub struct AssistantDocumentEditOp {
32    /// The exact bytes to replace. Must occur exactly once in the document as
33    /// it stands when this edit applies — zero matches means the edit was
34    /// written against text that is not there; two or more means it does not
35    /// say WHICH occurrence it meant.
36    pub old_string: String,
37    /// What replaces them.
38    pub new_string: String,
39}
40
41/// Why a batch of edits could not be applied.
42///
43/// Phrased around the ordinal (one-based, the position in the submitted batch)
44/// because the caller is a model reading its own batch back: "edit 2" is how it
45/// finds the operation it wrote.
46#[derive(thiserror::Error, Clone, Debug, PartialEq, Eq)]
47pub enum AssistantDocumentEditError {
48    /// The edit's `old_string` is empty, which matches nothing and everything.
49    #[error(
50        "edit {ordinal}: `old_string` is empty — an empty match names no place in the document. \
51         Quote the exact bytes to replace."
52    )]
53    EmptyOldString {
54        /// One-based position of the edit in its batch.
55        ordinal: usize,
56    },
57    /// The edit's `old_string` does not occur in the document.
58    #[error(
59        "edit {ordinal}: `old_string` was not found in the document (searched for {preview:?}). \
60         The document may have changed since it was read — read it again and re-quote the exact \
61         bytes."
62    )]
63    Absent {
64        /// One-based position of the edit in its batch.
65        ordinal: usize,
66        /// The start of the string that was searched for, for the message.
67        preview: String,
68    },
69    /// The edit's `old_string` occurs more than once, so it does not say which
70    /// occurrence it means.
71    #[error(
72        "edit {ordinal}: `old_string` occurs {matches} times (searched for {preview:?}), so it \
73         does not say which occurrence to replace. Quote more surrounding text to make it unique."
74    )]
75    Ambiguous {
76        /// One-based position of the edit in its batch.
77        ordinal: usize,
78        /// How many times it occurred.
79        matches: usize,
80        /// The start of the string that was searched for, for the message.
81        preview: String,
82    },
83}
84
85/// How much of a failed `old_string` the error message quotes back.
86const PREVIEW_CHARACTERS: usize = 80;
87
88/// Apply `edits` to `text`, in order, each replacing exactly one occurrence.
89///
90/// # Errors
91///
92/// [`AssistantDocumentEditError`] naming the first edit that could not apply.
93/// The batch is atomic: on error, nothing was applied.
94pub fn apply_document_edits(
95    text: &str,
96    edits: &[AssistantDocumentEditOp],
97) -> Result<String, AssistantDocumentEditError> {
98    let mut current = text.to_owned();
99    for (index, edit) in edits.iter().enumerate() {
100        let ordinal = index.saturating_add(1);
101        if edit.old_string.is_empty() {
102            return Err(AssistantDocumentEditError::EmptyOldString { ordinal });
103        }
104        match current.matches(edit.old_string.as_str()).count() {
105            0 => {
106                return Err(AssistantDocumentEditError::Absent {
107                    ordinal,
108                    preview: preview_of(&edit.old_string),
109                });
110            }
111            1 => {
112                current = current.replacen(edit.old_string.as_str(), &edit.new_string, 1);
113            }
114            matches => {
115                return Err(AssistantDocumentEditError::Ambiguous {
116                    ordinal,
117                    matches,
118                    preview: preview_of(&edit.old_string),
119                });
120            }
121        }
122    }
123    Ok(current)
124}
125
126/// The first [`PREVIEW_CHARACTERS`] characters, on character boundaries.
127fn preview_of(searched: &str) -> String {
128    searched.chars().take(PREVIEW_CHARACTERS).collect()
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    fn op(old: &str, new: &str) -> AssistantDocumentEditOp {
136        AssistantDocumentEditOp {
137            old_string: old.to_owned(),
138            new_string: new.to_owned(),
139        }
140    }
141
142    /// Edits apply in order, each against the text the previous one produced —
143    /// so a later edit may match text an earlier edit wrote.
144    #[test]
145    fn edits_apply_in_order_against_the_running_text() -> Result<(), AssistantDocumentEditError> {
146        let applied = apply_document_edits(
147            "workflow demo\nstep one\n",
148            &[op("step one", "step first"), op("first", "first_renamed")],
149        )?;
150        assert_eq!(applied, "workflow demo\nstep first_renamed\n");
151        Ok(())
152    }
153
154    /// Zero matches refuses the batch and names the edit, so a model that read
155    /// a stale document is told to read again rather than left guessing.
156    #[test]
157    fn an_absent_old_string_refuses_the_whole_batch() {
158        let refused = apply_document_edits(
159            "workflow demo\n",
160            &[op("workflow demo", "workflow renamed"), op("gone", "there")],
161        );
162        assert_eq!(
163            refused,
164            Err(AssistantDocumentEditError::Absent {
165                ordinal: 2,
166                preview: "gone".to_owned(),
167            }),
168            "the batch is atomic: the first edit matched, and none of it lands"
169        );
170    }
171
172    /// Two matches is a refusal, not a guess: the edit does not say which
173    /// occurrence it means, and picking the first would silently edit the
174    /// wrong place half the time.
175    #[test]
176    fn an_ambiguous_old_string_is_refused_with_its_count() {
177        let refused = apply_document_edits("a b a", &[op("a", "c")]);
178        assert_eq!(
179            refused,
180            Err(AssistantDocumentEditError::Ambiguous {
181                ordinal: 1,
182                matches: 2,
183                preview: "a".to_owned(),
184            })
185        );
186    }
187
188    /// An empty `old_string` matches everywhere and nowhere; it is refused by
189    /// name rather than reported as absent or applied at position zero.
190    #[test]
191    fn an_empty_old_string_is_refused_by_name() {
192        let refused = apply_document_edits("text", &[op("", "inserted")]);
193        assert_eq!(
194            refused,
195            Err(AssistantDocumentEditError::EmptyOldString { ordinal: 1 })
196        );
197    }
198
199    /// The error's preview is clipped on CHARACTER boundaries, so a long
200    /// multi-byte `old_string` cannot split a code point in the message.
201    #[test]
202    fn a_long_search_string_is_previewed_on_character_boundaries() -> Result<(), String> {
203        let long = "é".repeat(200);
204        let refused = apply_document_edits("text", &[op(&long, "x")]);
205        match refused {
206            Err(AssistantDocumentEditError::Absent { preview, .. }) => {
207                assert_eq!(preview.chars().count(), 80);
208                Ok(())
209            }
210            other => Err(format!("expected an Absent refusal, got {other:?}")),
211        }
212    }
213
214    /// Replacement is literal, exactly-once: regex metacharacters in either
215    /// string are bytes, not patterns.
216    #[test]
217    fn strings_are_literal_bytes_not_patterns() -> Result<(), AssistantDocumentEditError> {
218        let applied = apply_document_edits("route a|b", &[op("a|b", "a.b")])?;
219        assert_eq!(applied, "route a.b");
220        Ok(())
221    }
222
223    /// Occurrences are counted DISJOINT, exactly as [`str::matches`] counts
224    /// them: `"\n\n"` occurs ONCE in `"\n\n\n"`, so the edit applies. This is
225    /// the canonical exactly-once semantics; the console's twin applier must
226    /// count the same way, or the server records a batch the editor refuses.
227    #[test]
228    fn overlapping_candidates_count_as_disjoint_matches() -> Result<(), AssistantDocumentEditError>
229    {
230        let applied = apply_document_edits("a\n\n\nb", &[op("\n\n", "\n")])?;
231        assert_eq!(applied, "a\n\nb");
232        let applied = apply_document_edits("aaa", &[op("aa", "x")])?;
233        assert_eq!(applied, "xa");
234        Ok(())
235    }
236}