use std::sync::Arc;
use std::time::Instant;
use af_agent::ChatModel;
use af_agent_session::Event;
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 provider_name(&self) -> Option<&str> {
self.model.provider_name()
}
fn name(&self) -> &str {
"agentfactory.model-compactor"
}
async fn summarize(
&self,
request: &CompletionRequest,
cancellation: CancellationToken,
deadline: Instant,
) -> Result<CompactionResult, RuntimeError> {
let mut request = request.clone();
request
.messages
.insert(0, ChatMessage::system(COMPACTION_PROMPT));
request = request.temperature(0.1).max_tokens(COMPACTION_MAX_TOKENS);
af_llm::images::select_images(&mut request.messages, af_llm::images::MAX_INPUT_IMAGES)
.map_err(|error| RuntimeError::InvalidInput(error.to_string()))?;
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 {
source: if response.usage.is_some() {
af_agent_session::MeteringSource::Reported
} else {
af_agent_session::MeteringSource::Estimated
},
provider: self.model.provider_name().map(str::to_owned),
summary,
prompt_tokens,
completion_tokens,
})
}
}
impl AgentRuntime {
pub(super) async fn compact_if_needed(
&self,
writer: &dyn EventWriter,
turn: &crate::TurnRequest,
step: u32,
transcript: &mut Vec<ChatMessage>,
context: &[ChatMessage],
cancellation: CancellationToken,
) -> Result<(u64, u64), RuntimeError> {
let run_id = &turn.run_id;
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 !turn.compact && 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 !turn.compact && 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.last_seq().await?;
let mut compactable = transcript
.iter()
.cloned()
.chain(context.iter().cloned())
.collect::<Vec<_>>();
af_llm::images::select_images(&mut compactable, self.limits.max_images)
.map_err(|error| RuntimeError::InvalidInput(error.to_string()))?;
let mut compaction_request = CompletionRequest::new(&self.model_name, compactable.clone());
compaction_request.context = Some(turn.context.clone());
compaction_request.provider_attempt_id = Some(provider_attempt_id.clone());
let prepared = writer
.append_at(
source_through_seq,
vec![
Event::CompactionStarted {
run_id: run_id.clone(),
compaction_id: compaction_id.clone(),
source_through_seq,
},
Event::ModelRequestPrepared {
metering: Some(af_agent_session::MeteringDetails::Model {
model: self.model_name.clone(),
provider: compactor.provider_name().map(str::to_owned),
provider_attempt_id: provider_attempt_id.parse().map_err(
|error: af_context::ids::EmptyId| {
RuntimeError::Invariant(error.to_string())
},
)?,
source: af_agent_session::MeteringSource::Estimated,
outcome: af_agent_session::MeteringOutcome::Unknown,
}),
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(&compaction_request, 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 attribution = Some(af_agent_session::MeteringDetails::Model {
model: self.model_name.clone(),
provider: result.provider.clone(),
provider_attempt_id: provider_attempt_id.parse().map_err(
|error: af_context::ids::EmptyId| RuntimeError::Invariant(error.to_string()),
)?,
source: result.source,
outcome: af_agent_session::MeteringOutcome::Completed,
});
let events = vec![
Event::UsageRecorded {
metering: attribution.clone(),
run_id: run_id.clone(),
operation_id: operation_id.clone(),
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.clone(),
compactor: compactor.name().into(),
model: self.model_name.clone(),
},
Event::CompactionFinished {
run_id: run_id.clone(),
compaction_id: compaction_id.clone(),
status: "completed".into(),
error: None,
},
];
match writer.append_at(prepared_through_seq, events).await {
Ok(_) => crate::replay::replace_with_summary(transcript, &result.summary),
Err(RuntimeError::CompactionConflict) => {
writer
.append(vec![
Event::UsageRecorded {
metering: attribution.clone(),
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);
}
Err(error) => return Err(error),
}
Ok((result.prompt_tokens, result.completion_tokens))
}
}