zesdex/crates/zesdex-backend/src/app/runtime/context/shaping.rs

//! Budget-based message shaping: compacts long conversation histories so
//! they fit within the provider's context window before being sent to
//! the LLM API. Ported from the former `runtime::shortsend` — behavior
//! is unchanged, only its token-counting now goes through
//! `context::tokens` instead of an inline heuristic.
use std::sync::atomic::{AtomicBool, Ordering};

use super::tokens::count_tokens;
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 {
        (max_wire_tokens as f32 * 0.95) as usize
    } else {
        (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 that fit a 70%-of-budget target, with a
/// summary system message (LLM-generated, or static placeholder as
/// last resort) in between.
///
/// When `force=true` (manual `/compact`), a different policy applies:
/// always compact by keeping the system message + at most 15 recent
/// messages. This guarantees the user-requested compaction always has
/// an effect, unlike auto-compaction which only triggers when the 70%
/// budget is exceeded.
///
/// **Progressive summarization**: if the dropped messages include a
/// previous compaction summary (e.g. `[Summary of compacted prior
/// conversation:...]`), that existing summary is extracted and passed
/// alongside the newly dropped messages. The LLM then produces an
/// updated summary that builds on the old one instead of starting
/// from scratch — preserving context across multiple compactions.
///
/// The LLM summarization checks the abort flag before calling the LLM,
/// so a user-requested abort is respected promptly. The turn loop already
/// runs on a background thread, so the blocking call does not freeze the UI.
///
/// Why: keeps context-size overhead roughly constant regardless of
/// session length while progressively preserving high-level context.
///
/// Return: the shaped message list, or `messages` unchanged if shaping
/// wasn't needed.
const FORCE_KEEP_MAX: usize = 15;

/// Prefix of a previous compaction summary. Used to detect progressive
/// summarization opportunities and to scan for prior summaries.
const SUMMARY_PREFIX: &str = "[Summary of compacted prior conversation:";

/// Detect whether a message contains a previous compaction summary.
fn msg_has_prior_summary(m: &ChatMessage) -> bool {
    m.content
        .as_deref()
        .is_some_and(|c| c.starts_with(SUMMARY_PREFIX))
}

/// Format dropped messages for the summarization prompt, excluding any
/// messages that are themselves previous summaries (those are handled
/// separately by progressive summarization).
fn format_dropped_messages(dropped: &[ChatMessage]) -> String {
    dropped
        .iter()
        .filter(|m| !msg_has_prior_summary(m))
        .map(|m| {
            let role_label = match m.role {
                crate::dto::chat::message::Role::User => "User",
                crate::dto::chat::message::Role::Assistant => "Assistant",
                crate::dto::chat::message::Role::System => "System",
                crate::dto::chat::message::Role::Tool => "Tool",
            };
            let has_tool_calls = m.tool_calls.is_some()
                && m.tool_calls.as_ref().is_some_and(|c| !c.is_empty());
            let mut entry =
                format!("[{role_label}]: {}", m.content.as_deref().unwrap_or(""));
            if has_tool_calls {
                if let Some(calls) = &m.tool_calls {
                    let names: Vec<&str> =
                        calls.iter().map(|c| c.function.name.as_str()).collect();
                    entry.push_str(&format!("\n  [tool calls: {}]", names.join(", ")));
                }
            }
            entry
        })
        .collect::<Vec<_>>()
        .join("\n\n---\n\n")
}

/// Extract the content of a previous compaction summary from a message.
fn extract_prior_summary(m: &ChatMessage) -> Option<String> {
    let content = m.content.as_deref()?;
    if content.starts_with(SUMMARY_PREFIX) {
        // Strip the `[Summary of compacted prior conversation:\n` prefix
        // and the trailing `\n]`.
        let inner = content
            .strip_prefix(SUMMARY_PREFIX)?
            .strip_suffix(']')?
            .trim();
        Some(inner.to_string())
    } else {
        None
    }
}

/// Build the summarization prompt, supporting progressive compaction:
/// if the dropped messages contain a previous summary, it is extracted
/// and the new prompt asks the LLM to build on it.
fn build_summarization_prompt(
    dropped_msgs: &[ChatMessage],
    dropped_content: &str,
) -> String {
    // Check for a prior summary among dropped messages.
    let prior = dropped_msgs.iter().find_map(extract_prior_summary);

    let base = "You are a context-preservation summarizer for an AI coding assistant. \
         The following conversation history is being dropped to free up context window space. \
         Produce a structured summary that preserves the information the AI \
         agent needs to continue working seamlessly.\n\n\
         Structure your summary into these sections:\n\
         1. **Goals & Objectives** — what the user asked for, what tasks remain\n\
         2. **Key Decisions** — architectural choices, design decisions, approach changes\n\
         3. **Files Modified/Created** — paths and brief description of changes\n\
         4. **Findings & State** — important discoveries, test results, current state\n\
         5. **Open Items** — unresolved issues, pending tasks, next steps\n\n\
         Be concise but thorough. Preserve file paths, error messages, and \
         specific details the agent needs to continue.";

    match prior {
        Some(prev) => {
            // Progressive compaction: the LLM already summarized earlier
            // parts of this conversation. Build on it rather than starting
            // from scratch.
            format!(
                "{base}\n\n\
                 ### Previous summary (build on this, don't repeat it):\n\
                 {prev}\n\n\
                 ### New messages to merge into the summary:\n\
                 {dropped_content}\n\n\
                 Produce the **complete updated summary** with all 5 sections, \
                 incorporating both the previous summary and the new messages.",
            )
        }
        None => {
            format!(
                "{base}\n\n\
                 ### History to summarize:\n\
                 {dropped_content}",
            )
        }
    }
}

/// Build a structural summary of dropped messages when AI summarization
/// is unavailable or failed. This is far more useful than a static
/// `[prior conversation compacted]` placeholder — it tells the LLM how
/// many messages of each role were dropped and what tools were used,
/// preserving key structural context.
fn make_structural_summary(dropped: &[ChatMessage]) -> String {
    use std::fmt::Write;

    let user_count = dropped
        .iter()
        .filter(|m| m.role == crate::dto::chat::message::Role::User)
        .count();
    let assistant_count = dropped
        .iter()
        .filter(|m| m.role == crate::dto::chat::message::Role::Assistant)
        .count();
    let tool_count = dropped
        .iter()
        .filter(|m| m.role == crate::dto::chat::message::Role::Tool)
        .count();
    let system_count = dropped
        .iter()
        .filter(|m| m.role == crate::dto::chat::message::Role::System)
        .count();

    // Collect unique tool names used in dropped assistant messages.
    let mut tool_names: Vec<&str> = dropped
        .iter()
        .filter_map(|m| m.tool_calls.as_ref())
        .flatten()
        .map(|c| c.function.name.as_str())
        .collect();
    tool_names.sort_unstable();
    tool_names.dedup();

    // Extract the last user message content as a hint about what was
    // being discussed.
    let last_user_content = dropped
        .iter()
        .rev()
        .find(|m| m.role == crate::dto::chat::message::Role::User)
        .and_then(|m| m.content.as_deref());

    let mut summary = String::new();
    write!(
        summary,
        "[prior conversation: {} user, {} assistant, {} tool, {} system messages",
        user_count, assistant_count, tool_count, system_count,
    )
    .unwrap();
    if !tool_names.is_empty() {
        write!(summary, " | tools used: {}", tool_names.join(", ")).unwrap();
    }
    if let Some(content) = last_user_content {
        // Only include the first line to keep it compact.
        let hint = content.lines().next().unwrap_or(content);
        let truncated = if hint.len() > 120 {
            &hint[..120]
        } else {
            hint
        };
        write!(summary, " | last request: {truncated}").unwrap();
    }
    summary.push(']');
    summary
}

pub fn shape_messages(
    messages: &[ChatMessage],
    token_count: usize,
    max_wire_tokens: usize,
    force: bool,
    client: Option<&crate::service::provider::LlmClient>,
    abort_flag: Option<&AtomicBool>,
) -> Vec<ChatMessage> {
    if !force && (token_count <= max_wire_tokens || messages.len() < 5) {
        return messages.to_vec();
    }

    if force && messages.len() < 5 {
        return messages.to_vec();
    }

    let mut msgs_to_eval = messages.to_vec();
    let first = if msgs_to_eval.is_empty() {
        None
    } else {
        Some(msgs_to_eval.remove(0))
    };

    let mut keep_recent = Vec::new();
    let mut dropped_msgs = Vec::new();

    if force {
        // Manual compaction (/compact): keep system + at most N recent messages.
        // This guarantees compaction always has an effect regardless of
        // conversation size, unlike auto-compaction which depends on budget.
        for m in msgs_to_eval.into_iter().rev() {
            if keep_recent.len() < FORCE_KEEP_MAX {
                keep_recent.push(m);
            } else {
                dropped_msgs.push(m);
            }
        }
    } else {
        // Auto-compaction: keep messages that fit within 70% of context window.
        let target_tokens = (max_wire_tokens as f32 * 0.70) as usize;
        let mut current_tokens = 0;

        for m in msgs_to_eval.into_iter().rev() {
            let text = m.content.as_deref().unwrap_or("");
            let msg_tokens = count_tokens(text);

            if current_tokens + msg_tokens <= target_tokens {
                current_tokens += msg_tokens;
                keep_recent.push(m);
            } else {
                dropped_msgs.push(m);
            }
        }
    }

    dropped_msgs.reverse();

    let mut result = Vec::new();
    if let Some(f) = first {
        result.push(f);
    }

    if !dropped_msgs.is_empty() {
        // Always prefer AI-generated summary. The hardcoded placeholder
        // `[prior conversation compacted]` is never used — it provides
        // zero useful context to the LLM and defeats the purpose of
        // compaction. Instead we produce a minimal structural summary.
        let summary_text: String = if let Some(llm) = client {
            // Check abort before starting the blocking LLM summarization call.
            let aborted = abort_flag.is_some_and(|f| f.load(Ordering::SeqCst));
            if !aborted {
                let dropped_content = format_dropped_messages(&dropped_msgs);
                let prompt = build_summarization_prompt(&dropped_msgs, &dropped_content);
                let req_msgs = vec![ChatMessage::user(prompt)];

                // Try summarization with one retry on failure.
                let mut result: Option<String> = None;
                let mut last_err: Option<anyhow::Error> = None;
                for attempt in 0..2 {
                    match llm.chat_with_tools_non_streaming(&req_msgs, None, None, None, abort_flag) {
                        Ok(resp) => {
                            if let Some(content) = resp.0.content {
                                result = Some(format!(
                                    "[Summary of compacted prior conversation:\n{content}\n]"
                                ));
                                break;
                            }
                        }
                        Err(e) => {
                            last_err = Some(e);
                            if attempt == 0 {
                                tracing::info!(
                                    "[context::shaping] summarization attempt {} failed, retrying...",
                                    attempt + 1,
                                );
                            }
                        }
                    }
                }

                // If all retries failed, produce a basic structural summary
                // instead of a useless placeholder.
                match result {
                    Some(s) => s,
                    None => {
                        if let Some(e) = last_err {
                            tracing::warn!(
                                "[context::shaping] LLM summarization failed after retry: {}. \
                                 Falling back to structural summary.",
                                e,
                            );
                        }
                        make_structural_summary(&dropped_msgs)
                    }
                }
            } else {
                make_structural_summary(&dropped_msgs)
            }
        } else {
            // No LLM client available (tests / edge case with no provider).
            make_structural_summary(&dropped_msgs)
        };

        result.push(ChatMessage::system(summary_text));
    }

    result.extend(keep_recent.into_iter().rev());
    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dto::chat::message::ChatMessage;

    #[test]
    fn should_shape_triggers_at_85_percent_when_not_previously_shaped() {
        assert!(should_shape(850, 1000, false));
        assert!(!should_shape(849, 1000, false));
    }

    #[test]
    fn should_shape_uses_95_percent_threshold_once_already_shaped() {
        assert!(
            !should_shape(900, 1000, true),
            "below 95% and already shaped: no re-trigger yet"
        );
        assert!(should_shape(950, 1000, true));
    }

    #[test]
    fn shape_messages_is_a_noop_under_budget_and_not_forced() {
        let messages = vec![
            ChatMessage::system("sys"),
            ChatMessage::user("hi"),
            ChatMessage::assistant(Some("hello".to_string())),
        ];
        let result = shape_messages(&messages, 10, 1000, false, None, None);
        assert_eq!(result.len(), messages.len());
    }

    /// Build a message whose real BPE token count is large enough that 20
    /// of them (~49 tokens each, ~980 total — verified empirically with
    /// `context::tokens::count_tokens`) comfortably exceed
    /// `shape_messages`'s 70%-of-1000 = 700 token target, guaranteeing
    /// several get dropped. A short fixture like `format!("message {i}")`
    /// (~8 tokens each, ~160 total for 20) stays entirely under budget
    /// with real BPE counting and would make these tests pass vacuously
    /// (nothing ever gets dropped, so "must survive shaping" and "falls
    /// back to placeholder" hold trivially without exercising the actual
    /// drop logic) — this was a real bug caught during Task 5's first
    /// implementation attempt.
    fn padded_message(i: usize) -> String {
        format!(
            "message number {i} with some padding text {}",
            "additional padding content to increase token count substantially ".repeat(5),
        )
    }

    #[test]
    fn shape_messages_always_preserves_the_first_system_message() {
        let mut messages = vec![ChatMessage::system("system prompt")];
        for i in 0..20 {
            messages.push(ChatMessage::user(padded_message(i)));
        }
        let result = shape_messages(&messages, 100_000, 1000, true, None, None);
        assert_eq!(result[0].content.as_deref(), Some("system prompt"));
    }

    #[test]
    fn shape_messages_without_a_client_falls_back_to_structural_summary() {
        let mut messages = vec![ChatMessage::system("system prompt")];
        for i in 0..20 {
            messages.push(ChatMessage::user(padded_message(i)));
        }
        let result = shape_messages(&messages, 100_000, 1000, true, None, None);
        let summary_msg = result.iter().find(|m| {
            m.content
                .as_deref()
                .is_some_and(|c| c.starts_with("[prior conversation:"))
        });
        assert!(
            summary_msg.is_some(),
            "must contain a structural summary, not a hardcoded placeholder"
        );
        let content = summary_msg.unwrap().content.as_deref().unwrap();
        assert!(
            content.contains("user"),
            "structural summary must include message counts, got: {content}"
        );
        assert!(
            !content.contains("[prior conversation compacted]"),
            "must NOT contain the useless hardcoded placeholder"
        );
    }

    #[test]
    fn shape_messages_keeps_most_recent_messages_over_older_ones() {
        let mut messages = vec![ChatMessage::system("system prompt")];
        for i in 0..20 {
            messages.push(ChatMessage::user(padded_message(i)));
        }
        let result = shape_messages(&messages, 100_000, 1000, true, None, None);
        let last_content = messages.last().unwrap().content.clone();
        assert!(
            result.iter().any(|m| m.content == last_content),
            "most recent message must survive shaping"
        );
    }
}