use std::sync::Arc;
use std::time::Instant;
use af_agent::ChatModel;
use af_agent_session::Event;
use af_context::RunId;
use af_llm::{ChatMessage, CompletionRequest, Role};
use async_trait::async_trait;
use crate::{
AgentRuntime, CancellationToken, CompactionResult, Compactor, EventWriter, RuntimeError,
};
const 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.";
const COMPACTION_MAX_TOKENS: u32 = 2_048;
pub struct ModelCompactor {
model: Arc<dyn ChatModel>,
}
impl ModelCompactor {
pub fn new(model: Arc<dyn ChatModel>) -> Self {
Self { model }
}
}
#[async_trait]
impl Compactor for ModelCompactor {
fn name(&self) -> &str {
"agentfactory.model-compactor"
}
async fn summarize(
&self,
model: &str,
messages: &[ChatMessage],
operation_id: &str,
cancellation: CancellationToken,
deadline: Instant,
) -> Result<CompactionResult, RuntimeError> {
let mut request = CompletionRequest::new(
model,
std::iter::once(ChatMessage::system(COMPACTION_PROMPT))
.chain(messages.iter().cloned())
.collect(),
)
.temperature(0.1)
.max_tokens(COMPACTION_MAX_TOKENS);
request.provider_attempt_id = Some(operation_id.into());
let (delta_tx, _delta_rx) = tokio::sync::mpsc::unbounded_channel();
let completion = self.model.complete_streaming(&request, delta_tx);
tokio::pin!(completion);
let timeout = tokio::time::sleep_until(deadline.into());
tokio::pin!(timeout);
let response = tokio::select! {
response = &mut completion => response.map_err(|error| RuntimeError::Model(error.to_string()))?,
_ = &mut timeout => return Err(RuntimeError::Model("compaction deadline exceeded".into())),
_ = cancellation.cancelled() => return Err(RuntimeError::Cancelled),
};
let summary = response
.first_content()
.filter(|value| !value.trim().is_empty())
.ok_or_else(|| RuntimeError::Model("compactor returned no summary".into()))?
.trim()
.to_string();
let (prompt_tokens, completion_tokens) = response.usage.map_or_else(
|| {
(
request
.messages
.iter()
.filter_map(|message| message.content.as_deref())
.map(|content| content.chars().count().div_ceil(4) as u64)
.sum(),
summary.chars().count().div_ceil(4).max(1) as u64,
)
},
|usage| {
(
u64::from(usage.prompt_tokens),
u64::from(usage.completion_tokens),
)
},
);
Ok(CompactionResult {
summary,
prompt_tokens,
completion_tokens,
})
}
}
impl AgentRuntime {
pub(super) async fn compact_if_needed(
&self,
writer: &dyn EventWriter,
run_id: &RunId,
step: u32,
transcript: &mut Vec<ChatMessage>,
context: &[ChatMessage],
cancellation: CancellationToken,
) -> Result<(u64, u64), RuntimeError> {
let count = |transcript: &[ChatMessage]| {
self.meter.count(
&self.model_name,
&std::iter::once(ChatMessage::system(self.prompts.render()))
.chain(transcript.iter().cloned())
.chain(context.iter().cloned())
.collect::<Vec<_>>(),
)
};
if count(transcript) <= self.limits.max_tokens {
return Ok((0, 0));
}
let mut pruned = Vec::new();
for message in transcript
.iter_mut()
.filter(|message| message.role == Role::Tool)
{
let Some(content) = &message.content else {
continue;
};
if content.chars().count() <= 2_000 {
continue;
}
if let Some(call_id) = &message.tool_call_id {
pruned.push(call_id.clone());
}
message.content = Some(format!(
"{}… [tool result pruned]",
content.chars().take(512).collect::<String>()
));
}
if !pruned.is_empty() {
writer
.append(vec![Event::ToolResultsPruned {
run_id: run_id.clone(),
call_ids: pruned,
}])
.await?;
if count(transcript) <= self.limits.max_tokens {
return Ok((0, 0));
}
}
let compactor = &self.compactor;
let compaction_id = format!("compaction:{step}");
let operation_id = format!("{compaction_id}:attempt:1");
let provider_attempt_id = format!("{run_id}:{operation_id}");
let source_through_seq = writer
.load_after(0)
.await?
.last()
.map_or(0, |event| event.seq);
let compactable = transcript
.iter()
.cloned()
.chain(context.iter().cloned())
.collect::<Vec<_>>();
let prepared = writer
.append(vec![
Event::CompactionStarted {
run_id: run_id.clone(),
compaction_id: compaction_id.clone(),
source_through_seq,
},
Event::ModelRequestPrepared {
run_id: run_id.clone(),
step,
attempt: 1,
provider_attempt_id: provider_attempt_id.clone(),
operation_id: operation_id.clone(),
reserved_prompt_tokens: self.meter.count(
&self.model_name,
&std::iter::once(ChatMessage::system(COMPACTION_PROMPT))
.chain(compactable.iter().cloned())
.collect::<Vec<_>>(),
),
reserved_completion_tokens: u64::from(COMPACTION_MAX_TOKENS),
request: serde_json::to_value(&compactable)
.map_err(|error| RuntimeError::Invariant(error.to_string()))?,
prompt_sections: serde_json::Value::Null,
},
])
.await?;
let prepared_through_seq = prepared
.last()
.map_or(source_through_seq, |event| event.seq);
let extension_cancellation = cancellation.child();
let deadline = std::time::Instant::now() + self.limits.provider_deadline;
let result = match crate::extension::await_extension(
"context compactor",
extension_cancellation.clone(),
deadline,
compactor.summarize(
&self.model_name,
&compactable,
&provider_attempt_id,
extension_cancellation,
deadline,
),
)
.await
{
Ok(Ok(result)) => result,
Ok(Err(error)) | Err(error) => {
writer
.append(vec![Event::CompactionFinished {
run_id: run_id.clone(),
compaction_id,
status: "failed".into(),
error: Some(error.to_string()),
}])
.await?;
return Err(error);
}
};
let current_seq = writer
.load_after(0)
.await?
.last()
.map_or(0, |event| event.seq);
if current_seq != prepared_through_seq {
writer
.append(vec![
Event::UsageRecorded {
run_id: run_id.clone(),
operation_id,
prompt_tokens: result.prompt_tokens,
completion_tokens: result.completion_tokens,
cost_units: 0,
},
Event::CompactionFinished {
run_id: run_id.clone(),
compaction_id,
status: "failed".into(),
error: Some("surface_changed".into()),
},
])
.await?;
return Err(RuntimeError::CompactionConflict);
}
transcript.clear();
transcript.push(ChatMessage::system(format!(
"Conversation summary:\n{}",
result.summary
)));
let mut events = vec![
Event::UsageRecorded {
run_id: run_id.clone(),
operation_id,
prompt_tokens: result.prompt_tokens,
completion_tokens: result.completion_tokens,
cost_units: 0,
},
Event::SummaryReplaced {
run_id: run_id.clone(),
through_seq: source_through_seq,
summary: result.summary,
compactor: compactor.name().into(),
model: self.model_name.clone(),
},
Event::CompactionFinished {
run_id: run_id.clone(),
compaction_id,
status: "completed".into(),
error: None,
},
];
if result.prompt_tokens == 0 && result.completion_tokens == 0 {
events.remove(0);
}
writer.append(events).await?;
Ok((result.prompt_tokens, result.completion_tokens))
}
}