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}