use kernel::consts_policy::EXCHANGE_BUDGET_BYTES;
use kernel::{AxCode, AxError, ByteLen, ContentBlock};
use super::{Shrink, detect, plan, shorten};
#[derive(Debug, Clone, Default, PartialEq)]
pub struct Exchange {
assistant: Vec<ContentBlock>,
results: Vec<ContentBlock>,
}
impl Exchange {
#[must_use]
pub fn new() -> Exchange {
Exchange::default()
}
pub fn push_assistant(&mut self, content: Vec<ContentBlock>) {
self.assistant.extend(content);
}
pub fn push_result(&mut self, block: ContentBlock) {
self.results.push(block);
}
#[must_use]
pub fn assistant(&self) -> &[ContentBlock] {
&self.assistant
}
#[must_use]
pub fn results(&self) -> &[ContentBlock] {
&self.results
}
pub fn compact(&mut self) -> Result<(), AxError> {
let mut texts = Vec::new();
for block in self.assistant.iter_mut().chain(self.results.iter_mut()) {
match block {
ContentBlock::Text { text } => texts.push(text),
ContentBlock::ToolResult { content, .. } => texts.push(content),
ContentBlock::Thinking { .. }
| ContentBlock::RedactedThinking { .. }
| ContentBlock::ToolUse { .. }
| ContentBlock::Image(_) => {}
}
}
compact_texts(texts)
}
}
fn compact_texts(texts: Vec<&mut String>) -> Result<(), AxError> {
let budget = ByteLen::new(EXCHANGE_BUDGET_BYTES);
let mut total = 0u64;
let mut count = 0u64;
for text in &texts {
if text.is_empty() {
continue;
}
total = total.checked_add(bytes(text)?).ok_or_else(uncountable)?;
count = count.checked_add(1).ok_or_else(uncountable)?;
}
if count == 0 || total <= budget.get() {
return Ok(());
}
let share = budget.get().checked_div(count).ok_or_else(uncountable)?;
let share = ByteLen::new(share);
for text in texts {
if text.is_empty() {
continue;
}
let size = ByteLen::new(bytes(text)?);
match plan(detect(text), size, share) {
Shrink::Keep | Shrink::MustOffload => {}
Shrink::Cut(strategy) => {
let cut = shorten(text, strategy, share);
*text = cut.text;
}
}
}
Ok(())
}
fn bytes(text: &str) -> Result<u64, AxError> {
u64::try_from(text.len()).map_err(|_| uncountable())
}
fn uncountable() -> AxError {
AxError::failure(
AxCode::InvalidArgs,
"compact one turn's exchange",
"a window text larger than a u64 byte count",
)
.with_recovery(
"report this against runtime::compaction::exchange: no tool this city ships \
can put a text this size in the window",
)
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
clippy::indexing_slicing,
clippy::arithmetic_side_effects,
clippy::wildcard_enum_match_arm,
reason = "test code"
)]
mod tests;