1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
//! Async variant of context transformation.
//!
//! [`AsyncContextTransformer`] supports async operations like fetching summaries
//! from an LLM or database before compacting context. It complements the
//! synchronous [`ContextTransformer`](crate::ContextTransformer) used in the
//! hot loop.
use std::future::Future;
use std::pin::Pin;
use crate::context::CompactionReport;
use crate::types::AgentMessage;
/// A boxed future returned by an [`AsyncContextTransformer`].
pub type AsyncTransformFuture<'a> =
Pin<Box<dyn Future<Output = Option<CompactionReport>> + Send + 'a>>;
/// Async context transformer for operations that require I/O (summary fetching,
/// RAG retrieval, database lookups) before transforming the message context.
///
/// Unlike [`ContextTransformer`](crate::ContextTransformer), this trait's
/// `transform` method is async, making it suitable for pre-turn preparation
/// that involves network calls or other async work.
///
/// # Usage Pattern
///
/// The async transformer runs *before* the synchronous `ContextTransformer` in
/// the turn pipeline. It can inject summary messages, fetch relevant context
/// from a vector store, or perform any async preparation.
pub trait AsyncContextTransformer: Send + Sync {
/// Transform the context messages asynchronously.
///
/// Called before each LLM turn. The `overflow` flag is true when the
/// previous turn exceeded the context window.
///
/// Returns `Some(CompactionReport)` if messages were modified, `None` otherwise.
fn transform<'a>(
&'a self,
messages: &'a mut Vec<AgentMessage>,
overflow: bool,
) -> AsyncTransformFuture<'a>;
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{ContentBlock, LlmMessage, UserMessage};
fn text_message(text: &str) -> AgentMessage {
AgentMessage::Llm(LlmMessage::User(UserMessage {
content: vec![ContentBlock::Text {
text: text.to_owned(),
}],
timestamp: 0,
cache_hint: None,
}))
}
#[tokio::test]
async fn async_transformer_struct_impl() {
struct OverflowTruncator;
impl AsyncContextTransformer for OverflowTruncator {
fn transform<'a>(
&'a self,
messages: &'a mut Vec<AgentMessage>,
overflow: bool,
) -> AsyncTransformFuture<'a> {
Box::pin(async move {
if overflow && messages.len() > 2 {
let before = messages.len();
messages.truncate(2);
Some(CompactionReport {
dropped_count: before - 2,
tokens_before: 0,
tokens_after: 0,
overflow: true,
dropped_messages: Vec::new(),
})
} else {
None
}
})
}
}
let transformer = OverflowTruncator;
// No overflow — no change
let mut messages = vec![text_message("a"), text_message("b"), text_message("c")];
let report = transformer.transform(&mut messages, false).await;
assert!(report.is_none());
assert_eq!(messages.len(), 3);
// Overflow — truncate
let report = transformer.transform(&mut messages, true).await;
assert!(report.is_some());
let report = report.unwrap();
assert_eq!(report.dropped_count, 1);
assert!(report.overflow);
assert_eq!(messages.len(), 2);
}
#[tokio::test]
async fn async_transformer_trait_object() {
struct SummaryInjector;
impl AsyncContextTransformer for SummaryInjector {
fn transform<'a>(
&'a self,
messages: &'a mut Vec<AgentMessage>,
_overflow: bool,
) -> AsyncTransformFuture<'a> {
Box::pin(async move {
// Simulate injecting a summary at the start
messages.insert(0, text_message("[summary of prior context]"));
None // not compaction, just injection
})
}
}
let transformer: Box<dyn AsyncContextTransformer> = Box::new(SummaryInjector);
let mut messages = vec![text_message("hello")];
transformer.transform(&mut messages, false).await;
assert_eq!(messages.len(), 2);
if let AgentMessage::Llm(LlmMessage::User(u)) = &messages[0] {
assert_eq!(
ContentBlock::extract_text(&u.content),
"[summary of prior context]"
);
} else {
panic!("expected user message");
}
}
}