agent_framework_core/
history.rs

1//! Conversation-history context providers.
2//!
3//! Upstream moved conversation history out of the thread/session entirely:
4//! it is now just another [`ContextProvider`] — a [`HistoryProvider`] —
5//! that prepends its stored messages ahead of a run (`before_run`) and
6//! records the run's request + response messages after a successful run
7//! (`after_run`). [`InMemoryHistoryProvider`] is the in-process default;
8//! [`FileHistoryProvider`] persists to a JSON file on disk.
9//!
10//! [`Agent`](crate::agent::Agent) and
11//! [`WorkflowAgent`](crate::workflow::WorkflowAgent) auto-attach a fresh
12//! [`InMemoryHistoryProvider`] (via [`ensure_history_provider`]) to any
13//! non-service-managed [`AgentSession`] that doesn't already carry a history
14//! provider, so local multi-turn conversations keep accumulating history the
15//! way the old `AgentThread` message store used to.
16
17use std::path::PathBuf;
18use std::sync::{Arc, Mutex};
19
20use async_trait::async_trait;
21use serde_json::Value;
22
23use crate::error::{Error, Result};
24use crate::memory::{ContextProvider, SessionContext};
25use crate::session::AgentSession;
26use crate::types::Message;
27
28/// A [`ContextProvider`] that also manages conversation history.
29///
30/// This is a marker trait (over and above [`ContextProvider::is_history_provider`],
31/// which drives runtime detection via trait objects): implementing it
32/// documents that a type's `before_run`/`after_run` are the ones responsible
33/// for a session's conversation history, distinguishing it from a generic
34/// memory/RAG provider.
35pub trait HistoryProvider: ContextProvider {}
36
37/// Attach a fresh [`InMemoryHistoryProvider`] as the **first** context
38/// provider on `session` when it is not service-managed and does not already
39/// carry a history provider. A no-op for service-managed sessions (the
40/// service owns history server-side) and for sessions that already have one
41/// attached (detected via [`ContextProvider::is_history_provider`]).
42pub fn ensure_history_provider(session: &mut AgentSession) {
43    if session.service_session_id().is_none()
44        && !session
45            .context_providers
46            .iter()
47            .any(|p| p.is_history_provider())
48    {
49        session
50            .context_providers
51            .insert(0, Arc::new(InMemoryHistoryProvider::new()));
52    }
53}
54
55/// In-memory [`HistoryProvider`]: keeps history in an `Arc<Mutex<Vec<Message>>>`,
56/// shared across clones.
57#[derive(Default, Clone)]
58pub struct InMemoryHistoryProvider {
59    messages: Arc<Mutex<Vec<Message>>>,
60}
61
62impl InMemoryHistoryProvider {
63    /// An empty history provider.
64    pub fn new() -> Self {
65        Self::default()
66    }
67
68    /// A history provider seeded with `messages`.
69    pub fn with_messages(messages: Vec<Message>) -> Self {
70        Self {
71            messages: Arc::new(Mutex::new(messages)),
72        }
73    }
74
75    /// The stored messages, in chronological order.
76    pub fn list_messages(&self) -> Vec<Message> {
77        self.messages.lock().unwrap().clone()
78    }
79
80    /// Serialize the stored history to `{"messages": [...]}`.
81    pub fn to_dict(&self) -> Value {
82        serde_json::json!({ "messages": self.list_messages() })
83    }
84
85    /// Reconstruct a provider from state produced by [`InMemoryHistoryProvider::to_dict`].
86    pub fn from_dict(state: &Value) -> Result<Self> {
87        let messages = match state.get("messages") {
88            Some(v) if !v.is_null() => serde_json::from_value(v.clone()).map_err(|e| {
89                Error::Serialization(format!("failed to restore history provider: {e}"))
90            })?,
91            _ => Vec::new(),
92        };
93        Ok(Self::with_messages(messages))
94    }
95}
96
97#[async_trait]
98impl ContextProvider for InMemoryHistoryProvider {
99    async fn before_run(&self, ctx: &mut SessionContext) -> Result<()> {
100        let stored = self.messages.lock().unwrap().clone();
101        let existing = std::mem::take(&mut ctx.messages);
102        ctx.messages = stored.into_iter().chain(existing).collect();
103        Ok(())
104    }
105
106    async fn after_run(
107        &self,
108        request_messages: &[Message],
109        response_messages: &[Message],
110        error: Option<&Error>,
111    ) -> Result<()> {
112        if error.is_none() {
113            let mut guard = self.messages.lock().unwrap();
114            guard.extend(request_messages.iter().cloned());
115            guard.extend(response_messages.iter().cloned());
116        }
117        Ok(())
118    }
119
120    fn is_history_provider(&self) -> bool {
121        true
122    }
123}
124
125impl HistoryProvider for InMemoryHistoryProvider {}
126
127/// A [`HistoryProvider`] that persists to a JSON file on disk, loading any
128/// existing history from `path` on construction and rewriting the whole file
129/// after every successful run.
130#[derive(Clone)]
131pub struct FileHistoryProvider {
132    path: PathBuf,
133    messages: Arc<Mutex<Vec<Message>>>,
134}
135
136impl FileHistoryProvider {
137    /// Open (or create) a file-backed history provider at `path`. A missing
138    /// or empty file starts with no history; an existing file is parsed
139    /// eagerly, so a malformed file fails the constructor rather than
140    /// silently discarding history.
141    pub fn new(path: impl Into<PathBuf>) -> Result<Self> {
142        let path = path.into();
143        let messages = if path.exists() {
144            let data = std::fs::read_to_string(&path)
145                .map_err(|e| Error::other(format!("failed to read history file {path:?}: {e}")))?;
146            if data.trim().is_empty() {
147                Vec::new()
148            } else {
149                let value: Value = serde_json::from_str(&data).map_err(|e| {
150                    Error::Serialization(format!("failed to parse history file {path:?}: {e}"))
151                })?;
152                match value.get("messages") {
153                    Some(v) if !v.is_null() => serde_json::from_value(v.clone()).map_err(|e| {
154                        Error::Serialization(format!("failed to parse history file {path:?}: {e}"))
155                    })?,
156                    _ => Vec::new(),
157                }
158            }
159        } else {
160            Vec::new()
161        };
162        Ok(Self {
163            path,
164            messages: Arc::new(Mutex::new(messages)),
165        })
166    }
167
168    /// The path this provider persists to.
169    pub fn path(&self) -> &std::path::Path {
170        &self.path
171    }
172
173    /// The stored messages, in chronological order.
174    pub fn list_messages(&self) -> Vec<Message> {
175        self.messages.lock().unwrap().clone()
176    }
177
178    /// Serialize the stored history to `{"messages": [...]}`.
179    pub fn to_dict(&self) -> Value {
180        serde_json::json!({ "messages": self.list_messages() })
181    }
182
183    fn persist(&self) -> Result<()> {
184        let json = serde_json::to_string_pretty(&self.to_dict())
185            .map_err(|e| Error::Serialization(format!("failed to serialize history: {e}")))?;
186        std::fs::write(&self.path, json)
187            .map_err(|e| Error::other(format!("failed to write history file {:?}: {e}", self.path)))
188    }
189}
190
191#[async_trait]
192impl ContextProvider for FileHistoryProvider {
193    async fn before_run(&self, ctx: &mut SessionContext) -> Result<()> {
194        let stored = self.messages.lock().unwrap().clone();
195        let existing = std::mem::take(&mut ctx.messages);
196        ctx.messages = stored.into_iter().chain(existing).collect();
197        Ok(())
198    }
199
200    async fn after_run(
201        &self,
202        request_messages: &[Message],
203        response_messages: &[Message],
204        error: Option<&Error>,
205    ) -> Result<()> {
206        if error.is_none() {
207            {
208                let mut guard = self.messages.lock().unwrap();
209                guard.extend(request_messages.iter().cloned());
210                guard.extend(response_messages.iter().cloned());
211            }
212            self.persist()?;
213        }
214        Ok(())
215    }
216
217    fn is_history_provider(&self) -> bool {
218        true
219    }
220}
221
222impl HistoryProvider for FileHistoryProvider {}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227    use crate::types::Message;
228
229    #[tokio::test]
230    async fn before_run_prepends_stored_messages_ahead_of_existing_context_messages() {
231        let provider = InMemoryHistoryProvider::with_messages(vec![
232            Message::user("q1"),
233            Message::assistant("a1"),
234        ]);
235        let mut ctx = SessionContext::new(vec![Message::user("q2")]);
236        ctx.messages
237            .push(Message::system("injected by another provider"));
238        provider.before_run(&mut ctx).await.unwrap();
239        let texts: Vec<String> = ctx.messages.iter().map(|m| m.text()).collect();
240        assert_eq!(
241            texts,
242            vec![
243                "q1".to_string(),
244                "a1".to_string(),
245                "injected by another provider".to_string(),
246            ]
247        );
248    }
249
250    #[tokio::test]
251    async fn after_run_appends_only_on_success() {
252        let provider = InMemoryHistoryProvider::new();
253        provider
254            .after_run(&[Message::user("hi")], &[Message::assistant("hello")], None)
255            .await
256            .unwrap();
257        assert_eq!(provider.list_messages().len(), 2);
258
259        // A failed run must not record anything.
260        provider
261            .after_run(
262                &[Message::user("again")],
263                &[],
264                Some(&Error::service("boom")),
265            )
266            .await
267            .unwrap();
268        assert_eq!(provider.list_messages().len(), 2);
269    }
270
271    #[test]
272    fn to_dict_from_dict_round_trips_messages() {
273        let provider = InMemoryHistoryProvider::with_messages(vec![
274            Message::user("q1"),
275            Message::assistant("a1"),
276        ]);
277        let state = provider.to_dict();
278        let restored = InMemoryHistoryProvider::from_dict(&state).unwrap();
279        let msgs = restored.list_messages();
280        assert_eq!(msgs.len(), 2);
281        assert_eq!(msgs[0].text(), "q1");
282        assert_eq!(msgs[1].text(), "a1");
283    }
284
285    #[test]
286    fn from_dict_tolerates_a_missing_messages_key() {
287        let restored = InMemoryHistoryProvider::from_dict(&serde_json::json!({})).unwrap();
288        assert!(restored.list_messages().is_empty());
289    }
290
291    #[test]
292    fn ensure_history_provider_attaches_once_and_skips_service_managed() {
293        let mut local = AgentSession::new();
294        ensure_history_provider(&mut local);
295        assert_eq!(local.context_providers.len(), 1);
296        assert!(local.context_providers[0].is_history_provider());
297        // A second call must not attach a duplicate.
298        ensure_history_provider(&mut local);
299        assert_eq!(local.context_providers.len(), 1);
300
301        let mut service = AgentSession::service("svc-1");
302        ensure_history_provider(&mut service);
303        assert!(service.context_providers.is_empty());
304    }
305
306    #[tokio::test]
307    async fn file_history_provider_persists_and_reloads() {
308        let dir = std::env::temp_dir().join(format!("afr-history-test-{}", uuid::Uuid::new_v4()));
309        std::fs::create_dir_all(&dir).unwrap();
310        let path = dir.join("history.json");
311
312        let provider = FileHistoryProvider::new(&path).unwrap();
313        assert!(provider.list_messages().is_empty());
314        provider
315            .after_run(&[Message::user("hi")], &[Message::assistant("hello")], None)
316            .await
317            .unwrap();
318        assert_eq!(provider.list_messages().len(), 2);
319
320        // A fresh provider opened on the same path picks up the persisted
321        // history.
322        let reloaded = FileHistoryProvider::new(&path).unwrap();
323        let msgs = reloaded.list_messages();
324        assert_eq!(msgs.len(), 2);
325        assert_eq!(msgs[0].text(), "hi");
326        assert_eq!(msgs[1].text(), "hello");
327
328        std::fs::remove_dir_all(&dir).ok();
329    }
330}
agent_framework_core/history.rs

agent_framework_core/
history.rs
