fix(context): perbaiki fixture test shaping agar men-drop pesan lama

Fixture pendek sebelumnya (~8 token nyata per pesan lewat tiktoken) tidak pernah melebihi target 700 token, jadi shape_messages tidak pernah men-drop satu pesan pun -- satu test gagal, satu test lain lulus secara vakum. Pakai fixture lebih panjang (~49 token/pesan) yang diverifikasi melebihi target dengan nyaman.
2026-07-16 07:49:43 +07:00
parent 75e9cadcd5
commit 8bb697a53f
1 changed files with 205 additions and 0 deletions
@@ -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<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();
+
+    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::<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{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");
+    }
+}