autumn_plugin_agent/
session.rs

1//! Sessions: persisted conversations that outlive one run.
2//!
3//! A session is an ordered transcript keyed by a [`SessionId`]. Each run in
4//! a session loads the transcript, adds one turn, and saves it back, so the
5//! agent remembers what was said yesterday. The app chooses the key: a chat
6//! thread, an email chain, a user, or a standing job.
7//!
8//! Long sessions outgrow the context window. [`Compaction`] folds the older
9//! part of a transcript into one summary message (one extra model call) and
10//! keeps the recent turns verbatim.
11
12use std::collections::HashMap;
13use std::fmt::Write as _;
14use std::sync::{Arc, Mutex, PoisonError};
15
16use futures::future::BoxFuture;
17
18use crate::agent::estimate_messages;
19use crate::client::{ChatMessage, ChatRequest, ChatRole, ContentPart, LlmClient, TokenUsage};
20use crate::error::AgentError;
21use crate::ids::SessionId;
22
23/// Persistence for session transcripts.
24///
25/// Transcripts exclude the system prompt: the agent rebuilds it on every run
26/// so memory and skills stay current. [`ChatMessage`] is `Serialize`, so a
27/// database implementation can store the transcript as one JSON column.
28pub trait SessionStore: Send + Sync + std::fmt::Debug {
29    /// Load a transcript. An unknown session is an empty transcript.
30    fn load<'a>(&'a self, id: &'a SessionId)
31    -> BoxFuture<'a, Result<Vec<ChatMessage>, AgentError>>;
32
33    /// Replace a transcript.
34    fn save<'a>(
35        &'a self,
36        id: &'a SessionId,
37        messages: &'a [ChatMessage],
38    ) -> BoxFuture<'a, Result<(), AgentError>>;
39}
40
41/// A process-local [`SessionStore`]. Transcripts vanish on restart and are
42/// not shared between replicas.
43#[derive(Debug, Default)]
44pub struct InMemorySessionStore {
45    sessions: Mutex<HashMap<SessionId, Vec<ChatMessage>>>,
46}
47
48impl InMemorySessionStore {
49    /// An empty store.
50    #[must_use]
51    pub fn new() -> Self {
52        Self::default()
53    }
54
55    /// Share the store.
56    #[must_use]
57    pub fn shared(self) -> Arc<dyn SessionStore> {
58        Arc::new(self)
59    }
60}
61
62impl SessionStore for InMemorySessionStore {
63    fn load<'a>(
64        &'a self,
65        id: &'a SessionId,
66    ) -> BoxFuture<'a, Result<Vec<ChatMessage>, AgentError>> {
67        let messages = self
68            .sessions
69            .lock()
70            .unwrap_or_else(PoisonError::into_inner)
71            .get(id)
72            .cloned()
73            .unwrap_or_default();
74        Box::pin(std::future::ready(Ok(messages)))
75    }
76
77    fn save<'a>(
78        &'a self,
79        id: &'a SessionId,
80        messages: &'a [ChatMessage],
81    ) -> BoxFuture<'a, Result<(), AgentError>> {
82        self.sessions
83            .lock()
84            .unwrap_or_else(PoisonError::into_inner)
85            .insert(id.clone(), messages.to_vec());
86        Box::pin(std::future::ready(Ok(())))
87    }
88}
89
90/// When and how to summarize a long session transcript.
91#[derive(Debug, Clone, Copy, PartialEq, Eq)]
92pub struct Compaction {
93    /// Compact when the stored transcript estimates above this many tokens.
94    pub trigger_tokens: u32,
95    /// Keep roughly this many tokens of the newest turns verbatim.
96    pub keep_recent_tokens: u32,
97    /// Output cap for the summary call.
98    pub summary_max_tokens: u32,
99}
100
101impl Default for Compaction {
102    fn default() -> Self {
103        Self {
104            trigger_tokens: 24_000,
105            keep_recent_tokens: 8_000,
106            summary_max_tokens: 1_024,
107        }
108    }
109}
110
111/// What one compaction did.
112#[derive(Debug, Clone, Copy, PartialEq, Eq)]
113pub struct CompactionReport {
114    /// Messages folded into the summary.
115    pub summarized_messages: usize,
116    /// Estimated transcript tokens before compaction.
117    pub tokens_before: u32,
118    /// Estimated transcript tokens after compaction.
119    pub tokens_after: u32,
120    /// Tokens the summary call spent.
121    pub usage: TokenUsage,
122}
123
124/// The prefix marking a compaction summary message.
125pub const SUMMARY_PREFIX: &str = "[Summary of the earlier conversation]";
126
127const SUMMARY_INSTRUCTIONS: &str = "You compress conversation history for an AI agent that will continue the conversation. Write a dense summary of the transcript: the user's goals and standing instructions, decisions made, facts learned (with exact names, ids, numbers), tool results that still matter, and open tasks. Leave out small talk. Write plain prose or bullets; do not address the user.";
128
129/// Pick where the verbatim tail starts: the oldest user message such that
130/// the tail fits `keep_recent_tokens`. Returns `None` when no user message
131/// after index 0 qualifies (nothing to fold).
132fn split_point(messages: &[ChatMessage], keep_recent_tokens: u32) -> Option<usize> {
133    let mut tail_tokens: u32 = 0;
134    let mut split = None;
135    for index in (1..messages.len()).rev() {
136        let cost = estimate_messages(messages.get(index..=index).unwrap_or_default());
137        tail_tokens = tail_tokens.saturating_add(cost);
138        if tail_tokens > keep_recent_tokens {
139            break;
140        }
141        if messages
142            .get(index)
143            .is_some_and(|message| message.role == ChatRole::User)
144        {
145            split = Some(index);
146        }
147    }
148    split
149}
150
151/// Render a transcript slice as plain text for the summarizer.
152fn render_transcript(messages: &[ChatMessage]) -> String {
153    let mut out = String::new();
154    for message in messages {
155        let role = match message.role {
156            ChatRole::System => "system",
157            ChatRole::User => "user",
158            ChatRole::Assistant => "assistant",
159            ChatRole::Tool => "tool",
160        };
161        for part in &message.content {
162            match part {
163                ContentPart::Text(text) => {
164                    let _ = writeln!(out, "{role}: {text}");
165                }
166                ContentPart::ToolCall {
167                    name, arguments, ..
168                } => {
169                    let _ = writeln!(out, "{role} called {name} with {arguments}");
170                }
171                ContentPart::ToolResult { content, .. } => {
172                    let clipped: String = content.chars().take(600).collect();
173                    let _ = writeln!(out, "{role} result: {clipped}");
174                }
175            }
176        }
177    }
178    out
179}
180
181/// Summarize the older part of a transcript when it exceeds the trigger.
182///
183/// Returns `None` (and leaves `messages` alone) when the transcript is under
184/// the trigger or has no safe split point. Otherwise replaces the older
185/// messages with one user message that starts with [`SUMMARY_PREFIX`]. The
186/// split always lands on a user message, so no tool result loses its call.
187///
188/// # Errors
189///
190/// Returns [`AgentError`] when the summary call fails. The transcript is
191/// unchanged in that case.
192pub async fn compact(
193    client: &dyn LlmClient,
194    messages: &mut Vec<ChatMessage>,
195    settings: &Compaction,
196) -> Result<Option<CompactionReport>, AgentError> {
197    let tokens_before = estimate_messages(messages);
198    if tokens_before <= settings.trigger_tokens {
199        return Ok(None);
200    }
201    let Some(split) = split_point(messages, settings.keep_recent_tokens) else {
202        return Ok(None);
203    };
204    let head = messages.get(..split).unwrap_or_default();
205    let request = ChatRequest {
206        messages: vec![
207            ChatMessage::text(ChatRole::System, SUMMARY_INSTRUCTIONS),
208            ChatMessage::text(ChatRole::User, render_transcript(head)),
209        ],
210        tools: Vec::new(),
211        max_tokens: Some(settings.summary_max_tokens),
212        temperature: None,
213    };
214    let response = client.chat(&request).await?;
215    let summary = response
216        .content
217        .iter()
218        .filter_map(|part| match part {
219            ContentPart::Text(text) => Some(text.as_str()),
220            ContentPart::ToolCall { .. } | ContentPart::ToolResult { .. } => None,
221        })
222        .collect::<Vec<_>>()
223        .join("");
224    let mut compacted = Vec::with_capacity(messages.len().saturating_sub(split).saturating_add(1));
225    compacted.push(ChatMessage::text(
226        ChatRole::User,
227        format!("{SUMMARY_PREFIX}\n{}", summary.trim()),
228    ));
229    compacted.extend(messages.drain(split..));
230    *messages = compacted;
231    Ok(Some(CompactionReport {
232        summarized_messages: split,
233        tokens_before,
234        tokens_after: estimate_messages(messages),
235        usage: response.usage,
236    }))
237}
238
239#[cfg(test)]
240mod tests;
autumn_plugin_agent/session.rs

autumn_plugin_agent/
session.rs
