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)]
11use std::{
43 collections::HashMap,
44 sync::{Arc, Mutex as StdMutex},
45};
46
47pub 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
58pub trait MemoryPolicy: WasmCompatSend + WasmCompatSync {
64 fn apply(&self, messages: Vec<Message>) -> Result<Vec<Message>, MemoryError>;
67
68 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
121pub trait IntoFilter: MemoryPolicy + Sized + 'static {
128 #[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 #[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#[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#[derive(Debug, Clone, Copy)]
190pub struct SlidingWindowMemory {
191 max_messages: usize,
192}
193
194impl SlidingWindowMemory {
195 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 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
232pub trait TokenCounter: WasmCompatSend + WasmCompatSync {
238 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#[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 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 pub fn openai() -> Self {
340 Self::new(4.0, 4, 256)
341 }
342
343 pub fn anthropic() -> Self {
345 Self::new(3.5, 4, 256)
346 }
347
348 pub fn gemini() -> Self {
350 Self::new(4.0, 4, 256)
351 }
352
353 fn bytes_to_tokens(&self, bytes: usize) -> usize {
354 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 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
425pub struct TokenWindowMemory {
434 max_tokens: usize,
435 counter: Arc<dyn TokenCounter>,
436}
437
438impl TokenWindowMemory {
439 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#[derive(Debug, Clone, Copy)]
514pub struct PolicyMemory<M, P> {
515 inner: M,
516 policy: P,
517}
518
519impl<M, P> PolicyMemory<M, P> {
520 pub fn new(inner: M, policy: P) -> Self {
522 Self { inner, policy }
523 }
524
525 pub fn inner(&self) -> &M {
527 &self.inner
528 }
529
530 pub fn policy(&self) -> &P {
532 &self.policy
533 }
534
535 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
572pub 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 delivered: usize,
637 in_flight: Option<InFlightReservation>,
641}
642
643impl<M, P, H> DemotingPolicyMemory<M, P, H> {
644 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 pub fn inner(&self) -> &M {
657 &self.inner
658 }
659
660 pub fn policy(&self) -> &P {
662 &self.policy
663 }
664
665 pub fn hook(&self) -> &H {
667 &self.hook
668 }
669
670 pub fn into_inner(self) -> (M, P, H) {
672 (self.inner, self.policy, self.hook)
673 }
674
675 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 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 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 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 (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 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 {
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
837struct 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 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
891struct 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 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
946pub 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 summary: Option<A>,
1026 absorbed: usize,
1029 in_flight: Option<InFlightReservation>,
1033}
1034
1035impl<M, P, C: Compactor> CompactingMemory<M, P, C> {
1036 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 pub fn inner(&self) -> &M {
1049 &self.inner
1050 }
1051
1052 pub fn policy(&self) -> &P {
1054 &self.policy
1055 }
1056
1057 pub fn compactor(&self) -> &C {
1059 &self.compactor
1060 }
1061
1062 pub fn into_inner(self) -> (M, P, C) {
1064 (self.inner, self.policy, self.compactor)
1065 }
1066
1067 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 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 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 return Ok(splice(entry.summary.clone(), kept));
1132 }
1133 if demoted_count <= entry.absorbed {
1134 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 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 let CompactionPlan {
1172 carry_over,
1173 skip,
1174 reservation,
1175 } = plan;
1176
1177 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(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 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 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 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#[derive(Debug, Clone)]
1329pub struct TemplateCompactor {
1330 header: String,
1331 max_bytes: Option<usize>,
1332}
1333
1334impl TemplateCompactor {
1335 pub fn new() -> Self {
1338 Self::with_header("[Conversation summary so far]")
1339 }
1340
1341 pub fn with_header(header: impl Into<String>) -> Self {
1344 Self {
1345 header: header.into(),
1346 max_bytes: None,
1347 }
1348 }
1349
1350 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#[derive(Debug, Clone)]
1381pub struct TextSummary(String);
1382
1383impl TextSummary {
1384 pub fn as_str(&self) -> &str {
1386 &self.0
1387 }
1388
1389 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
1435fn truncate_summary(buf: &str, cap: usize) -> String {
1443 const MARKER: &str = "[\u{2026}truncated\u{2026}]\n";
1444 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 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 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
1479fn 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 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 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 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 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 rendezvous.notified().await;
2137 assert_eq!(calls.load(Ordering::SeqCst), 1);
2138
2139 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.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 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 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 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 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 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 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 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 #[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 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 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 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 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 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 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 #[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 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 #[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 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 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 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 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 assert_eq!(calls[1], (2, true));
2696 }
2697
2698 #[tokio::test]
2699 async fn compacting_serialises_concurrent_loads() {
2700 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 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 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 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 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 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 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 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 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 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 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 assert!(
2939 summary
2940 .as_str()
2941 .starts_with("[Conversation summary so far]\n")
2942 );
2943 assert!(summary.as_str().contains("[\u{2026}truncated\u{2026}]"));
2945 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 assert!(!summary.as_str().contains("[\u{2026}truncated\u{2026}]"));
2959 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 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 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 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 let mem_load = mem.clone();
3058 let load_handle = tokio::spawn(async move { mem_load.load("c").await });
3059
3060 while !compactor.entered.load(Ordering::SeqCst) {
3062 tokio::task::yield_now().await;
3063 }
3064
3065 mem.clear("c").await.unwrap();
3067
3068 compactor.release.notify_one();
3071 let _ = load_handle.await.unwrap();
3072
3073 assert_eq!(mem.tracked_conversations(), 0);
3074 assert!(mem.load("c").await.unwrap().is_empty());
3076 }
3077
3078 #[tokio::test]
3079 async fn compacting_dropped_load_releases_in_flight_gate() {
3080 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 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 let mem_load = mem.clone();
3136 let retry = tokio::spawn(async move { mem_load.load("c").await });
3137 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 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 assert!(text.starts_with("line one\n"));
3253 assert!(text.contains("[\u{2026}truncated\u{2026}]"));
3255 assert!(text.contains("message number 49"));
3256 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}