swink_agent/
async_context_transformer.rs

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

swink_agent/
async_context_transformer.rs
