Skip to main content

molo_agent/memory/
mod.rs

1//! Memory: manages the agent's context.
2//!
3//! The inference loop uses [`Memory`] to store conversation history and to
4//! retrieve the full context for each request. This module provides two
5//! ready-to-use implementations:
6//!
7//! - [`InMemoryMemory`] — stores all messages verbatim, never trims;
8//! - [`WindowMemory`] — trims to a recent window against a budget, with
9//!   pluggable token counting and trim strategies.
10//!
11//! Compaction, retrieval, persistence and similar strategies plug in by
12//! implementing [`Memory`] (or injecting a [`TrimStrategy`] into
13//! [`WindowMemory`]) without touching the inference loop. A built-in
14//! summarization strategy is provided as [`SummarizeStrategy`].
15
16mod in_memory;
17mod summarize;
18mod window;
19
20pub use in_memory::InMemoryMemory;
21pub use summarize::SummarizeStrategy;
22pub use window::{
23    Budget, CharTokenCounter, TokenCounter, TrimResult, TrimStrategy, WindowDrop, WindowMemory,
24};
25
26use crate::message::Message;
27
28/// Manages the agent's context: decides which messages the model sees on each
29/// turn.
30///
31/// The inference loop interacts with Memory in only two ways:
32/// - hands each turn's new messages (user input, assistant replies, tool
33///   results) to [`record`](Memory::record) /
34///   [`record_protected`](Memory::record_protected);
35/// - fetches the full context via [`context`](Memory::context) before each
36///   request and passes it to the Provider.
37///
38/// How messages are stored, trimmed, and compacted is up to the implementation;
39/// the inference loop is unaware of it. For example, [`WindowMemory`] trims to
40/// a budget on retrieval and guarantees that a tool message and its Assistant
41/// message stay paired; no implementation should split that pair (message
42/// interface constraint). `Send + Sync` ensures `Box<dyn Memory>` can be held
43/// across await points in agent implementations.
44#[async_trait::async_trait]
45pub trait Memory: Send + Sync {
46    /// Records a message, appending it to the end of the session context.
47    ///
48    /// # Errors
49    ///
50    /// Returns [`MemoryError`] when storing or counting fails (e.g., a remote
51    /// counting API is unavailable).
52    async fn record(&mut self, message: Message) -> Result<(), MemoryError>;
53
54    /// Records a **protected** message: implementations must not remove it when
55    /// trimming.
56    ///
57    /// Used for persistent behavior guidance such as skill bodies — trimming
58    /// one silently degrades behavior (the model keeps running but loses the
59    /// dedicated instruction, with no visible error). The default
60    /// implementation behaves like [`record`](Memory::record); implementations
61    /// that trim should override it to exempt protected messages.
62    ///
63    /// # Errors
64    ///
65    /// Same as [`record`](Memory::record).
66    async fn record_protected(&mut self, message: Message) -> Result<(), MemoryError> {
67        self.record(message).await
68    }
69
70    /// Returns the full message sequence of the current context (in order of
71    /// occurrence), which the agent hands to the Provider as-is.
72    ///
73    /// The returned value is a clone; mutating it does not affect Memory's
74    /// internal state.
75    ///
76    /// # Errors
77    ///
78    /// Returns [`MemoryError`] when reading from storage or trimming fails.
79    async fn context(&self) -> Result<Vec<Message>, MemoryError>;
80}
81
82/// Context access failed.
83///
84/// `#[non_exhaustive]`: new error variants can be added in the future without
85/// breaking changes; matches should include a wildcard arm.
86#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
87#[non_exhaustive]
88pub enum MemoryError {
89    /// Storage-layer error; in-memory implementations never produce it, but
90    /// storage-backed implementations (e.g., persistence) do.
91    ///
92    /// The message carries no "memory error" prefix: the type name already
93    /// conveys the domain, and it avoids a doubled prefix when wrapped by
94    /// [`AgentError::Memory`](crate::agent::AgentError).
95    #[error("{0}")]
96    Storage(String),
97    /// Token counting failed; remote counters (e.g., calling a counting API)
98    /// can produce it, local counters never do.
99    #[error("token counting failed: {0}")]
100    TokenCount(String),
101}