zesdex/src/app/runtime/shortsend.rs

//! Short-send / message shaping: compacts long conversation histories so
//! they fit within the provider's context window before being sent to the
//! LLM API.
use crate::dto::chat::message::ChatMessage;

/// Decide whether the message list should be shaped (compacted) before
/// sending to the LLM.
///
/// Flow: trigger based on token estimate. If `token_estimate` exceeds
/// the threshold, we shape. When `prev_shaped` is true, the threshold is
/// raised (95%) to avoid fluttering — compaction only re-triggers when
/// the context is genuinely full again. When `prev_shaped` is false, the
/// threshold is lower (85%) so compaction starts proactively.
///
/// Why: hysteresis prevents repeated compaction on every turn when the
/// token count hovers near the boundary.
///
/// Return: `true` if shaping should be applied.
pub fn should_shape(token_estimate: usize, max_wire_tokens: usize, prev_shaped: bool) -> bool {
    let threshold = if prev_shaped {
        // Higher threshold when already shaped — defer re-shaping until
        // the buffer is genuinely full again (95%).
        (max_wire_tokens as f32 * 0.95) as usize
    } else {
        // Lower threshold when not yet shaped — trigger shaping sooner
        // (85%) to avoid hitting the context window limit.
        (max_wire_tokens as f32 * 0.85) as usize
    };
    token_estimate >= threshold
}

/// Compact a long message list by dropping middle messages and inserting
/// a summary placeholder.
///
/// Flow: if the estimated token count is within budget and not forced, return messages
/// unchanged → otherwise keep the system message and the most recent
/// messages (up to `MAX_WIRE_TOKENS / 200` of them) with a `[prior
/// conversation compacted]` system message in between.
///
/// Why: keeps context-size overhead roughly constant regardless of
/// session length.
///
pub fn shape_messages(
    messages: &[ChatMessage],
    token_count: usize,
    max_wire_tokens: usize,
    force: bool,
    client: Option<&crate::service::provider::LlmClient>,
) -> Vec<ChatMessage> {
    if !force && (token_count <= max_wire_tokens || messages.len() < 5) {
        return messages.to_vec();
    }
    
    let target_tokens = (max_wire_tokens as f32 * 0.70) as usize;
    let mut current_tokens = 0;
    let mut keep_recent = Vec::new();
    let mut dropped_msgs = Vec::new();

    // Always keep the very first message (System Prompt) which we don't count here
    // as we just blindly preserve it later.
    let mut msgs_to_eval = messages.to_vec();
    let first = if !msgs_to_eval.is_empty() {
        Some(msgs_to_eval.remove(0))
    } else {
        None
    };

    // Iterate backwards from the most recent to oldest
    for m in msgs_to_eval.into_iter().rev() {
        let text = m.content.as_deref().unwrap_or("");
        // Estimate tokens: ~1 token per 3 bytes for mixed content (code,
    // prose, multi-byte). Conservative enough to stay under provider
    // limits while avoiding premature compaction.
    let msg_tokens = text.len() / 3;
        
        if current_tokens + msg_tokens <= target_tokens {
            current_tokens += msg_tokens;
            keep_recent.push(m);
        } else {
            dropped_msgs.push(m); // These will end up in reverse chronological order
        }
    }
    
    // Reverse dropped_msgs so they are back in chronological order
    dropped_msgs.reverse();
    
    let mut result = Vec::new();
    if let Some(f) = first {
        result.push(f);
    }
    
    if !dropped_msgs.is_empty() {
        let mut summary_text = "[prior conversation compacted]".to_string();
        
            if let Some(llm) = client {
                let prompt = format!(
                    "Summarize the following dropped conversation history briefly. Focus on main goals, decisions made, and files modified, so the context is preserved for future turns. Keep it concise.\n\nHistory:\n{}",
                    dropped_msgs.iter()
                        .map(|m| format!("[{}]: {}", if m.role == crate::dto::chat::message::Role::User { "User" } else { "Assistant" }, m.content.as_deref().unwrap_or("")))
                        .collect::<Vec<_>>()
                        .join("\n\n")
                );

                let req_msgs = vec![ChatMessage::user(prompt)];
                match llm.chat_with_tools_non_streaming(&req_msgs, None) {
                    Ok(resp) => {
                        if let Some(content) = resp.0.content {
                            summary_text = format!("[Summary of compacted prior conversation:\n{}\n]", content);
                        }
                    }
                    Err(e) => {
                        tracing::warn!(
                            "[shortsend] LLM summarization failed: {}. \
                             Prior conversation history is lost — no summary available. \
                             This means the model will lose context about earlier parts of \
                             the conversation.",
                            e,
                        );
                    }
                }
            }
        
        result.push(ChatMessage::system(summary_text));
    }
    
    result.extend(keep_recent.into_iter().rev());
    result
}