use rig::completion::message::{Message, ToolResultContent, UserContent};
use serde::{Deserialize, Serialize};
use super::compaction::{compact_text, count_tokens, message_content_text};
use crate::config;
use crate::tools::{TodoItem, ToolContext};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Transcript {
#[serde(default)]
pub summary: Option<String>,
pub messages: Vec<Message>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub struct TokenEstimate {
pub messages: usize,
pub system_tokens: usize,
pub message_tokens: usize,
pub total_tokens: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub struct CompactionStats {
pub removed_messages: usize,
pub compacted_tools: usize,
pub summarized: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub struct ContextStatus {
pub estimate: TokenEstimate,
pub limit_tokens: usize,
pub input_budget_tokens: usize,
pub trigger_tokens: usize,
pub summary_present: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ContextBudgetExceeded {
pub estimated_tokens: usize,
pub input_budget_tokens: usize,
pub limit_tokens: usize,
}
impl std::fmt::Display for ContextBudgetExceeded {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"context estimate {} exceeds input budget {}; use /compact, temporarily raise OY_CONTEXT_LIMIT, or force-truncate history",
self.estimated_tokens, self.input_budget_tokens
)
}
}
impl std::error::Error for ContextBudgetExceeded {}
impl Transcript {
pub fn new() -> Self {
Self {
summary: None,
messages: Vec::new(),
}
}
pub fn undo_last_turn(&mut self) -> bool {
for index in (0..self.messages.len()).rev() {
if is_user_prompt(&self.messages[index]) {
self.messages.truncate(index);
return true;
}
}
false
}
pub fn force_truncate_oldest_turns(&mut self) -> usize {
if self.messages.len() <= 1 {
return 0;
}
let remove_count = (self.messages.len() / 4)
.max(1)
.min(self.messages.len() - 1);
let keep_from = self.valid_keep_from(remove_count);
if keep_from == 0 || keep_from >= self.messages.len() {
return 0;
}
self.messages.drain(..keep_from);
keep_from
}
pub fn token_estimate(
&self,
model: &str,
system_prompt: &str,
todos: &[TodoItem],
) -> TokenEstimate {
let count_text = |text: &str| count_tokens(model, text);
let system_tokens = count_text(system_prompt) + if todos.is_empty() { 0 } else { 4 };
let summary_tokens = self
.summary
.as_ref()
.map(|summary| 4 + count_text(summary))
.unwrap_or(0);
let message_tokens = summary_tokens
+ self
.messages
.iter()
.map(|message| 4 + count_text(&message_content_text(message)))
.sum::<usize>();
TokenEstimate {
messages: self.messages.len() + usize::from(self.summary.is_some()),
system_tokens,
message_tokens,
total_tokens: system_tokens + message_tokens,
}
}
pub fn compact_tool_outputs(&mut self, max_bytes: usize) -> usize {
let mut compacted = 0;
for message in &mut self.messages {
let Message::User { content } = message else {
continue;
};
for item in content.iter_mut() {
let UserContent::ToolResult(result) = item else {
continue;
};
for part in result.content.iter_mut() {
let ToolResultContent::Text(text) = part else {
continue;
};
if text.text.contains("[tool output compacted]") {
continue;
}
text.text = compact_text(&text.text, max_bytes, "tool output compacted");
if text.text.contains("[tool output compacted]") {
compacted += 1;
}
}
}
}
compacted
}
pub fn deterministic_compact_old_turns(
&mut self,
recent_messages: usize,
summary_bytes: usize,
) -> Option<CompactionStats> {
if self.messages.len() <= 1 {
return None;
}
let protected = recent_messages.max(1).min(self.messages.len() - 1);
let keep_from = self.valid_keep_from(self.messages.len() - protected);
if keep_from == 0 {
return None;
}
let removed_messages = keep_from;
self.truncate_with_note(keep_from, summary_bytes);
Some(CompactionStats {
removed_messages,
compacted_tools: 0,
summarized: true,
})
}
pub fn replace_turn_from_rig(&mut self, start: usize, messages: Vec<Message>) {
self.messages.truncate(start.min(self.messages.len()));
self.messages.extend(messages);
}
pub fn to_messages(&self) -> Vec<Message> {
let mut messages = Vec::new();
if let Some(summary) = self.summary.as_ref().filter(|s| !s.trim().is_empty()) {
messages.push(Message::user(format!(
"[Compacted earlier conversation]\n{}",
summary.trim()
)));
}
messages.extend(self.messages.clone());
messages
}
pub fn request_preamble(&self, system_prompt: &str, tool_context: &ToolContext) -> String {
let mut prompt = system_prompt.to_string();
if !tool_context.todos.is_empty() {
let header = config::session_text_value("transcript", "todo_system")
.unwrap_or_else(|_| String::from("{todos}"));
let todos = crate::tools::format_todos(&tool_context.todos);
prompt.push_str("\n\n");
prompt.push_str(header.replace("{todos}", todos.trim_end()).trim());
}
prompt
}
fn valid_keep_from(&self, requested: usize) -> usize {
let mut keep_from = requested.min(self.messages.len());
while keep_from < self.messages.len() && !is_user_prompt(&self.messages[keep_from]) {
keep_from += 1;
}
keep_from
}
fn truncate_with_note(&mut self, keep_from: usize, summary_bytes: usize) {
let existing = self.summary.take();
let note = format!(
"[context truncated] Removed {keep_from} older messages; retained the most recent conversation window."
);
let merged = existing
.filter(|s| !s.trim().is_empty())
.map(|summary| format!("{}\n\n{}", summary.trim(), note))
.unwrap_or(note);
self.summary = Some(compact_text(&merged, summary_bytes, "truncated summary"));
self.messages = self.messages.split_off(keep_from.min(self.messages.len()));
}
}
fn is_user_prompt(message: &Message) -> bool {
let Message::User { content } = message else {
return false;
};
content
.iter()
.any(|item| matches!(item, UserContent::Text(text) if !text.text.trim().is_empty()))
}