1use std::sync::Arc;
2use std::time::{Duration, Instant};
3
4use af_agent::ChatModel;
5use af_agent_session::Event;
6use af_llm::{ChatMessage, CompletionRequest, Role};
7use async_trait::async_trait;
8
9use crate::{
10 AgentRuntime, CancellationToken, CompactionResult, Compactor, EventWriter, RuntimeError,
11};
12
13const COMPACTION_PROMPT: &str = "Create a compact conversation checkpoint. Preserve: current objective; confirmed decisions and constraints; concrete facts with citations; completed tool outcomes; unresolved questions; pending actions; and the exact next step. Do not invent facts. Use short labeled sections.";
14const COMPACTION_MAX_TOKENS: u32 = 2_048;
15
16pub struct ModelCompactor {
17 model: Arc<dyn ChatModel>,
18}
19
20impl ModelCompactor {
21 pub fn new(model: Arc<dyn ChatModel>) -> Self {
22 Self { model }
23 }
24}
25
26#[async_trait]
27impl Compactor for ModelCompactor {
28 fn name(&self) -> &str {
29 "agentfactory.model-compactor"
30 }
31
32 async fn summarize(
33 &self,
34 model: &str,
35 messages: &[ChatMessage],
36 operation_id: &str,
37 cancellation: CancellationToken,
38 deadline: Instant,
39 ) -> Result<CompactionResult, RuntimeError> {
40 let mut request = CompletionRequest::new(
41 model,
42 std::iter::once(ChatMessage::system(COMPACTION_PROMPT))
43 .chain(messages.iter().cloned())
44 .collect(),
45 )
46 .temperature(0.1)
47 .max_tokens(COMPACTION_MAX_TOKENS);
48 request.provider_attempt_id = Some(operation_id.into());
49 let (delta_tx, _delta_rx) = tokio::sync::mpsc::unbounded_channel();
50 let completion = self.model.complete_streaming(&request, delta_tx);
51 tokio::pin!(completion);
52 let timeout = tokio::time::sleep_until(deadline.into());
53 tokio::pin!(timeout);
54 let mut cancellation_poll = tokio::time::interval(Duration::from_millis(25));
55 let response = loop {
56 tokio::select! {
57 response = &mut completion => break response.map_err(|error| RuntimeError::Model(error.to_string()))?,
58 _ = &mut timeout => return Err(RuntimeError::Model("compaction deadline exceeded".into())),
59 _ = cancellation_poll.tick() => if cancellation.is_cancelled() {
60 return Err(RuntimeError::Cancelled);
61 },
62 }
63 };
64 let summary = response
65 .first_content()
66 .filter(|value| !value.trim().is_empty())
67 .ok_or_else(|| RuntimeError::Model("compactor returned no summary".into()))?
68 .trim()
69 .to_string();
70 let (prompt_tokens, completion_tokens) = response.usage.map_or_else(
71 || {
72 (
73 request
74 .messages
75 .iter()
76 .filter_map(|message| message.content.as_deref())
77 .map(|content| content.chars().count().div_ceil(4) as u64)
78 .sum(),
79 summary.chars().count().div_ceil(4).max(1) as u64,
80 )
81 },
82 |usage| {
83 (
84 u64::from(usage.prompt_tokens),
85 u64::from(usage.completion_tokens),
86 )
87 },
88 );
89 Ok(CompactionResult {
90 summary,
91 prompt_tokens,
92 completion_tokens,
93 })
94 }
95}
96
97impl AgentRuntime {
98 pub(super) async fn compact_if_needed(
99 &self,
100 writer: &dyn EventWriter,
101 run_id: &str,
102 step: u32,
103 transcript: &mut Vec<ChatMessage>,
104 context: &[ChatMessage],
105 cancellation: CancellationToken,
106 ) -> Result<(u64, u64), RuntimeError> {
107 let count = |transcript: &[ChatMessage]| {
108 self.meter.count(
109 &self.model_name,
110 &std::iter::once(ChatMessage::system(self.prompts.render()))
111 .chain(transcript.iter().cloned())
112 .chain(context.iter().cloned())
113 .collect::<Vec<_>>(),
114 )
115 };
116 if count(transcript) <= self.limits.max_tokens {
117 return Ok((0, 0));
118 }
119 let mut pruned = Vec::new();
120 for message in transcript
121 .iter_mut()
122 .filter(|message| message.role == Role::Tool)
123 {
124 let Some(content) = &message.content else {
125 continue;
126 };
127 if content.chars().count() <= 2_000 {
128 continue;
129 }
130 if let Some(call_id) = &message.tool_call_id {
131 pruned.push(call_id.clone());
132 }
133 message.content = Some(format!(
134 "{}… [tool result pruned]",
135 content.chars().take(512).collect::<String>()
136 ));
137 }
138 if !pruned.is_empty() {
139 writer
140 .append(vec![Event::ToolResultsPruned {
141 run_id: run_id.into(),
142 call_ids: pruned,
143 }])
144 .await?;
145 if count(transcript) <= self.limits.max_tokens {
146 return Ok((0, 0));
147 }
148 }
149 let compactor = &self.compactor;
150 let compaction_id = format!("compaction:{step}");
151 let operation_id = format!("{compaction_id}:attempt:1");
152 let provider_attempt_id = format!("{run_id}:{operation_id}");
153 let source_through_seq = writer
154 .load_after(0)
155 .await?
156 .last()
157 .map_or(0, |event| event.seq);
158 let compactable = transcript
159 .iter()
160 .cloned()
161 .chain(context.iter().cloned())
162 .collect::<Vec<_>>();
163 let prepared = writer
164 .append(vec![
165 Event::CompactionStarted {
166 run_id: run_id.into(),
167 compaction_id: compaction_id.clone(),
168 source_through_seq,
169 },
170 Event::ModelRequestPrepared {
171 run_id: run_id.into(),
172 step,
173 attempt: 1,
174 provider_attempt_id: provider_attempt_id.clone(),
175 operation_id: operation_id.clone(),
176 reserved_prompt_tokens: self.meter.count(
177 &self.model_name,
178 &std::iter::once(ChatMessage::system(COMPACTION_PROMPT))
179 .chain(compactable.iter().cloned())
180 .collect::<Vec<_>>(),
181 ),
182 reserved_completion_tokens: u64::from(COMPACTION_MAX_TOKENS),
183 request: serde_json::to_value(&compactable)
184 .map_err(|error| RuntimeError::Invariant(error.to_string()))?,
185 prompt_sections: serde_json::Value::Null,
186 },
187 ])
188 .await?;
189 let prepared_through_seq = prepared
190 .last()
191 .map_or(source_through_seq, |event| event.seq);
192 let extension_cancellation = cancellation.child();
193 let deadline = std::time::Instant::now() + self.limits.provider_deadline;
194 let result = match crate::extension::await_extension(
195 "context compactor",
196 extension_cancellation.clone(),
197 deadline,
198 compactor.summarize(
199 &self.model_name,
200 &compactable,
201 &provider_attempt_id,
202 extension_cancellation,
203 deadline,
204 ),
205 )
206 .await
207 {
208 Ok(Ok(result)) => result,
209 Ok(Err(error)) | Err(error) => {
210 writer
211 .append(vec![Event::CompactionFinished {
212 run_id: run_id.into(),
213 compaction_id,
214 status: "failed".into(),
215 error: Some(error.to_string()),
216 }])
217 .await?;
218 return Err(error);
219 }
220 };
221 let current_seq = writer
222 .load_after(0)
223 .await?
224 .last()
225 .map_or(0, |event| event.seq);
226 if current_seq != prepared_through_seq {
227 writer
228 .append(vec![
229 Event::UsageRecorded {
230 run_id: run_id.into(),
231 operation_id,
232 prompt_tokens: result.prompt_tokens,
233 completion_tokens: result.completion_tokens,
234 cost_units: 0,
235 },
236 Event::CompactionFinished {
237 run_id: run_id.into(),
238 compaction_id,
239 status: "failed".into(),
240 error: Some("surface_changed".into()),
241 },
242 ])
243 .await?;
244 return Err(RuntimeError::CompactionConflict);
245 }
246 transcript.clear();
247 transcript.push(ChatMessage::system(format!(
248 "Conversation summary:\n{}",
249 result.summary
250 )));
251 let mut events = vec![
252 Event::UsageRecorded {
253 run_id: run_id.into(),
254 operation_id,
255 prompt_tokens: result.prompt_tokens,
256 completion_tokens: result.completion_tokens,
257 cost_units: 0,
258 },
259 Event::SummaryReplaced {
260 run_id: run_id.into(),
261 through_seq: source_through_seq,
262 summary: result.summary,
263 compactor: compactor.name().into(),
264 model: self.model_name.clone(),
265 },
266 Event::CompactionFinished {
267 run_id: run_id.into(),
268 compaction_id,
269 status: "completed".into(),
270 error: None,
271 },
272 ];
273 if result.prompt_tokens == 0 && result.completion_tokens == 0 {
274 events.remove(0);
275 }
276 writer.append(events).await?;
277 Ok((result.prompt_tokens, result.completion_tokens))
278 }
279}