agent_framework_core/
memory.rs

1//! Context / memory providers.
2//!
3//! Rust equivalent of `agent_framework._memory`. A [`ContextProvider`] injects
4//! extra instructions, messages, and tools into an agent invocation without
5//! persisting them to the conversation history.
6//!
7//! Upstream renamed `ContextProvider.invoking`/`invoked` to `before_run`/
8//! `after_run` and removed the `thread_created` hook entirely; `before_run`
9//! mutates a [`SessionContext`] in place instead of returning a value. There
10//! is no aggregate wrapper any more — consumers hold a
11//! `Vec<Arc<dyn ContextProvider>>` and iterate it directly.
12
13use async_trait::async_trait;
14
15use crate::error::{Error, Result};
16use crate::tools::ToolDefinition;
17use crate::types::Message;
18
19/// Per-invocation context a provider contributes to a run. Providers mutate
20/// this in place in before_run. Rust equivalent of upstream SessionContext.
21#[derive(Debug, Clone, Default)]
22pub struct SessionContext {
23    /// Local session identifier (from the thread), for provider scoping.
24    pub session_id: Option<String>,
25    /// Service-managed session/conversation id, when applicable.
26    pub service_session_id: Option<String>,
27    /// The run's input messages (read-only for providers).
28    pub input_messages: Vec<Message>,
29    /// Extra system instructions to inject (providers append via add_instructions).
30    pub instructions: Option<String>,
31    /// Extra context messages to inject ahead of history.
32    pub messages: Vec<Message>,
33    /// Extra tools to make available for this run.
34    pub tools: Vec<ToolDefinition>,
35}
36
37impl SessionContext {
38    pub fn new(input_messages: Vec<Message>) -> Self {
39        Self {
40            input_messages,
41            ..Default::default()
42        }
43    }
44    /// Append instructions, newline-concatenating with any already present.
45    pub fn add_instructions(&mut self, s: impl Into<String>) {
46        let s = s.into();
47        self.instructions = match self.instructions.take() {
48            Some(existing) => Some(format!("{existing}\n{s}")),
49            None => Some(s),
50        };
51    }
52}
53
54/// A source of per-invocation context (memory, RAG, etc.).
55/// Upstream renamed invoking/invoked -> before_run/after_run and REMOVED
56/// thread_created. before_run mutates the SessionContext in place instead of
57/// returning a Context.
58#[async_trait]
59pub trait ContextProvider: Send + Sync {
60    /// Called before the model is invoked; mutate ctx to inject instructions,
61    /// messages, and/or tools. Read ctx.input_messages / ctx.session_id.
62    async fn before_run(&self, ctx: &mut SessionContext) -> Result<()>;
63
64    /// Called after an invocation completes, on BOTH success and failure.
65    /// On success, error is None and response_messages holds the output.
66    /// On failure, error is Some and response_messages is empty.
67    async fn after_run(
68        &self,
69        _request_messages: &[Message],
70        _response_messages: &[Message],
71        _error: Option<&Error>,
72    ) -> Result<()> {
73        Ok(())
74    }
75
76    /// Whether this provider manages conversation history (a
77    /// [`HistoryProvider`](crate::history::HistoryProvider)). [`Agent`](crate::agent::Agent)
78    /// and [`WorkflowAgent`](crate::workflow::WorkflowAgent) use this to
79    /// detect an already-attached history provider among a session's
80    /// `context_providers` and avoid auto-attaching a redundant
81    /// [`InMemoryHistoryProvider`](crate::history::InMemoryHistoryProvider).
82    /// Defaults to `false`; history providers override it to `true`.
83    fn is_history_provider(&self) -> bool {
84        false
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn add_instructions_sets_when_none() {
94        let mut ctx = SessionContext::new(vec![]);
95        assert!(ctx.instructions.is_none());
96        ctx.add_instructions("be brief");
97        assert_eq!(ctx.instructions.as_deref(), Some("be brief"));
98    }
99
100    #[test]
101    fn add_instructions_newline_concatenates() {
102        let mut ctx = SessionContext::new(vec![]);
103        ctx.add_instructions("first");
104        ctx.add_instructions("second");
105        ctx.add_instructions("third");
106        assert_eq!(ctx.instructions.as_deref(), Some("first\nsecond\nthird"));
107    }
108
109    #[test]
110    fn new_sets_input_messages_and_defaults_rest() {
111        let messages = vec![Message::user("hi")];
112        let ctx = SessionContext::new(messages.clone());
113        assert_eq!(ctx.input_messages.len(), messages.len());
114        assert_eq!(ctx.input_messages[0].text(), "hi");
115        assert!(ctx.session_id.is_none());
116        assert!(ctx.service_session_id.is_none());
117        assert!(ctx.instructions.is_none());
118        assert!(ctx.messages.is_empty());
119        assert!(ctx.tools.is_empty());
120    }
121}
agent_framework_core/memory.rs

agent_framework_core/
memory.rs
