Skip to main content

rig_memory/
lib.rs

1#![cfg_attr(
2    test,
3    allow(
4        clippy::expect_used,
5        clippy::indexing_slicing,
6        clippy::panic,
7        clippy::unwrap_used,
8        clippy::unreachable
9    )
10)]
11//! Conversation memory policies for the Rig agent framework.
12//!
13//! `rig-core` provides the [`ConversationMemory`] trait and an in-process
14//! [`InMemoryConversationMemory`] backend. This crate adds reusable, named
15//! transformations for shaping loaded history before it is sent to the model:
16//!
17//! - [`NoopMemoryPolicy`] — identity, returns input unchanged.
18//! - [`SlidingWindowMemory`] — retains the most recent `N` messages.
19//! - [`TokenWindowMemory`] — retains messages that fit within a token budget.
20//! - [`HeuristicTokenCounter`] — provider-agnostic, zero-dependency
21//!   [`TokenCounter`] that approximates token cost from character lengths.
22//! - [`DemotionHook`] + [`DemotingPolicyMemory`] — bridge truncated turns
23//!   from a [`MemoryPolicy`] into a long-tail store.
24//! - [`Compactor`] + [`CompactingMemory`] — replace truncated turns with a
25//!   derived summary artifact (rolling-summary semantics).
26//! - [`TemplateCompactor`] — zero-dependency reference [`Compactor`] that
27//!   produces a textual rollup without calling an LLM.
28//!
29//! All sliding policies drop a leading orphan tool-result message when the
30//! preceding assistant tool call has been truncated, since most providers
31//! reject unpaired tool results.
32//!
33//! # Example
34//!
35//! ```
36//! use rig_memory::{InMemoryConversationMemory, IntoFilter, SlidingWindowMemory};
37//!
38//! let memory = InMemoryConversationMemory::new()
39//!     .with_filter(SlidingWindowMemory::last_messages(20).into_filter());
40//! ```
41
42use std::{
43    collections::HashMap,
44    sync::{Arc, Mutex as StdMutex},
45};
46
47/// Re-exports of the core memory abstractions so callers only need a single
48/// dependency on `rig-memory` for both the trait/backend and the policies.
49pub use rig_core::memory::{
50    Compactor, ConversationMemory, DemotionHook, InMemoryConversationMemory, MemoryError,
51    NoopDemotionHook,
52};
53
54use rig_core::completion::Message;
55use rig_core::message::UserContent;
56use rig_core::wasm_compat::{WasmBoxedFuture, WasmCompatSend, WasmCompatSync};
57
58/// A transformation applied to messages loaded from a [`ConversationMemory`].
59///
60/// Policies typically truncate, summarize, or re-order history. They are
61/// pure, fallible message transformers: implementors that cannot fail should
62/// always return `Ok`.
63pub trait MemoryPolicy: WasmCompatSend + WasmCompatSync {
64    /// Transform `messages` into the history that should be returned to the
65    /// agent. This is the required method — every policy must implement it.
66    fn apply(&self, messages: Vec<Message>) -> Result<Vec<Message>, MemoryError>;
67
68    /// Transform `messages` and report which messages were demoted (excluded
69    /// from the returned history).
70    ///
71    /// Returns `(kept, demoted)`. The default implementation returns
72    /// `(self.apply(messages)?, Vec::new())`, which is correct for
73    /// non-truncating policies. Truncating policies (sliding window, token
74    /// window, …) override this method to populate `demoted` with the
75    /// messages they evicted.
76    ///
77    /// Implementors must guarantee that `demoted` is the prefix of the
78    /// original input not retained in `kept`, in original order. Composing
79    /// adapters such as [`DemotingPolicyMemory`] rely on this contract to
80    /// track delivery watermarks correctly.
81    fn apply_with_demoted(
82        &self,
83        messages: Vec<Message>,
84    ) -> Result<(Vec<Message>, Vec<Message>), MemoryError> {
85        Ok((self.apply(messages)?, Vec::new()))
86    }
87}
88
89impl<P> MemoryPolicy for Arc<P>
90where
91    P: MemoryPolicy + ?Sized,
92{
93    fn apply(&self, messages: Vec<Message>) -> Result<Vec<Message>, MemoryError> {
94        (**self).apply(messages)
95    }
96
97    fn apply_with_demoted(
98        &self,
99        messages: Vec<Message>,
100    ) -> Result<(Vec<Message>, Vec<Message>), MemoryError> {
101        (**self).apply_with_demoted(messages)
102    }
103}
104
105impl<P> MemoryPolicy for Box<P>
106where
107    P: MemoryPolicy + ?Sized,
108{
109    fn apply(&self, messages: Vec<Message>) -> Result<Vec<Message>, MemoryError> {
110        (**self).apply(messages)
111    }
112
113    fn apply_with_demoted(
114        &self,
115        messages: Vec<Message>,
116    ) -> Result<(Vec<Message>, Vec<Message>), MemoryError> {
117        (**self).apply_with_demoted(messages)
118    }
119}
120
121/// Adapt a [`MemoryPolicy`] into a closure suitable for
122/// [`InMemoryConversationMemory::with_filter`].
123///
124/// Errors raised by the policy are swallowed because `with_filter` does not
125/// propagate failures. Use [`MemoryPolicy::apply`] directly when you need to
126/// observe policy errors.
127pub trait IntoFilter: MemoryPolicy + Sized + 'static {
128    /// Convert this policy into a filter closure.
129    ///
130    /// On policy error the original input is returned unchanged and a
131    /// `tracing::warn!` is emitted, so a transient policy bug degrades
132    /// gracefully (the model still sees the unfiltered history) instead of
133    /// silently erasing context.
134    #[cfg(not(target_family = "wasm"))]
135    fn into_filter(self) -> Box<dyn Fn(Vec<Message>) -> Vec<Message> + Send + Sync> {
136        let policy = Arc::new(self);
137        Box::new(move |msgs| {
138            let fallback = msgs.clone();
139            match policy.apply(msgs) {
140                Ok(out) => out,
141                Err(err) => {
142                    tracing::warn!(error = %err, "memory policy failed; returning unfiltered history");
143                    fallback
144                }
145            }
146        })
147    }
148
149    /// Convert this policy into a filter closure.
150    ///
151    /// On policy error the original input is returned unchanged and a
152    /// `tracing::warn!` is emitted, so a transient policy bug degrades
153    /// gracefully (the model still sees the unfiltered history) instead of
154    /// silently erasing context.
155    #[cfg(target_family = "wasm")]
156    fn into_filter(self) -> Box<dyn Fn(Vec<Message>) -> Vec<Message>> {
157        let policy = Arc::new(self);
158        Box::new(move |msgs| {
159            let fallback = msgs.clone();
160            match policy.apply(msgs) {
161                Ok(out) => out,
162                Err(err) => {
163                    tracing::warn!(error = %err, "memory policy failed; returning unfiltered history");
164                    fallback
165                }
166            }
167        })
168    }
169}
170
171impl<P> IntoFilter for P where P: MemoryPolicy + 'static {}
172
173/// A [`MemoryPolicy`] that returns its input unchanged.
174#[derive(Debug, Default, Clone, Copy)]
175pub struct NoopMemoryPolicy;
176
177impl MemoryPolicy for NoopMemoryPolicy {
178    fn apply(&self, messages: Vec<Message>) -> Result<Vec<Message>, MemoryError> {
179        Ok(messages)
180    }
181}
182
183/// A [`MemoryPolicy`] that retains only the most recent `max_messages` entries.
184///
185/// When the window starts mid-conversation, a leading orphan tool-result
186/// message (a [`Message::User`] whose first content is a tool result without
187/// its preceding [`Message::Assistant`] tool call) is dropped to preserve the
188/// tool-call/result pairing required by most providers.
189#[derive(Debug, Clone, Copy)]
190pub struct SlidingWindowMemory {
191    max_messages: usize,
192}
193
194impl SlidingWindowMemory {
195    /// Keep at most `n` messages.
196    pub fn last_messages(n: usize) -> Self {
197        Self { max_messages: n }
198    }
199}
200
201impl MemoryPolicy for SlidingWindowMemory {
202    fn apply(&self, messages: Vec<Message>) -> Result<Vec<Message>, MemoryError> {
203        Ok(self.apply_with_demoted(messages)?.0)
204    }
205
206    fn apply_with_demoted(
207        &self,
208        messages: Vec<Message>,
209    ) -> Result<(Vec<Message>, Vec<Message>), MemoryError> {
210        if messages.len() <= self.max_messages {
211            return Ok((messages, Vec::new()));
212        }
213
214        let start = messages.len() - self.max_messages;
215        let mut iter = messages.into_iter();
216        let mut demoted: Vec<Message> = (&mut iter).take(start).collect();
217        let mut window: Vec<Message> = iter.collect();
218
219        // The orphan tool-result, if any, becomes part of the demoted set so
220        // it is preserved end-to-end through the demotion hook even though
221        // the model never sees it again.
222        if let Some(Message::User { content }) = window.first()
223            && matches!(content.first_ref(), UserContent::ToolResult(_))
224        {
225            demoted.push(window.remove(0));
226        }
227
228        Ok((window, demoted))
229    }
230}
231
232/// Counts the tokens contributed by a single [`Message`].
233///
234/// Implementors should pick a counting strategy appropriate for their target
235/// provider (for example, `tiktoken-rs` for OpenAI). Counting must be cheap;
236/// it runs once per message on every memory load.
237pub trait TokenCounter: WasmCompatSend + WasmCompatSync {
238    /// Approximate the number of tokens contributed by `message`.
239    fn count(&self, message: &Message) -> usize;
240}
241
242impl<F> TokenCounter for F
243where
244    F: Fn(&Message) -> usize + WasmCompatSend + WasmCompatSync,
245{
246    fn count(&self, message: &Message) -> usize {
247        (self)(message)
248    }
249}
250
251impl<C> TokenCounter for Arc<C>
252where
253    C: TokenCounter + ?Sized,
254{
255    fn count(&self, message: &Message) -> usize {
256        (**self).count(message)
257    }
258}
259
260impl TokenCounter for Box<dyn TokenCounter> {
261    fn count(&self, message: &Message) -> usize {
262        (**self).count(message)
263    }
264}
265
266/// A provider-agnostic [`TokenCounter`] that approximates token counts from
267/// UTF-8 byte lengths.
268///
269/// This is intended as a zero-dependency default. It is **not** a substitute
270/// for a tokenizer and will under- or over-count by up to ~30 % on real
271/// content, but it is monotonic in message size and stable across runs, which
272/// is enough for [`TokenWindowMemory`] to enforce a budget that *trends*
273/// with provider billing.
274///
275/// # Strategy
276///
277/// For every text-bearing block (`Text`, reasoning text, tool-result text)
278/// the counter sums UTF-8 byte lengths (`str::len`, an O(1) call) and divides
279/// by `bytes_per_token`, rounded up. Bytes are used instead of Unicode
280/// scalars because the cost is O(1), modern BPE tokenizers operate on byte
281/// sequences, and per-message budgeting only needs the rough order of
282/// magnitude. For ASCII text bytes and characters coincide; for non-ASCII
283/// text the counter slightly over-estimates, which is the safe direction
284/// for a hard budget.
285///
286/// Tool calls are charged the JSON-serialised length of their `ToolFunction`
287/// payload. Each message is charged a flat `per_message_overhead` to model
288/// the per-turn role/separator tokens that providers add internally. Non-text
289/// blocks (images, audio, video, documents) are charged
290/// `per_attachment_tokens` each because their real cost is provider-specific
291/// and rarely text-derived.
292///
293/// # Presets
294///
295/// The defaults match OpenAI's published rule of thumb (~4 bytes per token,
296/// ~4 tokens of per-message overhead). [`HeuristicTokenCounter::anthropic`]
297/// uses a slightly denser ratio that better fits Claude's tokenizer.
298///
299/// # Example
300///
301/// ```
302/// use rig_memory::{HeuristicTokenCounter, TokenWindowMemory};
303///
304/// let policy = TokenWindowMemory::new(2_000, HeuristicTokenCounter::default());
305/// # let _ = policy;
306/// ```
307#[derive(Debug, Clone, Copy)]
308pub struct HeuristicTokenCounter {
309    bytes_per_token: f32,
310    per_message_overhead: usize,
311    per_attachment_tokens: usize,
312}
313
314impl HeuristicTokenCounter {
315    /// Create a counter with explicit parameters.
316    ///
317    /// `bytes_per_token` is clamped to a minimum of `1.0` so the counter
318    /// never panics or produces zero-cost messages on degenerate input.
319    pub fn new(
320        bytes_per_token: f32,
321        per_message_overhead: usize,
322        per_attachment_tokens: usize,
323    ) -> Self {
324        let bytes_per_token = if bytes_per_token.is_finite() && bytes_per_token >= 1.0 {
325            bytes_per_token
326        } else {
327            1.0
328        };
329        Self {
330            bytes_per_token,
331            per_message_overhead,
332            per_attachment_tokens,
333        }
334    }
335
336    /// Preset matching OpenAI's chat-completion token rule of thumb.
337    ///
338    /// Equivalent to [`HeuristicTokenCounter::default`].
339    pub fn openai() -> Self {
340        Self::new(4.0, 4, 256)
341    }
342
343    /// Preset tuned for Anthropic Claude's tokenizer.
344    pub fn anthropic() -> Self {
345        Self::new(3.5, 4, 256)
346    }
347
348    /// Preset tuned for Google Gemini.
349    pub fn gemini() -> Self {
350        Self::new(4.0, 4, 256)
351    }
352
353    fn bytes_to_tokens(&self, bytes: usize) -> usize {
354        // `bytes_per_token` is clamped to >= 1.0 in the constructor, so the
355        // division is well-defined. We round up so a single non-empty
356        // input still costs at least one token.
357        let tokens = (bytes as f32) / self.bytes_per_token;
358        tokens.ceil() as usize
359    }
360
361    fn count_user(&self, content: &rig_core::message::UserContent) -> usize {
362        use rig_core::message::UserContent;
363        match content {
364            UserContent::Text(text) => self.bytes_to_tokens(text.text.len()),
365            UserContent::ToolResult(result) => result
366                .content
367                .iter()
368                .map(|c| match c {
369                    rig_core::message::ToolResultContent::Text(t) => {
370                        self.bytes_to_tokens(t.text.len())
371                    }
372                    rig_core::message::ToolResultContent::Json { value } => {
373                        self.bytes_to_tokens(value.to_string().len())
374                    }
375                    rig_core::message::ToolResultContent::Image(_) => self.per_attachment_tokens,
376                })
377                .sum(),
378            UserContent::Image(_)
379            | UserContent::Audio(_)
380            | UserContent::Video(_)
381            | UserContent::Document(_) => self.per_attachment_tokens,
382        }
383    }
384
385    fn count_assistant(&self, content: &rig_core::message::AssistantContent) -> usize {
386        use rig_core::message::AssistantContent;
387        match content {
388            AssistantContent::Text(text) => self.bytes_to_tokens(text.text.len()),
389            AssistantContent::Reasoning(reasoning) => {
390                self.bytes_to_tokens(reasoning.display_text().len())
391            }
392            AssistantContent::ToolCall(call) => {
393                let name_bytes = call.function.name.len();
394                // `serde_json::Value::to_string` is the canonical compact JSON
395                // encoding and never fails, so we charge tool calls by the
396                // length of their serialised arguments without pulling in a
397                // direct `serde_json` dependency.
398                let args_bytes = call.function.arguments.to_string().len();
399                self.bytes_to_tokens(name_bytes + args_bytes)
400            }
401            AssistantContent::Image(_) => self.per_attachment_tokens,
402        }
403    }
404}
405
406impl Default for HeuristicTokenCounter {
407    fn default() -> Self {
408        Self::openai()
409    }
410}
411
412impl TokenCounter for HeuristicTokenCounter {
413    fn count(&self, message: &Message) -> usize {
414        let content_tokens: usize = match message {
415            Message::User { content } => content.iter().map(|c| self.count_user(c)).sum(),
416            Message::Assistant { content, .. } => {
417                content.iter().map(|c| self.count_assistant(c)).sum()
418            }
419            Message::System { content } => self.bytes_to_tokens(content.len()),
420        };
421        content_tokens.saturating_add(self.per_message_overhead)
422    }
423}
424
425/// A [`MemoryPolicy`] that retains the most recent messages up to a token budget.
426///
427/// Messages are walked from newest to oldest, accumulating token counts
428/// produced by a [`TokenCounter`]. Once including a message would exceed
429/// `max_tokens`, the walk stops and the included messages are returned in
430/// original (oldest-first) order. As with [`SlidingWindowMemory`], a leading
431/// orphan tool-result is dropped when its paired assistant tool call has
432/// been truncated.
433pub struct TokenWindowMemory {
434    max_tokens: usize,
435    counter: Arc<dyn TokenCounter>,
436}
437
438impl TokenWindowMemory {
439    /// Create a new policy with a token budget and a counter.
440    pub fn new<C>(max_tokens: usize, counter: C) -> Self
441    where
442        C: TokenCounter + 'static,
443    {
444        Self {
445            max_tokens,
446            counter: Arc::new(counter),
447        }
448    }
449}
450
451impl std::fmt::Debug for TokenWindowMemory {
452    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
453        f.debug_struct("TokenWindowMemory")
454            .field("max_tokens", &self.max_tokens)
455            .field("counter", &"<counter>")
456            .finish()
457    }
458}
459
460impl MemoryPolicy for TokenWindowMemory {
461    fn apply(&self, messages: Vec<Message>) -> Result<Vec<Message>, MemoryError> {
462        Ok(self.apply_with_demoted(messages)?.0)
463    }
464
465    fn apply_with_demoted(
466        &self,
467        messages: Vec<Message>,
468    ) -> Result<(Vec<Message>, Vec<Message>), MemoryError> {
469        let mut budget = self.max_tokens;
470        let mut keep_from = messages.len();
471
472        for (idx, msg) in messages.iter().enumerate().rev() {
473            let cost = self.counter.count(msg);
474            if cost > budget {
475                break;
476            }
477            budget -= cost;
478            keep_from = idx;
479        }
480
481        let mut iter = messages.into_iter();
482        let mut demoted: Vec<Message> = (&mut iter).take(keep_from).collect();
483        let mut window: Vec<Message> = iter.collect();
484
485        if let Some(Message::User { content }) = window.first()
486            && matches!(content.first_ref(), UserContent::ToolResult(_))
487        {
488            demoted.push(window.remove(0));
489        }
490
491        Ok((window, demoted))
492    }
493}
494
495/// Wrap a [`ConversationMemory`] backend with a [`MemoryPolicy`], propagating
496/// policy errors to the caller as [`MemoryError::Policy`].
497///
498/// This is the hard-fail counterpart to
499/// [`InMemoryConversationMemory::with_filter`] + [`IntoFilter::into_filter`].
500/// `with_filter` swallows policy errors and returns the unfiltered history;
501/// `PolicyMemory` surfaces them so callers can decide how to react.
502///
503/// # Example
504///
505/// ```no_run
506/// use rig_memory::{InMemoryConversationMemory, PolicyMemory, SlidingWindowMemory};
507///
508/// let memory = PolicyMemory::new(
509///     InMemoryConversationMemory::new(),
510///     SlidingWindowMemory::last_messages(20),
511/// );
512/// ```
513#[derive(Debug, Clone, Copy)]
514pub struct PolicyMemory<M, P> {
515    inner: M,
516    policy: P,
517}
518
519impl<M, P> PolicyMemory<M, P> {
520    /// Wrap `inner` so every loaded history is run through `policy`.
521    pub fn new(inner: M, policy: P) -> Self {
522        Self { inner, policy }
523    }
524
525    /// Return a reference to the wrapped backend.
526    pub fn inner(&self) -> &M {
527        &self.inner
528    }
529
530    /// Return a reference to the wrapped policy.
531    pub fn policy(&self) -> &P {
532        &self.policy
533    }
534
535    /// Consume the wrapper and return the underlying backend and policy.
536    pub fn into_inner(self) -> (M, P) {
537        (self.inner, self.policy)
538    }
539}
540
541impl<M, P> ConversationMemory for PolicyMemory<M, P>
542where
543    M: ConversationMemory,
544    P: MemoryPolicy,
545{
546    fn load<'a>(
547        &'a self,
548        conversation_id: &'a str,
549    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
550        Box::pin(async move {
551            let messages = self.inner.load(conversation_id).await?;
552            self.policy.apply(messages)
553        })
554    }
555
556    fn append<'a>(
557        &'a self,
558        conversation_id: &'a str,
559        messages: Vec<Message>,
560    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
561        self.inner.append(conversation_id, messages)
562    }
563
564    fn clear<'a>(
565        &'a self,
566        conversation_id: &'a str,
567    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
568        self.inner.clear(conversation_id)
569    }
570}
571
572/// A [`ConversationMemory`] adapter that wraps a backend with a
573/// [`MemoryPolicy`] **and** a [`DemotionHook`], so messages truncated by the
574/// policy flow into the hook before the active window is returned.
575///
576/// `DemotingPolicyMemory` is the bridge between the recent-turn store
577/// ([`InMemoryConversationMemory`] or any other [`ConversationMemory`]) and a
578/// long-tail store (`MemvidPersistHook`, vector RAG, archival storage, …).
579/// Compose it with any [`MemoryPolicy`] that overrides
580/// [`MemoryPolicy::apply_with_demoted`]; policies that rely on the default
581/// implementation will still load correctly but will never demote anything.
582///
583/// # Concurrency
584///
585/// Concurrent [`ConversationMemory::load`] calls on the same
586/// `conversation_id` are serialised at the demotion seam: only one call at
587/// a time delivers messages to the hook for a given conversation. Other
588/// concurrent loads for that conversation observe the in-flight delivery
589/// and return the truncated `kept` history immediately without firing the
590/// hook again. Pending demotions that were skipped this way are picked up
591/// by the next `load` after the in-flight delivery completes.
592///
593/// **Failure visibility.** A hook error is returned only to the caller
594/// whose `load` actually drove the delivery. Concurrent callers that
595/// short-circuited on `in_flight` see `Ok(kept)` even if the in-flight
596/// delivery ultimately failed; the watermark stays unchanged so the next
597/// `load` retries. Callers that rely on the hook for durability should
598/// treat a successful `load` as best-effort with respect to demotion and
599/// surface hook failures through the hook's own observability (logs,
600/// metrics, dead-letter buffer) rather than the `load` return value.
601///
602/// # Persistence
603///
604/// Delivery watermarks are kept in process memory only. Across process
605/// restarts, the hook will receive previously-delivered demotions again;
606/// see the [`DemotionHook`] idempotency contract.
607///
608/// # Example
609///
610/// ```no_run
611/// use rig_memory::{
612///     DemotingPolicyMemory, DemotionHook, InMemoryConversationMemory,
613///     MemoryError, NoopDemotionHook, SlidingWindowMemory,
614/// };
615///
616/// let memory = DemotingPolicyMemory::new(
617///     InMemoryConversationMemory::new(),
618///     SlidingWindowMemory::last_messages(20),
619///     NoopDemotionHook,
620/// );
621/// # let _ = memory;
622/// ```
623pub struct DemotingPolicyMemory<M, P, H> {
624    inner: M,
625    policy: P,
626    hook: H,
627    state: StdMutex<HashMap<String, ConversationDemotionState>>,
628}
629
630type InFlightReservation = Arc<()>;
631
632#[derive(Debug, Default, Clone)]
633struct ConversationDemotionState {
634    /// Number of demoted messages already delivered to the hook within
635    /// this process lifetime. Advanced only on hook success.
636    delivered: usize,
637    /// Reservation held while a `load` is currently awaiting
638    /// `hook.on_demote(...)` for this conversation. Other concurrent loads
639    /// observe this and short-circuit without re-delivering the same messages.
640    in_flight: Option<InFlightReservation>,
641}
642
643impl<M, P, H> DemotingPolicyMemory<M, P, H> {
644    /// Wrap `inner` so every load runs through `policy` and demoted messages
645    /// flow into `hook`.
646    pub fn new(inner: M, policy: P, hook: H) -> Self {
647        Self {
648            inner,
649            policy,
650            hook,
651            state: StdMutex::new(HashMap::new()),
652        }
653    }
654
655    /// Return a reference to the wrapped backend.
656    pub fn inner(&self) -> &M {
657        &self.inner
658    }
659
660    /// Return a reference to the wrapped policy.
661    pub fn policy(&self) -> &P {
662        &self.policy
663    }
664
665    /// Return a reference to the demotion hook.
666    pub fn hook(&self) -> &H {
667        &self.hook
668    }
669
670    /// Consume the wrapper and return its three components.
671    pub fn into_inner(self) -> (M, P, H) {
672        (self.inner, self.policy, self.hook)
673    }
674
675    /// Drop the in-process delivery watermark for `conversation_id`.
676    ///
677    /// Call this when a conversation has ended to bound memory usage.
678    /// The watermark map is otherwise unbounded — entries persist for
679    /// the lifetime of the wrapper.
680    ///
681    /// If the internal state lock has been poisoned by a panic in another
682    /// thread, this is a no-op (the watermark will be dropped naturally
683    /// when the wrapper itself is dropped).
684    pub fn forget(&self, conversation_id: &str) {
685        if let Ok(mut guard) = self.state.lock() {
686            guard.remove(conversation_id);
687        }
688    }
689
690    /// Number of conversations currently tracked in the watermark map.
691    /// Useful for telemetry and leak detection. Returns `0` if the internal
692    /// state lock is poisoned.
693    pub fn tracked_conversations(&self) -> usize {
694        self.state.lock().map(|g| g.len()).unwrap_or(0)
695    }
696}
697
698impl<M, P, H> std::fmt::Debug for DemotingPolicyMemory<M, P, H>
699where
700    M: std::fmt::Debug,
701    P: std::fmt::Debug,
702{
703    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
704        f.debug_struct("DemotingPolicyMemory")
705            .field("inner", &self.inner)
706            .field("policy", &self.policy)
707            .field("hook", &"<hook>")
708            .finish()
709    }
710}
711
712impl<M, P, H> ConversationMemory for DemotingPolicyMemory<M, P, H>
713where
714    M: ConversationMemory,
715    P: MemoryPolicy,
716    H: DemotionHook,
717{
718    fn load<'a>(
719        &'a self,
720        conversation_id: &'a str,
721    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
722        Box::pin(async move {
723            let messages = self.inner.load(conversation_id).await?;
724            let (kept, mut demoted) = self.policy.apply_with_demoted(messages)?;
725            let demoted_count = demoted.len();
726
727            // Reserve a delivery slot atomically. Decide-and-mark must
728            // happen under one short-lived lock so concurrent loads on
729            // the same conversation_id can't both observe the same
730            // delivered watermark and double-fire the hook.
731            //
732            // Fast path: if the conversation is already tracked, mutate in
733            // place. Only allocate a new `String` key when we are about to
734            // record state for a conversation we have not seen before *and*
735            // there is actually demotion work to track.
736            let (pending, reservation) = {
737                let mut guard = self.state.lock().map_err(poisoned)?;
738                if let Some(entry) = guard.get_mut(conversation_id) {
739                    if entry.in_flight.is_some() {
740                        // Another load is mid-delivery for this conversation;
741                        // skip and let the next load see whatever it leaves
742                        // behind.
743                        return Ok(kept);
744                    }
745                    if entry.delivered >= demoted_count {
746                        (Vec::new(), None)
747                    } else {
748                        let split = entry.delivered;
749                        let reservation = Arc::new(());
750                        entry.in_flight = Some(reservation.clone());
751                        (demoted.split_off(split), Some(reservation))
752                    }
753                } else if demoted_count == 0 {
754                    // First load for this conversation and nothing was
755                    // demoted: no need to allocate a tracking entry yet.
756                    (Vec::new(), None)
757                } else {
758                    let reservation = Arc::new(());
759                    guard.insert(
760                        conversation_id.to_string(),
761                        ConversationDemotionState {
762                            delivered: 0,
763                            in_flight: Some(reservation.clone()),
764                        },
765                    );
766                    (std::mem::take(&mut demoted), Some(reservation))
767                }
768            };
769
770            let Some(reservation) = reservation else {
771                return Ok(kept);
772            };
773
774            // Arm an RAII guard so the in-flight gate is released even if
775            // this future is dropped mid-await (caller cancellation) or the
776            // hook panics. The reservation token prevents stale guards from
777            // clearing newer in-flight loads after clear()/forget() reuse the
778            // same conversation id.
779            let in_flight_guard =
780                DemotionInFlightGuard::new(&self.state, conversation_id, reservation.clone());
781
782            let result = self.hook.on_demote(conversation_id, pending).await;
783
784            // Reacquire briefly to advance the watermark on success and
785            // always clear the in-flight flag so a future load can retry.
786            //
787            // Only update if the entry still exists: a concurrent `clear`
788            // (and matching `forget`) for this `conversation_id` may have
789            // dropped the watermark entry while the hook was awaiting. In
790            // that case we must not resurrect it with a stale `delivered`
791            // count — the next load on a freshly-populated backend would
792            // then skip a real demotion.
793            {
794                let mut guard = self.state.lock().map_err(poisoned)?;
795                if let Some(entry) = guard.get_mut(conversation_id)
796                    && entry
797                        .in_flight
798                        .as_ref()
799                        .is_some_and(|current| Arc::ptr_eq(current, &reservation))
800                {
801                    entry.in_flight = None;
802                    if result.is_ok() {
803                        entry.delivered = demoted_count;
804                    }
805                }
806            }
807            in_flight_guard.disarm();
808            result?;
809            Ok(kept)
810        })
811    }
812
813    fn append<'a>(
814        &'a self,
815        conversation_id: &'a str,
816        messages: Vec<Message>,
817    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
818        self.inner.append(conversation_id, messages)
819    }
820
821    fn clear<'a>(
822        &'a self,
823        conversation_id: &'a str,
824    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
825        Box::pin(async move {
826            self.inner.clear(conversation_id).await?;
827            self.forget(conversation_id);
828            Ok(())
829        })
830    }
831}
832
833fn poisoned<E: std::fmt::Display>(err: E) -> MemoryError {
834    MemoryError::Internal(err.to_string())
835}
836
837/// RAII guard that clears the `in_flight` flag for a conversation in the
838/// shared demotion state map when dropped, unless the consumer explicitly
839/// disarms it after a successful post-await update.
840///
841/// This prevents the in-flight gate from leaking when the awaiting
842/// `load(...)` future is dropped (caller timeout, `tokio::select!`, etc.)
843/// or when the hook panics. A missing entry is a no-op, covering the case
844/// where a concurrent `clear` removed the conversation while delivery was
845/// awaiting.
846struct DemotionInFlightGuard<'a> {
847    state: &'a StdMutex<HashMap<String, ConversationDemotionState>>,
848    key: &'a str,
849    reservation: InFlightReservation,
850    armed: bool,
851}
852
853impl<'a> DemotionInFlightGuard<'a> {
854    fn new(
855        state: &'a StdMutex<HashMap<String, ConversationDemotionState>>,
856        key: &'a str,
857        reservation: InFlightReservation,
858    ) -> Self {
859        Self {
860            state,
861            key,
862            reservation,
863            armed: true,
864        }
865    }
866
867    /// Disable the `Drop` clean-up. Call after the post-await state
868    /// update has already cleared `in_flight` while holding the lock.
869    fn disarm(mut self) {
870        self.armed = false;
871    }
872}
873
874impl Drop for DemotionInFlightGuard<'_> {
875    fn drop(&mut self) {
876        if !self.armed {
877            return;
878        }
879        if let Ok(mut guard) = self.state.lock()
880            && let Some(entry) = guard.get_mut(self.key)
881            && entry
882                .in_flight
883                .as_ref()
884                .is_some_and(|current| Arc::ptr_eq(current, &self.reservation))
885        {
886            entry.in_flight = None;
887        }
888    }
889}
890
891/// RAII guard that clears the `in_flight` flag for a conversation in the
892/// shared compaction state map when dropped, unless the consumer
893/// explicitly disarms it after a successful post-await update.
894///
895/// This prevents the in-flight gate from leaking when the awaiting
896/// `load(...)` future is dropped (caller timeout, `tokio::select!`, etc.)
897/// or when the compactor panics: in either case `Drop` runs and releases
898/// the gate so subsequent loads can retry. A missing entry is a no-op,
899/// covering the case where a concurrent `clear` removed the conversation
900/// while compaction was awaiting.
901struct InFlightGuard<'a, A> {
902    state: &'a StdMutex<HashMap<String, ConversationCompactionState<A>>>,
903    key: &'a str,
904    reservation: InFlightReservation,
905    armed: bool,
906}
907
908impl<'a, A> InFlightGuard<'a, A> {
909    fn new(
910        state: &'a StdMutex<HashMap<String, ConversationCompactionState<A>>>,
911        key: &'a str,
912        reservation: InFlightReservation,
913    ) -> Self {
914        Self {
915            state,
916            key,
917            reservation,
918            armed: true,
919        }
920    }
921
922    /// Disable the `Drop` clean-up. Call after the post-await state
923    /// update has already cleared `in_flight` while holding the lock.
924    fn disarm(mut self) {
925        self.armed = false;
926    }
927}
928
929impl<A> Drop for InFlightGuard<'_, A> {
930    fn drop(&mut self) {
931        if !self.armed {
932            return;
933        }
934        if let Ok(mut guard) = self.state.lock()
935            && let Some(entry) = guard.get_mut(self.key)
936            && entry
937                .in_flight
938                .as_ref()
939                .is_some_and(|current| Arc::ptr_eq(current, &self.reservation))
940        {
941            entry.in_flight = None;
942        }
943    }
944}
945
946/// A [`ConversationMemory`] adapter that wraps a backend with a
947/// [`MemoryPolicy`] **and** a [`Compactor`], replacing truncated turns with
948/// a summary artifact spliced at the front of the loaded history.
949///
950/// `CompactingMemory` is the next layer above [`DemotingPolicyMemory`]: a
951/// demotion hook only *observes* what the policy evicted, while a compactor
952/// *substitutes* the evicted prefix with a derived [`Message`]. The loaded
953/// history shape is therefore `[summary_message, ...kept_window]` whenever
954/// any compaction has occurred for the conversation, and just `kept_window`
955/// otherwise. The summary itself is recomputed (rolled forward) on every
956/// load that produces newly-evicted messages, so older summaries are folded
957/// into newer ones via the compactor's `carry_over` parameter.
958///
959/// # Concurrency
960///
961/// Concurrent [`ConversationMemory::load`] calls on the same
962/// `conversation_id` are serialised at the compaction seam: only one call
963/// at a time invokes the compactor for a given conversation. Other
964/// concurrent loads observe the in-flight compaction and immediately
965/// return the previously-stored summary spliced in front of `kept`,
966/// without re-running the compactor. Newly-evicted messages skipped this
967/// way are folded into the next compaction.
968///
969/// **Failure visibility.** A compactor error is returned only to the
970/// caller whose `load` actually drove the compaction. Concurrent callers
971/// that short-circuited on `in_flight` see `Ok([old_summary?, ...kept])`
972/// even if the in-flight compaction ultimately failed; the watermark
973/// stays unchanged so the next `load` retries.
974///
975/// # Persistence
976///
977/// The carry-over summary and delivery watermarks are kept in process
978/// memory only. Across process restarts, the first load on each
979/// conversation re-evicts and re-compacts the same prefix; compactors
980/// that have side effects (LLM calls, persistent writes) should
981/// deduplicate.
982///
983/// # Prompt shape and budgets
984///
985/// `CompactingMemory` is **policy-agnostic**: the wrapped
986/// [`MemoryPolicy`] decides which messages are kept versus demoted, and
987/// only the kept window is bounded by that policy. The summary artifact
988/// produced by the [`Compactor`] is spliced **outside** that budget — so
989/// the loaded prompt has shape `[summary, ...kept_window]` where
990/// `kept_window` respects the policy's bounds and `summary` adds an
991/// extra message on top of it.
992///
993/// Callers that combine `CompactingMemory` with a token-budgeted policy
994/// (e.g. [`TokenWindowMemory`]) **must use a [`Compactor`] that bounds
995/// its own artifact**, or accept that the loaded prompt may exceed the
996/// policy's budget by the size of the summary. The reference
997/// [`TemplateCompactor`] grows monotonically by default; configure it
998/// with [`TemplateCompactor::with_max_bytes`] to cap the rolled-up text.
999///
1000/// # Example
1001///
1002/// ```no_run
1003/// use rig_memory::{
1004///     CompactingMemory, InMemoryConversationMemory, SlidingWindowMemory,
1005///     TemplateCompactor,
1006/// };
1007///
1008/// let memory = CompactingMemory::new(
1009///     InMemoryConversationMemory::new(),
1010///     SlidingWindowMemory::last_messages(20),
1011///     TemplateCompactor::new(),
1012/// );
1013/// # let _ = memory;
1014/// ```
1015pub struct CompactingMemory<M, P, C: Compactor> {
1016    inner: M,
1017    policy: P,
1018    compactor: C,
1019    state: StdMutex<HashMap<String, ConversationCompactionState<C::Artifact>>>,
1020}
1021
1022struct ConversationCompactionState<A> {
1023    /// Latest summary artifact for this conversation, if compaction has
1024    /// already happened. Cloned into the loaded history on every `load`.
1025    summary: Option<A>,
1026    /// Number of demoted messages already absorbed into `summary` within
1027    /// this process lifetime. Advanced only on compactor success.
1028    absorbed: usize,
1029    /// Reservation held while a `load` is currently awaiting the compactor for
1030    /// this conversation. Other concurrent loads observe this and short-circuit
1031    /// without re-running the compactor.
1032    in_flight: Option<InFlightReservation>,
1033}
1034
1035impl<M, P, C: Compactor> CompactingMemory<M, P, C> {
1036    /// Wrap `inner` so every load runs through `policy` and demoted messages
1037    /// are summarised by `compactor`.
1038    pub fn new(inner: M, policy: P, compactor: C) -> Self {
1039        Self {
1040            inner,
1041            policy,
1042            compactor,
1043            state: StdMutex::new(HashMap::new()),
1044        }
1045    }
1046
1047    /// Return a reference to the wrapped backend.
1048    pub fn inner(&self) -> &M {
1049        &self.inner
1050    }
1051
1052    /// Return a reference to the wrapped policy.
1053    pub fn policy(&self) -> &P {
1054        &self.policy
1055    }
1056
1057    /// Return a reference to the compactor.
1058    pub fn compactor(&self) -> &C {
1059        &self.compactor
1060    }
1061
1062    /// Consume the wrapper and return its three components.
1063    pub fn into_inner(self) -> (M, P, C) {
1064        (self.inner, self.policy, self.compactor)
1065    }
1066
1067    /// Drop the in-process compaction state for `conversation_id`.
1068    ///
1069    /// Call this when a conversation has ended to bound memory usage; the
1070    /// state map is otherwise unbounded. If the internal lock has been
1071    /// poisoned by a panic in another thread, this is a no-op.
1072    pub fn forget(&self, conversation_id: &str) {
1073        if let Ok(mut guard) = self.state.lock() {
1074            guard.remove(conversation_id);
1075        }
1076    }
1077
1078    /// Number of conversations currently tracked in the compaction state
1079    /// map. Useful for telemetry and leak detection. Returns `0` if the
1080    /// internal lock is poisoned.
1081    pub fn tracked_conversations(&self) -> usize {
1082        self.state.lock().map(|g| g.len()).unwrap_or(0)
1083    }
1084}
1085
1086impl<M, P, C> std::fmt::Debug for CompactingMemory<M, P, C>
1087where
1088    M: std::fmt::Debug,
1089    P: std::fmt::Debug,
1090    C: Compactor,
1091{
1092    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1093        f.debug_struct("CompactingMemory")
1094            .field("inner", &self.inner)
1095            .field("policy", &self.policy)
1096            .field("compactor", &"<compactor>")
1097            .finish()
1098    }
1099}
1100
1101impl<M, P, C> ConversationMemory for CompactingMemory<M, P, C>
1102where
1103    M: ConversationMemory,
1104    P: MemoryPolicy,
1105    C: Compactor,
1106{
1107    fn load<'a>(
1108        &'a self,
1109        conversation_id: &'a str,
1110    ) -> WasmBoxedFuture<'a, Result<Vec<Message>, MemoryError>> {
1111        Box::pin(async move {
1112            let messages = self.inner.load(conversation_id).await?;
1113            let (kept, demoted) = self.policy.apply_with_demoted(messages)?;
1114            let demoted_count = demoted.len();
1115
1116            // Decide-and-mark must happen under one short-lived lock so two
1117            // concurrent loads on the same conversation_id can't both
1118            // observe the same `absorbed` watermark and run the compactor
1119            // twice with the same input slice.
1120            //
1121            // Fast path: if the conversation is already tracked, mutate in
1122            // place. Only allocate a new `String` key when there is real
1123            // compaction work for a conversation we have not seen before.
1124            let plan = {
1125                let mut guard = self.state.lock().map_err(poisoned)?;
1126                if let Some(entry) = guard.get_mut(conversation_id) {
1127                    if entry.in_flight.is_some() {
1128                        // Another load is mid-compaction; return what we
1129                        // have so far. Newly-evicted messages will be
1130                        // folded in by the next load.
1131                        return Ok(splice(entry.summary.clone(), kept));
1132                    }
1133                    if demoted_count <= entry.absorbed {
1134                        // No new evictions to compact. Splice the existing
1135                        // summary (if any) and we're done.
1136                        return Ok(splice(entry.summary.clone(), kept));
1137                    }
1138                    let reservation = Arc::new(());
1139                    entry.in_flight = Some(reservation.clone());
1140                    CompactionPlan {
1141                        carry_over: entry.summary.clone(),
1142                        skip: entry.absorbed,
1143                        reservation,
1144                    }
1145                } else if demoted_count == 0 {
1146                    // First load for this conversation and nothing was
1147                    // demoted: no tracking entry needed yet.
1148                    return Ok(kept);
1149                } else {
1150                    let reservation = Arc::new(());
1151                    guard.insert(
1152                        conversation_id.to_string(),
1153                        ConversationCompactionState {
1154                            summary: None,
1155                            absorbed: 0,
1156                            in_flight: Some(reservation.clone()),
1157                        },
1158                    );
1159                    CompactionPlan {
1160                        carry_over: None,
1161                        skip: 0,
1162                        reservation,
1163                    }
1164                }
1165            };
1166
1167            // SAFETY: split_at(plan.skip) is sound because `plan.skip` was
1168            // sourced from the entry's `absorbed` watermark while we held
1169            // the lock, and we only set `absorbed = demoted_count` on
1170            // success — so `plan.skip <= demoted_count == demoted.len()`.
1171            let CompactionPlan {
1172                carry_over,
1173                skip,
1174                reservation,
1175            } = plan;
1176
1177            // Arm an RAII guard so the in-flight gate is released even if
1178            // this future is dropped mid-await (caller cancellation) or
1179            // the compactor panics. The guard is disarmed below once the
1180            // post-await state update has already cleared the flag under
1181            // the same lock acquisition that records the new watermark.
1182            let in_flight_guard =
1183                InFlightGuard::new(&self.state, conversation_id, reservation.clone());
1184
1185            let new_slice = match demoted.get(skip..) {
1186                Some(s) => s,
1187                None => {
1188                    // Drop the guard explicitly so the gate is released
1189                    // before we surface the invariant break.
1190                    drop(in_flight_guard);
1191                    return Err(MemoryError::Internal(
1192                        "compaction watermark exceeds demoted slice length".into(),
1193                    ));
1194                }
1195            };
1196
1197            let result = self
1198                .compactor
1199                .compact(conversation_id, new_slice, carry_over.as_ref())
1200                .await;
1201
1202            // Reacquire briefly to advance the watermark on success and
1203            // always clear the in-flight flag so a future load can retry.
1204            //
1205            // Only update if the entry still exists: a concurrent `clear`
1206            // (and matching `forget`) for this `conversation_id` may have
1207            // dropped the state entry while the compactor was awaiting. In
1208            // that case we must not resurrect it with stale state — the
1209            // next load on a freshly-populated backend would then start
1210            // from a non-zero watermark and skip a real compaction.
1211            let summary_for_splice = match result {
1212                Ok(artifact) => {
1213                    let mut guard = self.state.lock().map_err(poisoned)?;
1214                    if let Some(entry) = guard.get_mut(conversation_id) {
1215                        if entry
1216                            .in_flight
1217                            .as_ref()
1218                            .is_some_and(|current| Arc::ptr_eq(current, &reservation))
1219                        {
1220                            entry.in_flight = None;
1221                            entry.absorbed = demoted_count;
1222                            entry.summary = Some(artifact.clone());
1223                            Some(artifact)
1224                        } else {
1225                            None
1226                        }
1227                    } else {
1228                        // Conversation was cleared mid-compaction. Drop
1229                        // the artifact rather than reviving stale state.
1230                        None
1231                    }
1232                }
1233                Err(err) => {
1234                    let mut guard = self.state.lock().map_err(poisoned)?;
1235                    if let Some(entry) = guard.get_mut(conversation_id)
1236                        && entry
1237                            .in_flight
1238                            .as_ref()
1239                            .is_some_and(|current| Arc::ptr_eq(current, &reservation))
1240                    {
1241                        entry.in_flight = None;
1242                    }
1243                    return Err(err);
1244                }
1245            };
1246
1247            // Post-await state update completed under the lock above and
1248            // already cleared `in_flight`; disarm the RAII guard so its
1249            // `Drop` does not re-acquire the lock for a redundant clear.
1250            in_flight_guard.disarm();
1251
1252            Ok(splice(summary_for_splice, kept))
1253        })
1254    }
1255
1256    fn append<'a>(
1257        &'a self,
1258        conversation_id: &'a str,
1259        messages: Vec<Message>,
1260    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
1261        self.inner.append(conversation_id, messages)
1262    }
1263
1264    fn clear<'a>(
1265        &'a self,
1266        conversation_id: &'a str,
1267    ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
1268        Box::pin(async move {
1269            self.inner.clear(conversation_id).await?;
1270            self.forget(conversation_id);
1271            Ok(())
1272        })
1273    }
1274}
1275
1276struct CompactionPlan<A> {
1277    carry_over: Option<A>,
1278    skip: usize,
1279    reservation: InFlightReservation,
1280}
1281
1282fn splice<A>(summary: Option<A>, kept: Vec<Message>) -> Vec<Message>
1283where
1284    A: Into<Message>,
1285{
1286    match summary {
1287        Some(artifact) => {
1288            let mut out = Vec::with_capacity(kept.len() + 1);
1289            out.push(artifact.into());
1290            out.extend(kept);
1291            out
1292        }
1293        None => kept,
1294    }
1295}
1296
1297/// A zero-dependency reference [`Compactor`] that produces a textual
1298/// rollup of evicted messages without calling an LLM.
1299///
1300/// The artifact is a single [`Message::System`] whose body concatenates a
1301/// header, the previous summary (if any), and the textual content of each
1302/// newly-evicted message. It is intentionally simple: useful as a default
1303/// for tests and examples, and as a placeholder before wiring a real
1304/// summarising LLM through a custom [`Compactor`] implementation.
1305///
1306/// # Bounding the summary
1307///
1308/// By default the summary grows monotonically: every compaction pass
1309/// embeds the previous summary verbatim and appends newly-evicted lines.
1310/// Long-running conversations should call [`Self::with_max_bytes`] to
1311/// cap the rolled-up text. When the cap is exceeded, the oldest portion
1312/// of the body (after the header) is dropped at a UTF-8 boundary and
1313/// replaced with a `"[…truncated…]"` marker, preserving the most recent
1314/// context.
1315///
1316/// # Example
1317///
1318/// ```
1319/// use rig_memory::TemplateCompactor;
1320///
1321/// // Default header is "[Conversation summary so far]", unbounded.
1322/// let _compactor = TemplateCompactor::new();
1323///
1324/// // Custom header plus a 4 KiB cap for use with token-budgeted policies.
1325/// let _bounded = TemplateCompactor::with_header("Earlier context")
1326///     .with_max_bytes(4 * 1024);
1327/// ```
1328#[derive(Debug, Clone)]
1329pub struct TemplateCompactor {
1330    header: String,
1331    max_bytes: Option<usize>,
1332}
1333
1334impl TemplateCompactor {
1335    /// Create a [`TemplateCompactor`] with the default header
1336    /// `"[Conversation summary so far]"` and no size cap.
1337    pub fn new() -> Self {
1338        Self::with_header("[Conversation summary so far]")
1339    }
1340
1341    /// Create a [`TemplateCompactor`] with a custom header line and no
1342    /// size cap.
1343    pub fn with_header(header: impl Into<String>) -> Self {
1344        Self {
1345            header: header.into(),
1346            max_bytes: None,
1347        }
1348    }
1349
1350    /// Cap the rolled-up summary at `max_bytes` bytes (UTF-8). When the
1351    /// assembled body exceeds the cap, the oldest portion after the
1352    /// header is dropped at a char boundary and replaced with a
1353    /// `"[…truncated…]"` marker.
1354    ///
1355    /// `max_bytes` of `0` disables truncation (equivalent to the default
1356    /// unbounded behaviour). The header line plus the marker are always
1357    /// preserved even if they exceed the cap.
1358    pub fn with_max_bytes(mut self, max_bytes: usize) -> Self {
1359        self.max_bytes = if max_bytes == 0 {
1360            None
1361        } else {
1362            Some(max_bytes)
1363        };
1364        self
1365    }
1366}
1367
1368impl Default for TemplateCompactor {
1369    fn default() -> Self {
1370        Self::new()
1371    }
1372}
1373
1374/// Plain-text artifact produced by [`TemplateCompactor`].
1375///
1376/// Convertible into a [`Message::System`] whose body is the rolled-up
1377/// text. The system role is used because the rollup represents
1378/// out-of-band context about the prior conversation, not a turn from
1379/// any participant.
1380#[derive(Debug, Clone)]
1381pub struct TextSummary(String);
1382
1383impl TextSummary {
1384    /// Borrow the underlying summary text.
1385    pub fn as_str(&self) -> &str {
1386        &self.0
1387    }
1388
1389    /// Consume the wrapper and return the underlying `String`.
1390    pub fn into_string(self) -> String {
1391        self.0
1392    }
1393}
1394
1395impl From<TextSummary> for Message {
1396    fn from(value: TextSummary) -> Self {
1397        Message::System { content: value.0 }
1398    }
1399}
1400
1401impl Compactor for TemplateCompactor {
1402    type Artifact = TextSummary;
1403
1404    fn compact<'a>(
1405        &'a self,
1406        _conversation_id: &'a str,
1407        evicted: &'a [Message],
1408        carry_over: Option<&'a Self::Artifact>,
1409    ) -> WasmBoxedFuture<'a, Result<Self::Artifact, MemoryError>> {
1410        Box::pin(async move {
1411            let mut buf = String::new();
1412            buf.push_str(&self.header);
1413            buf.push('\n');
1414            if let Some(prev) = carry_over {
1415                buf.push_str(prev.as_str());
1416                buf.push('\n');
1417            }
1418            for msg in evicted {
1419                let line = render_message_line(msg);
1420                if !line.is_empty() {
1421                    buf.push_str(&line);
1422                    buf.push('\n');
1423                }
1424            }
1425            if let Some(cap) = self.max_bytes
1426                && buf.len() > cap
1427            {
1428                buf = truncate_summary(&buf, cap);
1429            }
1430            Ok(TextSummary(buf))
1431        })
1432    }
1433}
1434
1435/// Truncate `buf` to fit within `cap` bytes by dropping the oldest
1436/// content after the header line. Always preserves the header plus a
1437/// `"[\u{2026}truncated\u{2026}]"` marker, even if they alone exceed `cap`.
1438///
1439/// The header boundary is located by scanning `buf` for the first `\n`
1440/// rather than by trusting any caller-supplied header length, so a
1441/// header containing embedded newlines does not mis-locate the body.
1442fn truncate_summary(buf: &str, cap: usize) -> String {
1443    const MARKER: &str = "[\u{2026}truncated\u{2026}]\n";
1444    // Body starts right after the first newline in `buf`. If `buf` has
1445    // no newline at all there is no body to drop, so return as-is.
1446    let header_prefix_len = match buf.find('\n') {
1447        Some(i) => i + 1,
1448        None => return buf.to_string(),
1449    };
1450    if buf.len() <= header_prefix_len {
1451        return buf.to_string();
1452    }
1453    let preserved = header_prefix_len + MARKER.len();
1454    // Number of bytes of the body we can keep after the marker.
1455    let keep_bytes = cap.saturating_sub(preserved);
1456    let body_start = header_prefix_len;
1457    let body = match buf.get(body_start..) {
1458        Some(b) => b,
1459        None => return buf.to_string(),
1460    };
1461    // Take the suffix of `body` whose length is at most `keep_bytes`,
1462    // walking forward to a UTF-8 char boundary.
1463    let mut cut = body.len().saturating_sub(keep_bytes);
1464    while cut < body.len() && !body.is_char_boundary(cut) {
1465        cut += 1;
1466    }
1467    let suffix: &str = body.get(cut..).unwrap_or_default();
1468    let header_with_nl = match buf.get(..header_prefix_len) {
1469        Some(h) => h,
1470        None => return buf.to_string(),
1471    };
1472    let mut out = String::with_capacity(header_prefix_len + MARKER.len() + suffix.len());
1473    out.push_str(header_with_nl);
1474    out.push_str(MARKER);
1475    out.push_str(suffix);
1476    out
1477}
1478
1479/// Render a single message as a `"role: text"` line for [`TemplateCompactor`].
1480///
1481/// Non-textual content (tool calls, tool results, attachments) is rendered
1482/// as a short marker so the rollup does not silently drop them but also
1483/// does not balloon with serialized JSON.
1484fn render_message_line(msg: &Message) -> String {
1485    use rig_core::message::AssistantContent;
1486
1487    match msg {
1488        Message::System { content } => {
1489            if content.is_empty() {
1490                String::new()
1491            } else {
1492                format!("system: {content}")
1493            }
1494        }
1495        Message::User { content } => {
1496            let mut text = String::new();
1497            for c in content.iter() {
1498                match c {
1499                    UserContent::Text(t) => {
1500                        if !text.is_empty() {
1501                            text.push(' ');
1502                        }
1503                        text.push_str(&t.text);
1504                    }
1505                    UserContent::ToolResult(_) => {
1506                        if !text.is_empty() {
1507                            text.push(' ');
1508                        }
1509                        text.push_str("[tool result]");
1510                    }
1511                    _ => {
1512                        if !text.is_empty() {
1513                            text.push(' ');
1514                        }
1515                        text.push_str("[attachment]");
1516                    }
1517                }
1518            }
1519            if text.is_empty() {
1520                String::new()
1521            } else {
1522                format!("user: {text}")
1523            }
1524        }
1525        Message::Assistant { content, .. } => {
1526            let mut text = String::new();
1527            for c in content.iter() {
1528                match c {
1529                    AssistantContent::Text(t) => {
1530                        if !text.is_empty() {
1531                            text.push(' ');
1532                        }
1533                        text.push_str(&t.text);
1534                    }
1535                    AssistantContent::ToolCall(call) => {
1536                        if !text.is_empty() {
1537                            text.push(' ');
1538                        }
1539                        text.push_str(&format!("[tool call: {}]", call.function.name));
1540                    }
1541                    _ => {
1542                        if !text.is_empty() {
1543                            text.push(' ');
1544                        }
1545                        text.push_str("[reasoning]");
1546                    }
1547                }
1548            }
1549            if text.is_empty() {
1550                String::new()
1551            } else {
1552                format!("assistant: {text}")
1553            }
1554        }
1555    }
1556}
1557
1558#[cfg(test)]
1559mod tests {
1560    use super::*;
1561    use rig_core::OneOrMany;
1562    use rig_core::message::{
1563        AssistantContent, ToolCall, ToolFunction, ToolResult, ToolResultContent, UserContent,
1564    };
1565    use std::sync::Mutex;
1566
1567    fn user(text: &str) -> Message {
1568        Message::user(text)
1569    }
1570
1571    fn assistant(text: &str) -> Message {
1572        Message::assistant(text)
1573    }
1574
1575    fn tool_call_msg() -> Message {
1576        Message::Assistant {
1577            id: None,
1578            content: OneOrMany::one(AssistantContent::ToolCall(ToolCall::new(
1579                "call_1".into(),
1580                ToolFunction::new("t".into(), serde_json::json!({})),
1581            ))),
1582        }
1583    }
1584
1585    fn tool_result_msg() -> Message {
1586        Message::User {
1587            content: OneOrMany::one(UserContent::ToolResult(ToolResult {
1588                id: "call_1".into(),
1589                call_id: None,
1590                content: OneOrMany::one(ToolResultContent::text("ok")),
1591            })),
1592        }
1593    }
1594
1595    #[test]
1596    fn noop_policy_is_identity() {
1597        let msgs = vec![user("a"), assistant("b")];
1598        let out = NoopMemoryPolicy.apply(msgs).unwrap();
1599        assert_eq!(out.len(), 2);
1600    }
1601
1602    #[test]
1603    fn sliding_window_passthrough_when_under_limit() {
1604        let policy = SlidingWindowMemory::last_messages(5);
1605        let out = policy.apply(vec![user("1"), assistant("2")]).unwrap();
1606        assert_eq!(out.len(), 2);
1607    }
1608
1609    #[tokio::test]
1610    async fn sliding_window_truncates_via_filter() {
1611        let mem = InMemoryConversationMemory::new()
1612            .with_filter(SlidingWindowMemory::last_messages(2).into_filter());
1613
1614        mem.append(
1615            "c",
1616            vec![user("1"), assistant("2"), user("3"), assistant("4")],
1617        )
1618        .await
1619        .unwrap();
1620
1621        let loaded = mem.load("c").await.unwrap();
1622        assert_eq!(loaded.len(), 2);
1623    }
1624
1625    #[test]
1626    fn sliding_window_drops_leading_orphan_tool_result() {
1627        let policy = SlidingWindowMemory::last_messages(3);
1628        let out = policy
1629            .apply(vec![
1630                tool_call_msg(),
1631                tool_result_msg(),
1632                user("after"),
1633                assistant("done"),
1634            ])
1635            .unwrap();
1636
1637        assert_eq!(out.len(), 2);
1638        assert!(matches!(out.first(), Some(Message::User { content })
1639            if matches!(content.first(), UserContent::Text(_))));
1640    }
1641
1642    #[test]
1643    fn token_window_keeps_within_budget() {
1644        let msgs = vec![
1645            user("aaaa"),
1646            assistant("bbbb"),
1647            user("cccc"),
1648            assistant("dddd"),
1649        ];
1650        let policy = TokenWindowMemory::new(2, |_: &Message| 1);
1651        let out = policy.apply(msgs).unwrap();
1652        assert_eq!(out.len(), 2);
1653    }
1654
1655    #[test]
1656    fn token_window_passes_through_when_under_budget() {
1657        let msgs = vec![user("a"), assistant("b")];
1658        let policy = TokenWindowMemory::new(usize::MAX, |_: &Message| 1);
1659        let out = policy.apply(msgs).unwrap();
1660        assert_eq!(out.len(), 2);
1661    }
1662
1663    #[test]
1664    fn token_window_drops_leading_orphan_tool_result() {
1665        let policy = TokenWindowMemory::new(25, |_: &Message| 10);
1666        let out = policy
1667            .apply(vec![tool_call_msg(), tool_result_msg(), user("after")])
1668            .unwrap();
1669        assert_eq!(out.len(), 1);
1670        assert!(matches!(out.first(), Some(Message::User { content })
1671            if matches!(content.first(), UserContent::Text(_))));
1672    }
1673
1674    #[test]
1675    fn token_window_skips_message_larger_than_budget() {
1676        let policy = TokenWindowMemory::new(5, |_: &Message| 10);
1677        let out = policy.apply(vec![user("anything")]).unwrap();
1678        assert!(out.is_empty());
1679    }
1680
1681    #[test]
1682    fn heuristic_counter_charges_overhead_per_message() {
1683        let counter = HeuristicTokenCounter::default();
1684        let empty = counter.count(&user(""));
1685        assert!(
1686            empty >= 4,
1687            "default per-message overhead is at least 4 tokens"
1688        );
1689    }
1690
1691    #[test]
1692    fn heuristic_counter_is_monotonic_in_text_length() {
1693        let counter = HeuristicTokenCounter::default();
1694        let small = counter.count(&user("hi"));
1695        let big = counter.count(&user(&"x".repeat(400)));
1696        assert!(big > small);
1697    }
1698
1699    #[test]
1700    fn heuristic_counter_handles_tool_calls() {
1701        let counter = HeuristicTokenCounter::default();
1702        let cost = counter.count(&tool_call_msg());
1703        assert!(cost > 0);
1704    }
1705
1706    #[test]
1707    fn heuristic_counter_handles_system_messages() {
1708        let counter = HeuristicTokenCounter::default();
1709        let cost = counter.count(&Message::System {
1710            content: "you are helpful".into(),
1711        });
1712        assert!(cost > 0);
1713    }
1714
1715    #[test]
1716    fn heuristic_counter_clamps_invalid_bytes_per_token() {
1717        // Zero/NaN/negative ratios fall back to 1.0 instead of panicking.
1718        let counter = HeuristicTokenCounter::new(0.0, 0, 0);
1719        assert!(counter.count(&user("abcd")) >= 4);
1720        let nan = HeuristicTokenCounter::new(f32::NAN, 0, 0);
1721        assert!(nan.count(&user("abcd")) >= 4);
1722    }
1723
1724    #[test]
1725    fn heuristic_counter_drives_token_window() {
1726        let policy = TokenWindowMemory::new(100, HeuristicTokenCounter::default());
1727        let msgs = vec![user(&"a".repeat(2_000)), user("short")];
1728        let out = policy.apply(msgs).unwrap();
1729        // The huge message must be evicted; the short one retained.
1730        assert_eq!(out.len(), 1);
1731    }
1732
1733    #[test]
1734    fn arc_token_counter_can_drive_token_window() {
1735        let counter: Arc<dyn TokenCounter> = Arc::new(|_: &Message| 1);
1736        let policy = TokenWindowMemory::new(2, counter);
1737        let out = policy
1738            .apply(vec![user("a"), assistant("b"), user("c")])
1739            .unwrap();
1740
1741        assert_eq!(out.len(), 2);
1742    }
1743
1744    #[test]
1745    fn boxed_token_counter_forwards_count() {
1746        let counter: Box<dyn TokenCounter> = Box::new(|_: &Message| 7);
1747        assert_eq!(counter.count(&user("a")), 7);
1748    }
1749
1750    #[test]
1751    fn into_filter_returns_input_on_policy_error() {
1752        struct FailingPolicy;
1753        impl MemoryPolicy for FailingPolicy {
1754            fn apply(&self, _: Vec<Message>) -> Result<Vec<Message>, MemoryError> {
1755                Err(MemoryError::Policy("intentional failure".into()))
1756            }
1757        }
1758
1759        let filter = FailingPolicy.into_filter();
1760        let input = vec![user("a"), assistant("b"), user("c")];
1761        let out = filter(input.clone());
1762        assert_eq!(
1763            out.len(),
1764            input.len(),
1765            "history must be preserved on policy error"
1766        );
1767    }
1768
1769    #[tokio::test]
1770    async fn policy_memory_truncates_loaded_history() {
1771        let mem = PolicyMemory::new(
1772            InMemoryConversationMemory::new(),
1773            SlidingWindowMemory::last_messages(2),
1774        );
1775
1776        mem.append(
1777            "c",
1778            vec![user("1"), assistant("2"), user("3"), assistant("4")],
1779        )
1780        .await
1781        .unwrap();
1782
1783        let loaded = mem.load("c").await.unwrap();
1784        assert_eq!(loaded.len(), 2);
1785    }
1786
1787    #[tokio::test]
1788    async fn policy_memory_propagates_policy_errors() {
1789        struct FailingPolicy;
1790        impl MemoryPolicy for FailingPolicy {
1791            fn apply(&self, _: Vec<Message>) -> Result<Vec<Message>, MemoryError> {
1792                Err(MemoryError::Policy("intentional failure".into()))
1793            }
1794        }
1795
1796        let mem = PolicyMemory::new(InMemoryConversationMemory::new(), FailingPolicy);
1797        mem.append("c", vec![user("1"), assistant("2")])
1798            .await
1799            .unwrap();
1800
1801        let result = mem.load("c").await;
1802        assert!(matches!(result, Err(MemoryError::Policy(_))));
1803    }
1804
1805    #[tokio::test]
1806    async fn policy_memory_append_and_clear_delegate_to_inner() {
1807        let mem = PolicyMemory::new(InMemoryConversationMemory::new(), NoopMemoryPolicy);
1808        mem.append("c", vec![user("hi"), assistant("ok")])
1809            .await
1810            .unwrap();
1811        assert_eq!(mem.load("c").await.unwrap().len(), 2);
1812
1813        mem.clear("c").await.unwrap();
1814        assert!(mem.load("c").await.unwrap().is_empty());
1815    }
1816
1817    #[test]
1818    fn sliding_window_reports_demoted_prefix() {
1819        let policy = SlidingWindowMemory::last_messages(2);
1820        let (kept, demoted) = policy
1821            .apply_with_demoted(vec![
1822                user("oldest"),
1823                assistant("old"),
1824                user("recent"),
1825                assistant("latest"),
1826            ])
1827            .unwrap();
1828        assert_eq!(kept.len(), 2);
1829        assert_eq!(demoted.len(), 2);
1830    }
1831
1832    #[test]
1833    fn token_window_reports_demoted_prefix() {
1834        let policy = TokenWindowMemory::new(2, |_: &Message| 1);
1835        let (kept, demoted) = policy
1836            .apply_with_demoted(vec![user("a"), assistant("b"), user("c"), assistant("d")])
1837            .unwrap();
1838        assert_eq!(kept.len(), 2);
1839        assert_eq!(demoted.len(), 2);
1840    }
1841
1842    #[test]
1843    fn noop_policy_demotes_nothing() {
1844        let (kept, demoted) = NoopMemoryPolicy
1845            .apply_with_demoted(vec![user("a"), assistant("b")])
1846            .unwrap();
1847        assert_eq!(kept.len(), 2);
1848        assert!(demoted.is_empty());
1849    }
1850
1851    #[test]
1852    fn arc_memory_policy_preserves_demoted_metadata() {
1853        let policy: Arc<dyn MemoryPolicy> = Arc::new(SlidingWindowMemory::last_messages(1));
1854        let (kept, demoted) = policy
1855            .apply_with_demoted(vec![user("old"), assistant("new")])
1856            .unwrap();
1857
1858        assert_eq!(kept.len(), 1);
1859        assert_eq!(demoted.len(), 1);
1860    }
1861
1862    #[test]
1863    fn boxed_memory_policy_preserves_demoted_metadata() {
1864        let policy: Box<dyn MemoryPolicy> = Box::new(SlidingWindowMemory::last_messages(1));
1865        let (kept, demoted) = policy
1866            .apply_with_demoted(vec![user("old"), assistant("new")])
1867            .unwrap();
1868
1869        assert_eq!(kept.len(), 1);
1870        assert_eq!(demoted.len(), 1);
1871    }
1872
1873    #[test]
1874    fn sliding_window_demotes_orphan_tool_result_with_prefix() {
1875        // Window keeps the last 2 messages, but the leading message of that
1876        // window is an orphan tool result; it must be moved into `demoted`
1877        // so the hook can preserve it.
1878        let policy = SlidingWindowMemory::last_messages(2);
1879        let (kept, demoted) = policy
1880            .apply_with_demoted(vec![
1881                tool_call_msg(),
1882                tool_result_msg(),
1883                user("after"),
1884                assistant("done"),
1885            ])
1886            .unwrap();
1887        assert_eq!(kept.len(), 2);
1888        assert!(matches!(kept.first(), Some(Message::User { content })
1889            if matches!(content.first(), UserContent::Text(_))));
1890        assert_eq!(demoted.len(), 2);
1891    }
1892
1893    #[derive(Default)]
1894    struct CountingHook {
1895        seen: Mutex<Vec<(String, Vec<Message>)>>,
1896    }
1897
1898    impl CountingHook {
1899        fn calls(&self) -> usize {
1900            self.seen.lock().unwrap().len()
1901        }
1902        fn last_demoted_count(&self) -> usize {
1903            self.seen
1904                .lock()
1905                .unwrap()
1906                .last()
1907                .map(|(_, m)| m.len())
1908                .unwrap_or(0)
1909        }
1910    }
1911
1912    impl DemotionHook for CountingHook {
1913        fn on_demote<'a>(
1914            &'a self,
1915            conversation_id: &'a str,
1916            messages: Vec<Message>,
1917        ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
1918            Box::pin(async move {
1919                self.seen
1920                    .lock()
1921                    .unwrap()
1922                    .push((conversation_id.to_string(), messages));
1923                Ok(())
1924            })
1925        }
1926    }
1927
1928    #[tokio::test]
1929    async fn demoting_policy_memory_invokes_hook_on_truncation() {
1930        let hook = Arc::new(CountingHook::default());
1931        let mem = DemotingPolicyMemory::new(
1932            InMemoryConversationMemory::new(),
1933            SlidingWindowMemory::last_messages(2),
1934            hook.clone(),
1935        );
1936
1937        mem.append(
1938            "c",
1939            vec![user("1"), assistant("2"), user("3"), assistant("4")],
1940        )
1941        .await
1942        .unwrap();
1943
1944        let kept = mem.load("c").await.unwrap();
1945        assert_eq!(kept.len(), 2);
1946        assert_eq!(hook.calls(), 1);
1947        assert_eq!(hook.last_demoted_count(), 2);
1948    }
1949
1950    #[tokio::test]
1951    async fn demoting_policy_memory_does_not_replay_demotions() {
1952        let hook = Arc::new(CountingHook::default());
1953        let mem = DemotingPolicyMemory::new(
1954            InMemoryConversationMemory::new(),
1955            SlidingWindowMemory::last_messages(2),
1956            hook.clone(),
1957        );
1958
1959        mem.append(
1960            "c",
1961            vec![user("1"), assistant("2"), user("3"), assistant("4")],
1962        )
1963        .await
1964        .unwrap();
1965
1966        mem.load("c").await.unwrap();
1967        mem.load("c").await.unwrap();
1968        assert_eq!(hook.calls(), 1);
1969        assert_eq!(hook.last_demoted_count(), 2);
1970    }
1971
1972    #[tokio::test]
1973    async fn demoting_policy_memory_only_reports_newly_demoted_messages() {
1974        let hook = Arc::new(CountingHook::default());
1975        let mem = DemotingPolicyMemory::new(
1976            InMemoryConversationMemory::new(),
1977            SlidingWindowMemory::last_messages(2),
1978            hook.clone(),
1979        );
1980
1981        mem.append(
1982            "c",
1983            vec![user("1"), assistant("2"), user("3"), assistant("4")],
1984        )
1985        .await
1986        .unwrap();
1987        mem.load("c").await.unwrap();
1988
1989        mem.append("c", vec![user("5")]).await.unwrap();
1990        mem.load("c").await.unwrap();
1991
1992        assert_eq!(hook.calls(), 2);
1993        assert_eq!(hook.last_demoted_count(), 1);
1994    }
1995
1996    #[derive(Default)]
1997    struct FailingHook {
1998        calls: Mutex<usize>,
1999    }
2000
2001    impl DemotionHook for FailingHook {
2002        fn on_demote<'a>(
2003            &'a self,
2004            _conversation_id: &'a str,
2005            _messages: Vec<Message>,
2006        ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
2007            Box::pin(async move {
2008                *self.calls.lock().unwrap() += 1;
2009                Err(MemoryError::backend(std::io::Error::other("hook failed")))
2010            })
2011        }
2012    }
2013
2014    #[tokio::test]
2015    async fn demoting_policy_memory_does_not_advance_watermark_on_hook_failure() {
2016        let hook = Arc::new(FailingHook::default());
2017        let mem = DemotingPolicyMemory::new(
2018            InMemoryConversationMemory::new(),
2019            SlidingWindowMemory::last_messages(1),
2020            hook.clone(),
2021        );
2022        mem.append("c", vec![user("1"), assistant("2")])
2023            .await
2024            .unwrap();
2025
2026        assert!(mem.load("c").await.is_err());
2027        assert!(mem.load("c").await.is_err());
2028        assert_eq!(*hook.calls.lock().unwrap(), 2);
2029    }
2030
2031    #[tokio::test]
2032    async fn demoting_policy_memory_clear_resets_watermark() {
2033        let hook = Arc::new(CountingHook::default());
2034        let mem = DemotingPolicyMemory::new(
2035            InMemoryConversationMemory::new(),
2036            SlidingWindowMemory::last_messages(1),
2037            hook.clone(),
2038        );
2039
2040        mem.append("c", vec![user("1"), assistant("2")])
2041            .await
2042            .unwrap();
2043        mem.load("c").await.unwrap();
2044        mem.clear("c").await.unwrap();
2045        mem.append("c", vec![user("3"), assistant("4")])
2046            .await
2047            .unwrap();
2048        mem.load("c").await.unwrap();
2049
2050        assert_eq!(hook.calls(), 2);
2051        assert_eq!(hook.last_demoted_count(), 1);
2052    }
2053
2054    #[tokio::test]
2055    async fn demoting_policy_memory_skips_hook_when_nothing_evicted() {
2056        let hook = Arc::new(CountingHook::default());
2057        let mem = DemotingPolicyMemory::new(
2058            InMemoryConversationMemory::new(),
2059            SlidingWindowMemory::last_messages(10),
2060            hook.clone(),
2061        );
2062
2063        mem.append("c", vec![user("1"), assistant("2")])
2064            .await
2065            .unwrap();
2066        mem.load("c").await.unwrap();
2067        assert_eq!(hook.calls(), 0);
2068    }
2069
2070    #[tokio::test]
2071    async fn demoting_policy_memory_with_noop_hook_behaves_like_policy_memory() {
2072        let mem = DemotingPolicyMemory::new(
2073            InMemoryConversationMemory::new(),
2074            SlidingWindowMemory::last_messages(1),
2075            NoopDemotionHook,
2076        );
2077        mem.append("c", vec![user("a"), assistant("b"), user("c")])
2078            .await
2079            .unwrap();
2080        assert_eq!(mem.load("c").await.unwrap().len(), 1);
2081    }
2082
2083    /// Hook that blocks until the test releases it. Used to provoke the
2084    /// concurrent-load race against the in-flight gate.
2085    struct GatedHook {
2086        calls: Arc<std::sync::atomic::AtomicUsize>,
2087        rendezvous: Arc<tokio::sync::Notify>,
2088        release: Arc<tokio::sync::Notify>,
2089    }
2090
2091    impl DemotionHook for GatedHook {
2092        fn on_demote<'a>(
2093            &'a self,
2094            _conversation_id: &'a str,
2095            _messages: Vec<Message>,
2096        ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
2097            let calls = self.calls.clone();
2098            let rendezvous = self.rendezvous.clone();
2099            let release = self.release.clone();
2100            Box::pin(async move {
2101                calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2102                rendezvous.notify_one();
2103                release.notified().await;
2104                Ok(())
2105            })
2106        }
2107    }
2108
2109    #[tokio::test]
2110    async fn demoting_policy_memory_serialises_concurrent_loads() {
2111        use std::sync::atomic::{AtomicUsize, Ordering};
2112
2113        let calls = Arc::new(AtomicUsize::new(0));
2114        let rendezvous = Arc::new(tokio::sync::Notify::new());
2115        let release = Arc::new(tokio::sync::Notify::new());
2116        let hook = GatedHook {
2117            calls: calls.clone(),
2118            rendezvous: rendezvous.clone(),
2119            release: release.clone(),
2120        };
2121
2122        let mem = Arc::new(DemotingPolicyMemory::new(
2123            InMemoryConversationMemory::new(),
2124            SlidingWindowMemory::last_messages(1),
2125            hook,
2126        ));
2127
2128        mem.append("c", vec![user("1"), assistant("2"), user("3")])
2129            .await
2130            .unwrap();
2131
2132        let m1 = mem.clone();
2133        let first = tokio::spawn(async move { m1.load("c").await });
2134
2135        // Wait until the first load has entered the hook.
2136        rendezvous.notified().await;
2137        assert_eq!(calls.load(Ordering::SeqCst), 1);
2138
2139        // Second concurrent load on the same conversation must skip the
2140        // hook entirely (in-flight gate) and return the truncated view.
2141        let kept = mem.load("c").await.unwrap();
2142        assert_eq!(kept.len(), 1);
2143        assert_eq!(calls.load(Ordering::SeqCst), 1, "hook must not double-fire");
2144
2145        // Release the first load and confirm it completes successfully.
2146        release.notify_one();
2147        let kept_first = first.await.unwrap().unwrap();
2148        assert_eq!(kept_first.len(), 1);
2149        assert_eq!(calls.load(Ordering::SeqCst), 1);
2150
2151        // Subsequent loads observe the watermark and don't re-fire.
2152        mem.load("c").await.unwrap();
2153        assert_eq!(calls.load(Ordering::SeqCst), 1);
2154    }
2155
2156    #[tokio::test]
2157    async fn demoting_policy_memory_dropped_load_releases_in_flight_gate() {
2158        // If a `load(...)` future is dropped while awaiting the hook, the
2159        // in-flight gate must not leak: subsequent loads on the same
2160        // conversation must be able to retry demotion.
2161        use std::sync::atomic::{AtomicUsize, Ordering};
2162
2163        let calls = Arc::new(AtomicUsize::new(0));
2164        let rendezvous = Arc::new(tokio::sync::Notify::new());
2165        let release = Arc::new(tokio::sync::Notify::new());
2166        let hook = GatedHook {
2167            calls: calls.clone(),
2168            rendezvous,
2169            release: release.clone(),
2170        };
2171
2172        let mem = Arc::new(DemotingPolicyMemory::new(
2173            InMemoryConversationMemory::new(),
2174            SlidingWindowMemory::last_messages(1),
2175            hook,
2176        ));
2177
2178        mem.append("c", vec![user("1"), assistant("2"), user("3")])
2179            .await
2180            .unwrap();
2181
2182        // Kick off a load that will block inside the hook, then abort it
2183        // while awaiting — simulating a caller-side timeout or
2184        // `tokio::select!` cancellation.
2185        let mem_load = mem.clone();
2186        let handle = tokio::spawn(async move { mem_load.load("c").await });
2187        while calls.load(Ordering::SeqCst) == 0 {
2188            tokio::task::yield_now().await;
2189        }
2190        handle.abort();
2191        let _ = handle.await;
2192
2193        // The aborted future was dropped without clearing in_flight via
2194        // the success/error branches; the RAII guard's `Drop` should have
2195        // released it. A new load must therefore be able to drive a fresh
2196        // demotion rather than short-circuiting forever.
2197        let mem_load = mem.clone();
2198        let retry = tokio::spawn(async move { mem_load.load("c").await });
2199        for _ in 0..1_000 {
2200            if calls.load(Ordering::SeqCst) >= 2 {
2201                break;
2202            }
2203            tokio::task::yield_now().await;
2204        }
2205        assert_eq!(
2206            calls.load(Ordering::SeqCst),
2207            2,
2208            "retry must re-enter the hook after cancellation"
2209        );
2210
2211        release.notify_one();
2212        let kept = retry.await.unwrap().unwrap();
2213        assert_eq!(kept.len(), 1);
2214
2215        // The successful retry advances the watermark, so future loads
2216        // should not fire the hook again.
2217        mem.load("c").await.unwrap();
2218        assert_eq!(calls.load(Ordering::SeqCst), 2);
2219    }
2220
2221    #[tokio::test]
2222    async fn demoting_stale_cancelled_load_does_not_clear_new_reservation() {
2223        use std::sync::atomic::{AtomicUsize, Ordering};
2224
2225        let calls = Arc::new(AtomicUsize::new(0));
2226        let rendezvous = Arc::new(tokio::sync::Notify::new());
2227        let release = Arc::new(tokio::sync::Notify::new());
2228        let hook = GatedHook {
2229            calls: calls.clone(),
2230            rendezvous: rendezvous.clone(),
2231            release: release.clone(),
2232        };
2233
2234        let mem = Arc::new(DemotingPolicyMemory::new(
2235            InMemoryConversationMemory::new(),
2236            SlidingWindowMemory::last_messages(1),
2237            hook,
2238        ));
2239
2240        mem.append("c", vec![user("old 1"), assistant("old 2"), user("old 3")])
2241            .await
2242            .unwrap();
2243
2244        let mem_load = mem.clone();
2245        let stale = tokio::spawn(async move { mem_load.load("c").await });
2246        rendezvous.notified().await;
2247        assert_eq!(calls.load(Ordering::SeqCst), 1);
2248
2249        mem.clear("c").await.unwrap();
2250        mem.append(
2251            "c",
2252            vec![user("fresh 1"), assistant("fresh 2"), user("fresh 3")],
2253        )
2254        .await
2255        .unwrap();
2256
2257        let mem_load = mem.clone();
2258        let fresh = tokio::spawn(async move { mem_load.load("c").await });
2259        rendezvous.notified().await;
2260        assert_eq!(calls.load(Ordering::SeqCst), 2);
2261
2262        stale.abort();
2263        let _ = stale.await;
2264
2265        let mem_load = mem.clone();
2266        let mut concurrent = tokio::spawn(async move { mem_load.load("c").await });
2267        let concurrent_kept = tokio::select! {
2268            result = &mut concurrent => result.unwrap().unwrap(),
2269            _ = rendezvous.notified() => {
2270                panic!("stale guard must not clear the fresh in-flight reservation")
2271            }
2272        };
2273        assert_eq!(
2274            calls.load(Ordering::SeqCst),
2275            2,
2276            "stale guard must not clear the fresh in-flight reservation"
2277        );
2278
2279        release.notify_one();
2280        assert_eq!(fresh.await.unwrap().unwrap().len(), 1);
2281        assert_eq!(concurrent_kept.len(), 1);
2282        assert_eq!(calls.load(Ordering::SeqCst), 2);
2283    }
2284
2285    #[tokio::test]
2286    async fn demoting_stale_successful_load_does_not_clear_new_reservation() {
2287        #[derive(Default)]
2288        struct IndividuallyGatedHook {
2289            releases: Mutex<Vec<Arc<tokio::sync::Notify>>>,
2290        }
2291
2292        impl IndividuallyGatedHook {
2293            fn call_count(&self) -> usize {
2294                self.releases.lock().unwrap().len()
2295            }
2296
2297            async fn wait_for_call_count(&self, expected: usize) {
2298                while self.call_count() < expected {
2299                    tokio::task::yield_now().await;
2300                }
2301            }
2302
2303            fn release_call(&self, index: usize) {
2304                let release = self.releases.lock().unwrap()[index].clone();
2305                release.notify_one();
2306            }
2307        }
2308
2309        impl DemotionHook for IndividuallyGatedHook {
2310            fn on_demote<'a>(
2311                &'a self,
2312                _conversation_id: &'a str,
2313                _messages: Vec<Message>,
2314            ) -> WasmBoxedFuture<'a, Result<(), MemoryError>> {
2315                let release = Arc::new(tokio::sync::Notify::new());
2316                self.releases.lock().unwrap().push(release.clone());
2317                Box::pin(async move {
2318                    release.notified().await;
2319                    Ok(())
2320                })
2321            }
2322        }
2323
2324        let hook = Arc::new(IndividuallyGatedHook::default());
2325        let mem = Arc::new(DemotingPolicyMemory::new(
2326            InMemoryConversationMemory::new(),
2327            SlidingWindowMemory::last_messages(1),
2328            hook.clone(),
2329        ));
2330
2331        mem.append("c", vec![user("old 1"), assistant("old 2"), user("old 3")])
2332            .await
2333            .unwrap();
2334
2335        let mem_load = mem.clone();
2336        let stale = tokio::spawn(async move { mem_load.load("c").await });
2337        hook.wait_for_call_count(1).await;
2338
2339        mem.clear("c").await.unwrap();
2340        mem.append(
2341            "c",
2342            vec![user("fresh 1"), assistant("fresh 2"), user("fresh 3")],
2343        )
2344        .await
2345        .unwrap();
2346
2347        let mem_load = mem.clone();
2348        let fresh = tokio::spawn(async move { mem_load.load("c").await });
2349        hook.wait_for_call_count(2).await;
2350
2351        // Let the stale load finish successfully after the conversation id has
2352        // been reused. Its post-await update must not clear the fresh in-flight
2353        // reservation.
2354        hook.release_call(0);
2355        assert_eq!(stale.await.unwrap().unwrap().len(), 1);
2356        assert_eq!(hook.call_count(), 2);
2357
2358        let mem_load = mem.clone();
2359        let mut concurrent = tokio::spawn(async move { mem_load.load("c").await });
2360        let hook_wait = hook.clone();
2361        let concurrent_kept = tokio::select! {
2362            result = &mut concurrent => result.unwrap().unwrap(),
2363            _ = hook_wait.wait_for_call_count(3) => {
2364                panic!("stale successful load must not clear the fresh in-flight reservation")
2365            }
2366        };
2367        assert_eq!(
2368            hook.call_count(),
2369            2,
2370            "stale successful load must not clear the fresh in-flight reservation"
2371        );
2372
2373        hook.release_call(1);
2374        assert_eq!(fresh.await.unwrap().unwrap().len(), 1);
2375        assert_eq!(concurrent_kept.len(), 1);
2376
2377        mem.load("c").await.unwrap();
2378        assert_eq!(hook.call_count(), 2);
2379    }
2380
2381    #[tokio::test]
2382    async fn forget_drops_in_process_watermark() {
2383        let hook = Arc::new(CountingHook::default());
2384        let mem = DemotingPolicyMemory::new(
2385            InMemoryConversationMemory::new(),
2386            SlidingWindowMemory::last_messages(1),
2387            hook.clone(),
2388        );
2389
2390        mem.append("c", vec![user("1"), assistant("2")])
2391            .await
2392            .unwrap();
2393        mem.load("c").await.unwrap();
2394        assert_eq!(mem.tracked_conversations(), 1);
2395        assert_eq!(hook.calls(), 1);
2396
2397        // After forgetting, the next load on the same (still-populated)
2398        // backend re-delivers the demotion. This is the documented
2399        // contract: forget()/restart re-fire the hook, hooks must be
2400        // idempotent.
2401        mem.forget("c");
2402        assert_eq!(mem.tracked_conversations(), 0);
2403        mem.load("c").await.unwrap();
2404        assert_eq!(hook.calls(), 2);
2405    }
2406
2407    // ----------------------------------------------------------------
2408    // CompactingMemory tests
2409    // ----------------------------------------------------------------
2410
2411    #[tokio::test]
2412    async fn compacting_no_demotion_returns_kept_only() {
2413        let mem = CompactingMemory::new(
2414            InMemoryConversationMemory::new(),
2415            SlidingWindowMemory::last_messages(10),
2416            TemplateCompactor::new(),
2417        );
2418
2419        mem.append("c", vec![user("hi"), assistant("hello")])
2420            .await
2421            .unwrap();
2422        let loaded = mem.load("c").await.unwrap();
2423        assert_eq!(loaded.len(), 2);
2424        // No tracking entry needed when nothing was demoted on the first load.
2425        // (We may have inserted a default entry; what matters is that no
2426        // summary message was spliced in.)
2427        assert!(matches!(&loaded[0], Message::User { .. }));
2428    }
2429
2430    #[tokio::test]
2431    async fn compacting_splices_summary_when_demoted() {
2432        let mem = CompactingMemory::new(
2433            InMemoryConversationMemory::new(),
2434            SlidingWindowMemory::last_messages(2),
2435            TemplateCompactor::new(),
2436        );
2437
2438        mem.append(
2439            "c",
2440            vec![
2441                user("first"),
2442                assistant("second"),
2443                user("third"),
2444                assistant("fourth"),
2445            ],
2446        )
2447        .await
2448        .unwrap();
2449
2450        let loaded = mem.load("c").await.unwrap();
2451        // Expected shape: [summary, third, fourth]
2452        assert_eq!(loaded.len(), 3);
2453        let Message::System { content } = &loaded[0] else {
2454            panic!("expected summary as system message");
2455        };
2456        assert!(content.contains("[Conversation summary so far]"));
2457        assert!(content.contains("user: first"));
2458        assert!(content.contains("assistant: second"));
2459        // The kept window is intact.
2460        let Message::User { content } = &loaded[1] else {
2461            panic!("expected kept user message");
2462        };
2463        let UserContent::Text(t) = content.first_ref() else {
2464            panic!("expected text");
2465        };
2466        assert_eq!(t.text, "third");
2467    }
2468
2469    #[tokio::test]
2470    async fn compacting_rolls_summary_forward() {
2471        let mem = CompactingMemory::new(
2472            InMemoryConversationMemory::new(),
2473            SlidingWindowMemory::last_messages(2),
2474            TemplateCompactor::new(),
2475        );
2476
2477        mem.append(
2478            "c",
2479            vec![user("a"), assistant("b"), user("c"), assistant("d")],
2480        )
2481        .await
2482        .unwrap();
2483
2484        let first = mem.load("c").await.unwrap();
2485        let Message::System { content } = &first[0] else {
2486            panic!("summary missing");
2487        };
2488        let first_summary = content.clone();
2489        assert!(first_summary.contains("user: a"));
2490        assert!(first_summary.contains("assistant: b"));
2491
2492        // Append more turns; the next load should fold the previous summary
2493        // into a new one that also covers the newly-evicted prefix.
2494        mem.append("c", vec![user("e"), assistant("f")])
2495            .await
2496            .unwrap();
2497        let second = mem.load("c").await.unwrap();
2498        let Message::System { content } = &second[0] else {
2499            panic!("summary missing");
2500        };
2501        // The new summary contains the old summary text (carry_over) plus
2502        // the freshly-evicted lines.
2503        assert!(content.contains(&first_summary));
2504        assert!(content.contains("user: c"));
2505        assert!(content.contains("assistant: d"));
2506    }
2507
2508    #[tokio::test]
2509    async fn compacting_idempotent_within_process() {
2510        // Loading twice with no new evictions reuses the stored summary
2511        // and does not re-run the compactor (we observe this via the
2512        // produced text: a re-run with a non-None carry_over would double
2513        // the header line).
2514        let mem = CompactingMemory::new(
2515            InMemoryConversationMemory::new(),
2516            SlidingWindowMemory::last_messages(1),
2517            TemplateCompactor::new(),
2518        );
2519        mem.append("c", vec![user("a"), assistant("b"), user("c")])
2520            .await
2521            .unwrap();
2522
2523        let first = mem.load("c").await.unwrap();
2524        let second = mem.load("c").await.unwrap();
2525        assert_eq!(first.len(), second.len());
2526        let Message::System { content: c1 } = &first[0] else {
2527            panic!()
2528        };
2529        let Message::System { content: c2 } = &second[0] else {
2530            panic!()
2531        };
2532        assert_eq!(c1, c2);
2533    }
2534
2535    #[tokio::test]
2536    async fn compacting_clear_drops_summary() {
2537        let mem = CompactingMemory::new(
2538            InMemoryConversationMemory::new(),
2539            SlidingWindowMemory::last_messages(1),
2540            TemplateCompactor::new(),
2541        );
2542        mem.append("c", vec![user("a"), assistant("b"), user("c")])
2543            .await
2544            .unwrap();
2545        mem.load("c").await.unwrap();
2546        assert_eq!(mem.tracked_conversations(), 1);
2547
2548        mem.clear("c").await.unwrap();
2549        assert_eq!(mem.tracked_conversations(), 0);
2550        assert!(mem.load("c").await.unwrap().is_empty());
2551    }
2552
2553    // A compactor that fails the first call and succeeds afterwards, so we
2554    // can verify failure is propagated and the watermark is not advanced.
2555    #[derive(Default)]
2556    struct FlakyCompactor {
2557        calls: std::sync::atomic::AtomicUsize,
2558    }
2559
2560    impl Compactor for FlakyCompactor {
2561        type Artifact = TextSummary;
2562
2563        fn compact<'a>(
2564            &'a self,
2565            _conversation_id: &'a str,
2566            evicted: &'a [Message],
2567            _carry_over: Option<&'a Self::Artifact>,
2568        ) -> WasmBoxedFuture<'a, Result<Self::Artifact, MemoryError>> {
2569            Box::pin(async move {
2570                let n = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
2571                if n == 0 {
2572                    Err(MemoryError::Policy("flaky".into()))
2573                } else {
2574                    Ok(TextSummary(format!("compacted {} messages", evicted.len())))
2575                }
2576            })
2577        }
2578    }
2579
2580    #[tokio::test]
2581    async fn compacting_failure_does_not_advance_watermark() {
2582        let mem = CompactingMemory::new(
2583            InMemoryConversationMemory::new(),
2584            SlidingWindowMemory::last_messages(1),
2585            FlakyCompactor::default(),
2586        );
2587        mem.append("c", vec![user("a"), assistant("b"), user("c")])
2588            .await
2589            .unwrap();
2590
2591        let err = mem.load("c").await.unwrap_err();
2592        assert!(matches!(err, MemoryError::Policy(_)));
2593
2594        // Retry should succeed and produce a summary.
2595        let loaded = mem.load("c").await.unwrap();
2596        assert_eq!(loaded.len(), 2);
2597        let Message::System { content } = &loaded[0] else {
2598            panic!("expected summary")
2599        };
2600        assert!(content.contains("compacted"));
2601    }
2602
2603    // A compactor that records every invocation, including the lengths of
2604    // its `evicted` slice and whether `carry_over` was supplied.
2605    #[derive(Default)]
2606    struct CountingCompactor {
2607        log: Mutex<Vec<(usize, bool)>>,
2608    }
2609
2610    impl CountingCompactor {
2611        fn calls(&self) -> Vec<(usize, bool)> {
2612            self.log.lock().unwrap().clone()
2613        }
2614    }
2615
2616    impl Compactor for CountingCompactor {
2617        type Artifact = TextSummary;
2618
2619        fn compact<'a>(
2620            &'a self,
2621            _conversation_id: &'a str,
2622            evicted: &'a [Message],
2623            carry_over: Option<&'a Self::Artifact>,
2624        ) -> WasmBoxedFuture<'a, Result<Self::Artifact, MemoryError>> {
2625            Box::pin(async move {
2626                self.log
2627                    .lock()
2628                    .unwrap()
2629                    .push((evicted.len(), carry_over.is_some()));
2630                let prev = carry_over.map(|s| s.as_str()).unwrap_or("");
2631                Ok(TextSummary(format!("{prev}|{}", evicted.len())))
2632            })
2633        }
2634    }
2635
2636    #[tokio::test]
2637    async fn compacting_no_demotion_does_not_invoke_compactor() {
2638        let compactor = Arc::new(CountingCompactor::default());
2639        let mem = CompactingMemory::new(
2640            InMemoryConversationMemory::new(),
2641            SlidingWindowMemory::last_messages(10),
2642            compactor.clone(),
2643        );
2644
2645        mem.append("c", vec![user("a"), assistant("b")])
2646            .await
2647            .unwrap();
2648        mem.load("c").await.unwrap();
2649        mem.load("c").await.unwrap();
2650        mem.load("c").await.unwrap();
2651        assert!(compactor.calls().is_empty());
2652        // Fast path means we never installed a tracking entry either.
2653        assert_eq!(mem.tracked_conversations(), 0);
2654    }
2655
2656    #[tokio::test]
2657    async fn compacting_invokes_compactor_only_on_new_demotions() {
2658        let compactor = Arc::new(CountingCompactor::default());
2659        let mem = CompactingMemory::new(
2660            InMemoryConversationMemory::new(),
2661            SlidingWindowMemory::last_messages(2),
2662            compactor.clone(),
2663        );
2664
2665        // First eviction: 2 messages demoted.
2666        mem.append(
2667            "c",
2668            vec![user("a"), assistant("b"), user("c"), assistant("d")],
2669        )
2670        .await
2671        .unwrap();
2672        mem.load("c").await.unwrap();
2673        // Re-load: nothing new evicted; compactor must NOT run again.
2674        mem.load("c").await.unwrap();
2675        mem.load("c").await.unwrap();
2676        let calls = compactor.calls();
2677        assert_eq!(
2678            calls.len(),
2679            1,
2680            "compactor invoked more than once: {calls:?}"
2681        );
2682        assert_eq!(calls[0], (2, false));
2683
2684        // Append two more turns → another 2 demoted; compactor runs once
2685        // more, and this time `carry_over` must be present.
2686        mem.append("c", vec![user("e"), assistant("f")])
2687            .await
2688            .unwrap();
2689        mem.load("c").await.unwrap();
2690        mem.load("c").await.unwrap();
2691        let calls = compactor.calls();
2692        assert_eq!(calls.len(), 2, "expected exactly one new call: {calls:?}");
2693        // Second call only compacts the *newly* evicted prefix (2 msgs)
2694        // with the previous summary as carry-over.
2695        assert_eq!(calls[1], (2, true));
2696    }
2697
2698    #[tokio::test]
2699    async fn compacting_serialises_concurrent_loads() {
2700        // Many concurrent loads on the same conversation must produce at
2701        // most ONE compactor invocation per "epoch" of new evictions.
2702        let compactor = Arc::new(CountingCompactor::default());
2703        let mem = Arc::new(CompactingMemory::new(
2704            InMemoryConversationMemory::new(),
2705            SlidingWindowMemory::last_messages(2),
2706            compactor.clone(),
2707        ));
2708        mem.append(
2709            "c",
2710            vec![user("a"), assistant("b"), user("c"), assistant("d")],
2711        )
2712        .await
2713        .unwrap();
2714
2715        let mut handles = Vec::new();
2716        for _ in 0..32 {
2717            let mem = mem.clone();
2718            handles.push(tokio::spawn(async move {
2719                mem.load("c").await.unwrap();
2720            }));
2721        }
2722        for h in handles {
2723            h.await.unwrap();
2724        }
2725
2726        // Exactly one invocation: the first to acquire the lock runs the
2727        // compactor; the others see in_flight or the advanced watermark.
2728        let calls = compactor.calls();
2729        assert_eq!(calls.len(), 1, "expected exactly 1 call: {calls:?}");
2730    }
2731
2732    #[tokio::test]
2733    async fn compacting_clear_drops_summary_carry_over() {
2734        // After clear, the next load on a freshly-populated backend must
2735        // start compaction from scratch (carry_over=None), not roll the
2736        // old summary forward.
2737        let compactor = Arc::new(CountingCompactor::default());
2738        let mem = CompactingMemory::new(
2739            InMemoryConversationMemory::new(),
2740            SlidingWindowMemory::last_messages(1),
2741            compactor.clone(),
2742        );
2743        mem.append("c", vec![user("a"), assistant("b"), user("c")])
2744            .await
2745            .unwrap();
2746        mem.load("c").await.unwrap();
2747        assert_eq!(compactor.calls()[0], (2, false));
2748
2749        mem.clear("c").await.unwrap();
2750        assert_eq!(mem.tracked_conversations(), 0);
2751
2752        mem.append("c", vec![user("x"), assistant("y"), user("z")])
2753            .await
2754            .unwrap();
2755        mem.load("c").await.unwrap();
2756        let calls = compactor.calls();
2757        assert_eq!(calls.len(), 2);
2758        // Crucial: no carry_over after clear.
2759        assert_eq!(calls[1], (2, false));
2760    }
2761
2762    #[tokio::test]
2763    async fn compacting_forget_drops_summary() {
2764        let compactor = Arc::new(CountingCompactor::default());
2765        let mem = CompactingMemory::new(
2766            InMemoryConversationMemory::new(),
2767            SlidingWindowMemory::last_messages(1),
2768            compactor.clone(),
2769        );
2770        mem.append("c", vec![user("a"), assistant("b"), user("c")])
2771            .await
2772            .unwrap();
2773        mem.load("c").await.unwrap();
2774        assert_eq!(mem.tracked_conversations(), 1);
2775        mem.forget("c");
2776        assert_eq!(mem.tracked_conversations(), 0);
2777
2778        // Next load on the still-populated backend re-compacts from
2779        // scratch — same documented contract as DemotionHook.
2780        mem.load("c").await.unwrap();
2781        let calls = compactor.calls();
2782        assert_eq!(calls.len(), 2);
2783        assert_eq!(calls[1], (2, false));
2784    }
2785
2786    #[tokio::test]
2787    async fn compacting_arc_compactor_works() {
2788        // Arc<C> forwarding impl exists on Compactor, so CompactingMemory
2789        // must accept it.
2790        let compactor: Arc<dyn Compactor<Artifact = TextSummary>> =
2791            Arc::new(TemplateCompactor::new());
2792        let mem = CompactingMemory::new(
2793            InMemoryConversationMemory::new(),
2794            SlidingWindowMemory::last_messages(1),
2795            compactor,
2796        );
2797        mem.append("c", vec![user("a"), assistant("b"), user("c")])
2798            .await
2799            .unwrap();
2800        let loaded = mem.load("c").await.unwrap();
2801        assert_eq!(loaded.len(), 2);
2802        assert!(matches!(&loaded[0], Message::System { .. }));
2803    }
2804
2805    #[tokio::test]
2806    async fn compacting_into_inner_returns_components() {
2807        let mem = CompactingMemory::new(
2808            InMemoryConversationMemory::new(),
2809            SlidingWindowMemory::last_messages(1),
2810            TemplateCompactor::new(),
2811        );
2812        let (_inner, _policy, _compactor) = mem.into_inner();
2813    }
2814
2815    #[tokio::test]
2816    async fn compacting_isolates_conversations() {
2817        let compactor = Arc::new(CountingCompactor::default());
2818        let mem = CompactingMemory::new(
2819            InMemoryConversationMemory::new(),
2820            SlidingWindowMemory::last_messages(1),
2821            compactor.clone(),
2822        );
2823        mem.append("a", vec![user("a1"), assistant("a2"), user("a3")])
2824            .await
2825            .unwrap();
2826        mem.append("b", vec![user("b1"), assistant("b2"), user("b3")])
2827            .await
2828            .unwrap();
2829
2830        let a = mem.load("a").await.unwrap();
2831        let b = mem.load("b").await.unwrap();
2832        // Each conversation gets its own summary.
2833        assert_eq!(a.len(), 2);
2834        assert_eq!(b.len(), 2);
2835        assert_eq!(compactor.calls().len(), 2);
2836        assert_eq!(mem.tracked_conversations(), 2);
2837    }
2838
2839    #[tokio::test]
2840    async fn compacting_composes_with_token_window() {
2841        // Verify CompactingMemory is policy-agnostic: works over a
2842        // TokenWindowMemory just as well as a SlidingWindowMemory.
2843        let mem = CompactingMemory::new(
2844            InMemoryConversationMemory::new(),
2845            TokenWindowMemory::new(30, HeuristicTokenCounter::openai()),
2846            TemplateCompactor::new(),
2847        );
2848        mem.append(
2849            "c",
2850            vec![
2851                user("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"),
2852                assistant("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"),
2853                user("cccccccccccccccccccc"),
2854                assistant("d"),
2855            ],
2856        )
2857        .await
2858        .unwrap();
2859        let loaded = mem.load("c").await.unwrap();
2860        // Some prefix should have been evicted; expect a summary in front.
2861        assert!(loaded.len() >= 2);
2862        assert!(matches!(&loaded[0], Message::System { .. }));
2863    }
2864
2865    #[tokio::test]
2866    async fn template_compactor_renders_system_messages() {
2867        let compactor = TemplateCompactor::new();
2868        let evicted = vec![
2869            Message::System {
2870                content: "you are helpful".into(),
2871            },
2872            user("hi"),
2873            assistant("hello"),
2874        ];
2875        let summary = compactor.compact("c", &evicted, None).await.unwrap();
2876        let s = summary.as_str();
2877        assert!(s.contains("system: you are helpful"), "got: {s}");
2878        assert!(s.contains("user: hi"));
2879        assert!(s.contains("assistant: hello"));
2880    }
2881
2882    #[tokio::test]
2883    async fn template_compactor_renders_tool_call_marker() {
2884        let compactor = TemplateCompactor::new();
2885        let evicted = vec![tool_call_msg(), tool_result_msg()];
2886        let summary = compactor.compact("c", &evicted, None).await.unwrap();
2887        let s = summary.as_str();
2888        assert!(s.contains("[tool call: t]"), "got: {s}");
2889        assert!(s.contains("[tool result]"), "got: {s}");
2890    }
2891
2892    #[tokio::test]
2893    async fn template_compactor_carry_over_threaded() {
2894        let compactor = TemplateCompactor::new();
2895        let first = compactor
2896            .compact("c", &[user("hello")], None)
2897            .await
2898            .unwrap();
2899        assert!(!first.as_str().is_empty());
2900
2901        let second = compactor
2902            .compact("c", &[assistant("world")], Some(&first))
2903            .await
2904            .unwrap();
2905        // Carry-over text appears in the new summary.
2906        assert!(second.as_str().contains(first.as_str()));
2907        assert!(second.as_str().contains("assistant: world"));
2908    }
2909
2910    #[tokio::test]
2911    async fn template_compactor_artifact_into_message() {
2912        let s = TextSummary("rolled-up text".into());
2913        let msg: Message = s.into();
2914        let Message::System { content } = msg else {
2915            panic!("expected system message");
2916        };
2917        assert_eq!(content, "rolled-up text");
2918    }
2919
2920    #[tokio::test]
2921    async fn template_compactor_caps_summary_at_max_bytes() {
2922        let cap = 256;
2923        let compactor = TemplateCompactor::new().with_max_bytes(cap);
2924        // Build an evicted history large enough to exceed `cap` on its own.
2925        let mut evicted = Vec::new();
2926        for i in 0..50 {
2927            evicted.push(user(&format!("message number {i} with some filler")));
2928        }
2929        let summary = compactor.compact("c", &evicted, None).await.unwrap();
2930        assert!(
2931            summary.as_str().len()
2932                <= cap + "[Conversation summary so far]\n[\u{2026}truncated\u{2026}]\n".len(),
2933            "summary len {} exceeds cap {} (plus header+marker)",
2934            summary.as_str().len(),
2935            cap,
2936        );
2937        // Header is preserved.
2938        assert!(
2939            summary
2940                .as_str()
2941                .starts_with("[Conversation summary so far]\n")
2942        );
2943        // Truncation marker is present.
2944        assert!(summary.as_str().contains("[\u{2026}truncated\u{2026}]"));
2945        // Most recent line survives.
2946        assert!(summary.as_str().contains("message number 49"));
2947    }
2948
2949    #[tokio::test]
2950    async fn template_compactor_unbounded_by_default() {
2951        let compactor = TemplateCompactor::new();
2952        let mut evicted = Vec::new();
2953        for i in 0..200 {
2954            evicted.push(user(&format!("msg {i}")));
2955        }
2956        let summary = compactor.compact("c", &evicted, None).await.unwrap();
2957        // Without a cap, no truncation marker should appear.
2958        assert!(!summary.as_str().contains("[\u{2026}truncated\u{2026}]"));
2959        // Both ends are present.
2960        assert!(summary.as_str().contains("msg 0"));
2961        assert!(summary.as_str().contains("msg 199"));
2962    }
2963
2964    #[tokio::test]
2965    async fn template_compactor_with_max_bytes_zero_is_unbounded() {
2966        let compactor = TemplateCompactor::new().with_max_bytes(0);
2967        let mut evicted = Vec::new();
2968        for i in 0..200 {
2969            evicted.push(user(&format!("msg {i}")));
2970        }
2971        let summary = compactor.compact("c", &evicted, None).await.unwrap();
2972        assert!(!summary.as_str().contains("[\u{2026}truncated\u{2026}]"));
2973    }
2974
2975    #[tokio::test]
2976    async fn compacting_summary_stays_bounded_across_rolls() {
2977        // With a capped TemplateCompactor, repeated rolling must not let
2978        // the summary grow without bound.
2979        let cap = 512;
2980        let mem = CompactingMemory::new(
2981            InMemoryConversationMemory::new(),
2982            SlidingWindowMemory::last_messages(2),
2983            TemplateCompactor::new().with_max_bytes(cap),
2984        );
2985        mem.append("c", vec![user("seed-a"), assistant("seed-b")])
2986            .await
2987            .unwrap();
2988        for i in 0..30 {
2989            mem.append(
2990                "c",
2991                vec![
2992                    user(&format!("user line {i} ----- padding padding padding")),
2993                    assistant(&format!("assistant line {i} ----- padding padding")),
2994                ],
2995            )
2996            .await
2997            .unwrap();
2998            mem.load("c").await.unwrap();
2999        }
3000        let loaded = mem.load("c").await.unwrap();
3001        let Message::System { content } = &loaded[0] else {
3002            panic!("expected summary");
3003        };
3004        // Allow some slack for header + marker overhead.
3005        let slack = "[Conversation summary so far]\n[\u{2026}truncated\u{2026}]\n".len();
3006        assert!(
3007            content.len() <= cap + slack,
3008            "summary grew to {} bytes (cap {}, slack {})",
3009            content.len(),
3010            cap,
3011            slack,
3012        );
3013    }
3014
3015    #[tokio::test]
3016    async fn compacting_concurrent_with_clear_does_not_resurrect_state() {
3017        // A clear that lands while compaction is in flight must not be
3018        // overwritten by the post-await state update.
3019        use std::sync::atomic::{AtomicBool, Ordering};
3020
3021        struct GatedCompactor {
3022            release: tokio::sync::Notify,
3023            entered: AtomicBool,
3024        }
3025
3026        impl Compactor for GatedCompactor {
3027            type Artifact = TextSummary;
3028
3029            fn compact<'a>(
3030                &'a self,
3031                _conversation_id: &'a str,
3032                _evicted: &'a [Message],
3033                _carry_over: Option<&'a Self::Artifact>,
3034            ) -> WasmBoxedFuture<'a, Result<Self::Artifact, MemoryError>> {
3035                Box::pin(async move {
3036                    self.entered.store(true, Ordering::SeqCst);
3037                    self.release.notified().await;
3038                    Ok(TextSummary("late summary".into()))
3039                })
3040            }
3041        }
3042
3043        let compactor = Arc::new(GatedCompactor {
3044            release: tokio::sync::Notify::new(),
3045            entered: AtomicBool::new(false),
3046        });
3047        let mem = Arc::new(CompactingMemory::new(
3048            InMemoryConversationMemory::new(),
3049            SlidingWindowMemory::last_messages(1),
3050            compactor.clone(),
3051        ));
3052        mem.append("c", vec![user("a"), assistant("b"), user("c")])
3053            .await
3054            .unwrap();
3055
3056        // Kick off a load that will block inside the compactor.
3057        let mem_load = mem.clone();
3058        let load_handle = tokio::spawn(async move { mem_load.load("c").await });
3059
3060        // Wait for the compactor to have entered.
3061        while !compactor.entered.load(Ordering::SeqCst) {
3062            tokio::task::yield_now().await;
3063        }
3064
3065        // Clear while the compaction is in flight.
3066        mem.clear("c").await.unwrap();
3067
3068        // Release the compactor; it should complete and *not* resurrect
3069        // the cleared state.
3070        compactor.release.notify_one();
3071        let _ = load_handle.await.unwrap();
3072
3073        assert_eq!(mem.tracked_conversations(), 0);
3074        // A subsequent load on the empty backend returns nothing.
3075        assert!(mem.load("c").await.unwrap().is_empty());
3076    }
3077
3078    #[tokio::test]
3079    async fn compacting_dropped_load_releases_in_flight_gate() {
3080        // If a `load(...)` future is dropped while awaiting the
3081        // compactor, the in-flight gate must not leak: subsequent loads
3082        // on the same conversation must be able to retry compaction.
3083        use std::sync::atomic::{AtomicUsize, Ordering};
3084
3085        struct GatedCompactor {
3086            release: tokio::sync::Notify,
3087            entered: AtomicUsize,
3088        }
3089
3090        impl Compactor for GatedCompactor {
3091            type Artifact = TextSummary;
3092
3093            fn compact<'a>(
3094                &'a self,
3095                _conversation_id: &'a str,
3096                _evicted: &'a [Message],
3097                _carry_over: Option<&'a Self::Artifact>,
3098            ) -> WasmBoxedFuture<'a, Result<Self::Artifact, MemoryError>> {
3099                Box::pin(async move {
3100                    self.entered.fetch_add(1, Ordering::SeqCst);
3101                    self.release.notified().await;
3102                    Ok(TextSummary("ran".into()))
3103                })
3104            }
3105        }
3106
3107        let compactor = Arc::new(GatedCompactor {
3108            release: tokio::sync::Notify::new(),
3109            entered: AtomicUsize::new(0),
3110        });
3111        let mem = Arc::new(CompactingMemory::new(
3112            InMemoryConversationMemory::new(),
3113            SlidingWindowMemory::last_messages(1),
3114            compactor.clone(),
3115        ));
3116        mem.append("c", vec![user("a"), assistant("b"), user("c")])
3117            .await
3118            .unwrap();
3119
3120        // Kick off a load that will block inside the compactor, then
3121        // abort it while awaiting — simulating a caller-side timeout
3122        // or `tokio::select!` cancellation.
3123        let mem_load = mem.clone();
3124        let handle = tokio::spawn(async move { mem_load.load("c").await });
3125        while compactor.entered.load(Ordering::SeqCst) == 0 {
3126            tokio::task::yield_now().await;
3127        }
3128        handle.abort();
3129        let _ = handle.await;
3130
3131        // The aborted future was dropped without clearing in_flight via
3132        // the success/error branches; the RAII guard's `Drop` should
3133        // have released it. A new load must therefore be able to drive
3134        // a fresh compaction rather than short-circuiting forever.
3135        let mem_load = mem.clone();
3136        let retry = tokio::spawn(async move { mem_load.load("c").await });
3137        // Wait for the compactor to be entered a second time. If the
3138        // gate had leaked, this would never happen — the load would
3139        // short-circuit on `in_flight = true` and return immediately.
3140        while compactor.entered.load(Ordering::SeqCst) < 2 {
3141            tokio::task::yield_now().await;
3142        }
3143        compactor.release.notify_one();
3144        let loaded = retry.await.unwrap().unwrap();
3145        assert_eq!(loaded.len(), 2);
3146        let Message::System { content } = &loaded[0] else {
3147            panic!("expected summary")
3148        };
3149        assert_eq!(content, "ran");
3150    }
3151
3152    #[tokio::test]
3153    async fn compacting_stale_cancelled_load_does_not_clear_new_reservation() {
3154        use std::sync::atomic::{AtomicUsize, Ordering};
3155
3156        struct GatedCompactor {
3157            release: tokio::sync::Notify,
3158            rendezvous: tokio::sync::Notify,
3159            entered: AtomicUsize,
3160        }
3161
3162        impl Compactor for GatedCompactor {
3163            type Artifact = TextSummary;
3164
3165            fn compact<'a>(
3166                &'a self,
3167                _conversation_id: &'a str,
3168                _evicted: &'a [Message],
3169                _carry_over: Option<&'a Self::Artifact>,
3170            ) -> WasmBoxedFuture<'a, Result<Self::Artifact, MemoryError>> {
3171                Box::pin(async move {
3172                    self.entered.fetch_add(1, Ordering::SeqCst);
3173                    self.rendezvous.notify_one();
3174                    self.release.notified().await;
3175                    Ok(TextSummary("ran".into()))
3176                })
3177            }
3178        }
3179
3180        let compactor = Arc::new(GatedCompactor {
3181            release: tokio::sync::Notify::new(),
3182            rendezvous: tokio::sync::Notify::new(),
3183            entered: AtomicUsize::new(0),
3184        });
3185        let mem = Arc::new(CompactingMemory::new(
3186            InMemoryConversationMemory::new(),
3187            SlidingWindowMemory::last_messages(1),
3188            compactor.clone(),
3189        ));
3190
3191        mem.append("c", vec![user("old 1"), assistant("old 2"), user("old 3")])
3192            .await
3193            .unwrap();
3194
3195        let mem_load = mem.clone();
3196        let stale = tokio::spawn(async move { mem_load.load("c").await });
3197        compactor.rendezvous.notified().await;
3198        assert_eq!(compactor.entered.load(Ordering::SeqCst), 1);
3199
3200        mem.clear("c").await.unwrap();
3201        mem.append(
3202            "c",
3203            vec![user("fresh 1"), assistant("fresh 2"), user("fresh 3")],
3204        )
3205        .await
3206        .unwrap();
3207
3208        let mem_load = mem.clone();
3209        let fresh = tokio::spawn(async move { mem_load.load("c").await });
3210        compactor.rendezvous.notified().await;
3211        assert_eq!(compactor.entered.load(Ordering::SeqCst), 2);
3212
3213        stale.abort();
3214        let _ = stale.await;
3215
3216        let mem_load = mem.clone();
3217        let mut concurrent = tokio::spawn(async move { mem_load.load("c").await });
3218        let concurrent_kept = tokio::select! {
3219            result = &mut concurrent => result.unwrap().unwrap(),
3220            _ = compactor.rendezvous.notified() => {
3221                panic!("stale guard must not clear the fresh in-flight reservation")
3222            }
3223        };
3224        assert_eq!(
3225            compactor.entered.load(Ordering::SeqCst),
3226            2,
3227            "stale guard must not clear the fresh in-flight reservation"
3228        );
3229
3230        compactor.release.notify_one();
3231        assert_eq!(fresh.await.unwrap().unwrap().len(), 2);
3232        assert_eq!(concurrent_kept.len(), 1);
3233        assert_eq!(compactor.entered.load(Ordering::SeqCst), 2);
3234    }
3235
3236    #[tokio::test]
3237    async fn template_compactor_caps_summary_with_multiline_header() {
3238        // A header containing embedded newlines must not break the
3239        // truncation boundary calculation. The first newline in the
3240        // assembled buffer marks the header/body split, regardless of
3241        // how the caller chose to format the header.
3242        let cap = 256;
3243        let compactor = TemplateCompactor::with_header("line one\nline two").with_max_bytes(cap);
3244        let mut evicted = Vec::new();
3245        for i in 0..50 {
3246            evicted.push(user(&format!("message number {i} with some filler")));
3247        }
3248        let summary = compactor.compact("c", &evicted, None).await.unwrap();
3249        let text = summary.as_str();
3250
3251        // The first line of the header is preserved as the header line.
3252        assert!(text.starts_with("line one\n"));
3253        // Truncation marker is present and the most recent line survives.
3254        assert!(text.contains("[\u{2026}truncated\u{2026}]"));
3255        assert!(text.contains("message number 49"));
3256        // Cap is honoured up to the header+marker overhead.
3257        let overhead = "line one\n".len() + "[\u{2026}truncated\u{2026}]\n".len();
3258        assert!(
3259            text.len() <= cap + overhead,
3260            "summary len {} exceeds cap {} plus overhead {}",
3261            text.len(),
3262            cap,
3263            overhead,
3264        );
3265    }
3266}