diff --git a/src/app/runtime/context/shaping.rs b/src/app/runtime/context/shaping.rs new file mode 100644 index 0000000..451df31 --- /dev/null +++ b/src/app/runtime/context/shaping.rs @@ -0,0 +1,205 @@ +//! 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 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 +/// `[prior conversation compacted]` (or LLM-generated summary, if +/// `client` is `Some`) system message in between. +/// +/// Why: keeps context-size overhead roughly constant regardless of +/// session length. +/// +/// Return: the shaped message list, or `messages` unchanged if shaping +/// wasn't needed. +pub fn shape_messages( + messages: &[ChatMessage], + token_count: usize, + max_wire_tokens: usize, + force: bool, + client: Option<&crate::service::provider::LlmClient>, +) -> Vec { + 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(); + + let mut msgs_to_eval = messages.to_vec(); + let first = if msgs_to_eval.is_empty() { + None + } else { + Some(msgs_to_eval.remove(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() { + 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::>() + .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{content}\n]"); + } + } + Err(e) => { + tracing::warn!( + "[context::shaping] 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 +} + +#[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); + 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); + assert_eq!(result[0].content.as_deref(), Some("system prompt")); + } + + #[test] + fn shape_messages_without_a_client_falls_back_to_placeholder_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); + let has_placeholder = result.iter().any(|m| { + m.content.as_deref() == Some("[prior conversation compacted]") + }); + assert!(has_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); + let last_content = messages.last().unwrap().content.clone(); + assert!(result.iter().any(|m| m.content == last_content), "most recent message must survive shaping"); + } +}