autumn_plugin_agent/
memory.rs

1//! Bounded memory blocks the agent curates itself.
2//!
3//! The model keeps durable notes in labelled [`MemoryBlock`]s, each with a
4//! character limit: `memory` for what it learned, `user` for who it works
5//! for. The limit forces the agent to consolidate instead of hoarding.
6//!
7//! Two rules keep memory cheap and safe:
8//!
9//! * **Frozen snapshot.** The loop renders the blocks into the system prompt
10//!   once, when a run starts. Writes during the run persist at once but show
11//!   up in the prompt only on the next run, so the prompt prefix stays stable
12//!   and the provider's prompt cache keeps hitting.
13//! * **Full means error.** An `add` that would overflow a block fails with a
14//!   message that tells the model to remove or merge entries first.
15//!
16//! [`MemoryStore`] is the persistence seam. [`InMemoryMemoryStore`] serves
17//! tests and single-process apps; back it with your database for anything
18//! that must survive a restart. [`apply_op`] holds the edit rules, so every
19//! store applies them the same way.
20
21use std::collections::HashMap;
22use std::fmt::Write as _;
23use std::future::Future;
24use std::pin::Pin;
25use std::sync::{Arc, Mutex, PoisonError};
26
27use futures::future::BoxFuture;
28use serde::{Deserialize, Serialize};
29
30use crate::error::{AgentError, ErrorKind};
31use crate::tools::{Tool, ToolContext, ToolEffect};
32
33/// Whose memory a run reads and writes.
34///
35/// One app can keep one memory for the whole agent, or one per user or per
36/// tenant: the scope is the key.
37#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
38#[serde(transparent)]
39pub struct MemoryScope(String);
40
41impl MemoryScope {
42    /// Wrap an app-chosen scope key.
43    #[must_use]
44    pub fn new(scope: impl Into<String>) -> Self {
45        Self(scope.into())
46    }
47
48    /// The agent-wide scope, `"agent"`.
49    #[must_use]
50    pub fn agent() -> Self {
51        Self::new("agent")
52    }
53
54    /// The scope as a string slice.
55    #[must_use]
56    pub fn as_str(&self) -> &str {
57        &self.0
58    }
59}
60
61/// One labelled, size-limited list of memory entries.
62#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
63pub struct MemoryBlock {
64    /// Block label the model uses to address it (`memory`, `user`, ...).
65    pub label: String,
66    /// What the block is for, shown to the model.
67    pub description: String,
68    /// The entries, oldest first.
69    pub entries: Vec<String>,
70    /// Maximum total characters across all entries.
71    pub limit_chars: usize,
72}
73
74impl MemoryBlock {
75    /// An empty block.
76    #[must_use]
77    pub fn new(
78        label: impl Into<String>,
79        description: impl Into<String>,
80        limit_chars: usize,
81    ) -> Self {
82        Self {
83            label: label.into(),
84            description: description.into(),
85            entries: Vec::new(),
86            limit_chars,
87        }
88    }
89
90    /// Characters used by all entries.
91    #[must_use]
92    pub fn used_chars(&self) -> usize {
93        self.entries.iter().map(|entry| entry.chars().count()).sum()
94    }
95
96    /// The default pair of blocks: `memory` (2200 chars) and `user`
97    /// (1375 chars).
98    #[must_use]
99    pub fn defaults() -> Vec<Self> {
100        vec![
101            Self::new(
102                "memory",
103                "Facts, decisions, and lessons you want to remember across runs.",
104                2_200,
105            ),
106            Self::new(
107                "user",
108                "Who you work for: preferences, style, standing instructions.",
109                1_375,
110            ),
111        ]
112    }
113}
114
115/// One edit to a memory block.
116#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
117#[serde(tag = "action", rename_all = "snake_case")]
118pub enum MemoryOp {
119    /// Append a new entry.
120    Add {
121        /// Target block label.
122        block: String,
123        /// The entry text.
124        text: String,
125    },
126    /// Replace the one entry that contains `old` with `text`.
127    Replace {
128        /// Target block label.
129        block: String,
130        /// A substring that identifies exactly one entry.
131        old: String,
132        /// The new entry text.
133        text: String,
134    },
135    /// Remove the one entry that contains `old`.
136    Remove {
137        /// Target block label.
138        block: String,
139        /// A substring that identifies exactly one entry.
140        old: String,
141    },
142}
143
144impl MemoryOp {
145    /// The label of the block this op edits.
146    #[must_use]
147    pub fn block(&self) -> &str {
148        match self {
149            Self::Add { block, .. } | Self::Replace { block, .. } | Self::Remove { block, .. } => {
150                block
151            }
152        }
153    }
154}
155
156/// Apply one edit to a block, enforcing the size limit and unique matches.
157///
158/// # Errors
159///
160/// Returns [`ErrorKind::Tool`] errors written for the model: an empty
161/// entry, an overflow (with the numbers), or a `old` substring that matches
162/// no entry or more than one.
163pub fn apply_op(block: &mut MemoryBlock, op: &MemoryOp) -> Result<(), AgentError> {
164    match op {
165        MemoryOp::Add { text, .. } => {
166            let text = non_empty(text)?;
167            check_fits(block, None, text)?;
168            block.entries.push(text.to_owned());
169        }
170        MemoryOp::Replace { old, text, .. } => {
171            let text = non_empty(text)?;
172            let index = find_unique(block, old)?;
173            check_fits(block, Some(index), text)?;
174            if let Some(entry) = block.entries.get_mut(index) {
175                text.clone_into(entry);
176            }
177        }
178        MemoryOp::Remove { old, .. } => {
179            let index = find_unique(block, old)?;
180            block.entries.remove(index);
181        }
182    }
183    Ok(())
184}
185
186fn non_empty(text: &str) -> Result<&str, AgentError> {
187    let text = text.trim();
188    if text.is_empty() {
189        return Err(AgentError::new(
190            ErrorKind::Tool,
191            "memory entries must not be empty",
192        ));
193    }
194    Ok(text)
195}
196
197/// Fail when `text` (replacing entry `replacing`, if any) overflows the block.
198fn check_fits(block: &MemoryBlock, replacing: Option<usize>, text: &str) -> Result<(), AgentError> {
199    let freed = replacing
200        .and_then(|index| block.entries.get(index))
201        .map_or(0, |entry| entry.chars().count());
202    let after = block
203        .used_chars()
204        .saturating_sub(freed)
205        .saturating_add(text.chars().count());
206    if after > block.limit_chars {
207        return Err(AgentError::new(
208            ErrorKind::Tool,
209            format!(
210                "memory block {:?} is full: this edit needs {after} of {} characters. Remove or merge entries first, then retry.",
211                block.label, block.limit_chars
212            ),
213        ));
214    }
215    Ok(())
216}
217
218fn find_unique(block: &MemoryBlock, old: &str) -> Result<usize, AgentError> {
219    let old = old.trim();
220    if old.is_empty() {
221        return Err(AgentError::new(
222            ErrorKind::Tool,
223            "give `old`: a substring of the entry to change",
224        ));
225    }
226    let matches: Vec<usize> = block
227        .entries
228        .iter()
229        .enumerate()
230        .filter(|(_, entry)| entry.contains(old))
231        .map(|(index, _)| index)
232        .collect();
233    match matches.as_slice() {
234        [index] => Ok(*index),
235        [] => Err(AgentError::new(
236            ErrorKind::Tool,
237            format!("no entry in {:?} contains {old:?}", block.label),
238        )),
239        _ => Err(AgentError::new(
240            ErrorKind::Tool,
241            format!(
242                "{} entries in {:?} contain {old:?}: give a longer substring",
243                matches.len(),
244                block.label
245            ),
246        )),
247    }
248}
249
250/// Render blocks as the frozen system-prompt section.
251#[must_use]
252pub fn render_snapshot(blocks: &[MemoryBlock]) -> String {
253    let mut out = String::from(
254        "## Memory\nYour persistent memory, as it was when this run started. Edit it with the `memory` tool; edits show here from the next run.\n",
255    );
256    for block in blocks {
257        let _ = write!(
258            out,
259            "\n<memory block=\"{}\" used=\"{}/{}\">\n{}\n",
260            block.label,
261            block.used_chars(),
262            block.limit_chars,
263            block.description
264        );
265        for entry in &block.entries {
266            out.push_str("- ");
267            out.push_str(entry);
268            out.push('\n');
269        }
270        out.push_str("</memory>\n");
271    }
272    out
273}
274
275/// Persistence for memory blocks.
276pub trait MemoryStore: Send + Sync + std::fmt::Debug {
277    /// Load every block for a scope. A new scope gets the store's default
278    /// (empty) blocks.
279    fn load<'a>(
280        &'a self,
281        scope: &'a MemoryScope,
282    ) -> BoxFuture<'a, Result<Vec<MemoryBlock>, AgentError>>;
283
284    /// Apply one edit and persist it. Implementations call [`apply_op`] so
285    /// the rules match everywhere.
286    fn apply<'a>(
287        &'a self,
288        scope: &'a MemoryScope,
289        op: MemoryOp,
290    ) -> BoxFuture<'a, Result<MemoryBlock, AgentError>>;
291}
292
293/// A process-local [`MemoryStore`]. Contents vanish on restart.
294#[derive(Debug)]
295pub struct InMemoryMemoryStore {
296    template: Vec<MemoryBlock>,
297    scopes: Mutex<HashMap<MemoryScope, Vec<MemoryBlock>>>,
298}
299
300impl Default for InMemoryMemoryStore {
301    fn default() -> Self {
302        Self::new(MemoryBlock::defaults())
303    }
304}
305
306impl InMemoryMemoryStore {
307    /// A store whose new scopes start with copies of `template`.
308    #[must_use]
309    pub fn new(template: Vec<MemoryBlock>) -> Self {
310        Self {
311            template,
312            scopes: Mutex::new(HashMap::new()),
313        }
314    }
315
316    /// Share the store.
317    #[must_use]
318    pub fn shared(self) -> Arc<dyn MemoryStore> {
319        Arc::new(self)
320    }
321}
322
323impl MemoryStore for InMemoryMemoryStore {
324    fn load<'a>(
325        &'a self,
326        scope: &'a MemoryScope,
327    ) -> BoxFuture<'a, Result<Vec<MemoryBlock>, AgentError>> {
328        let blocks = self
329            .scopes
330            .lock()
331            .unwrap_or_else(PoisonError::into_inner)
332            .get(scope)
333            .cloned()
334            .unwrap_or_else(|| self.template.clone());
335        Box::pin(std::future::ready(Ok(blocks)))
336    }
337
338    fn apply<'a>(
339        &'a self,
340        scope: &'a MemoryScope,
341        op: MemoryOp,
342    ) -> BoxFuture<'a, Result<MemoryBlock, AgentError>> {
343        let mut scopes = self.scopes.lock().unwrap_or_else(PoisonError::into_inner);
344        let blocks = scopes
345            .entry(scope.clone())
346            .or_insert_with(|| self.template.clone());
347        let result = match blocks.iter().position(|block| block.label == op.block()) {
348            Some(index) => blocks
349                .get_mut(index)
350                .ok_or_else(|| unknown_block(op.block(), &[]))
351                .and_then(|block| apply_op(block, &op).map(|()| block.clone())),
352            None => Err(unknown_block(op.block(), blocks)),
353        };
354        drop(scopes);
355        Box::pin(std::future::ready(result))
356    }
357}
358
359fn unknown_block(label: &str, blocks: &[MemoryBlock]) -> AgentError {
360    let known: Vec<&str> = blocks.iter().map(|block| block.label.as_str()).collect();
361    AgentError::new(
362        ErrorKind::Tool,
363        format!("no memory block {label:?}; blocks: {known:?}"),
364    )
365}
366
367/// The built-in `memory` tool, bound to one store and scope.
368///
369/// The agent adds it to a run on its own when memory is configured; you
370/// rarely build one by hand.
371#[derive(Debug, Clone)]
372pub struct MemoryTool {
373    store: Arc<dyn MemoryStore>,
374    scope: MemoryScope,
375}
376
377impl MemoryTool {
378    /// Bind the tool to a store and scope.
379    #[must_use]
380    pub fn new(store: Arc<dyn MemoryStore>, scope: MemoryScope) -> Self {
381        Self { store, scope }
382    }
383}
384
385impl Tool for MemoryTool {
386    fn name(&self) -> &'static str {
387        "memory"
388    }
389
390    fn description(&self) -> &'static str {
391        "Edit your persistent memory. action=add appends `text` to `block`; action=replace swaps the entry containing `old` for `text`; action=remove deletes the entry containing `old`. Save durable facts and lessons, not logs."
392    }
393
394    fn input_schema(&self) -> serde_json::Value {
395        serde_json::json!({
396            "type": "object",
397            "properties": {
398                "action": {"type": "string", "enum": ["add", "replace", "remove"]},
399                "block": {"type": "string", "description": "Block label, e.g. memory or user."},
400                "text": {"type": "string", "description": "New entry text (add, replace)."},
401                "old": {"type": "string", "description": "Substring of the entry to change (replace, remove)."}
402            },
403            "required": ["action", "block"]
404        })
405    }
406
407    fn effect(&self) -> ToolEffect {
408        ToolEffect::Internal
409    }
410
411    fn execute<'a>(
412        &'a self,
413        input: serde_json::Value,
414        _ctx: &'a ToolContext,
415    ) -> Pin<Box<dyn Future<Output = Result<serde_json::Value, AgentError>> + Send + 'a>> {
416        Box::pin(async move {
417            let op: MemoryOp = serde_json::from_value(input).map_err(|err| {
418                AgentError::with_source(ErrorKind::Tool, format!("invalid memory edit: {err}"), err)
419            })?;
420            let block = self.store.apply(&self.scope, op).await?;
421            Ok(serde_json::json!({
422                "ok": true,
423                "block": block.label,
424                "used": block.used_chars(),
425                "limit": block.limit_chars,
426            }))
427        })
428    }
429}
430
431#[cfg(test)]
432mod tests;
autumn_plugin_agent/memory.rs

autumn_plugin_agent/
memory.rs
