From 796bb09c5b4c5d90b9374cb66aa36fe30af46a34 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Fri, 17 Jul 2026 10:29:21 +0700 Subject: [PATCH] refactor: enhance message shaping with progressive summarization and improved handling of dropped messages --- .../src/app/runtime/actions/handlers.rs | 108 ++++- .../src/app/runtime/context/shaping.rs | 368 ++++++++++++++---- 2 files changed, 383 insertions(+), 93 deletions(-) diff --git a/crates/zesdex-backend/src/app/runtime/actions/handlers.rs b/crates/zesdex-backend/src/app/runtime/actions/handlers.rs index 16bf6ca..ff7a2fe 100644 --- a/crates/zesdex-backend/src/app/runtime/actions/handlers.rs +++ b/crates/zesdex-backend/src/app/runtime/actions/handlers.rs @@ -209,28 +209,102 @@ pub(super) fn handle_abort_turn(state: &mut AppStateRest) { pub(super) fn handle_compact(state: &mut AppStateRest) { let max_wire_tokens = window::resolve(&state.app_config, &state.settings); - if let Some(ref mut rt) = state.session_runtime { - let token_estimate: usize = rt - .messages - .iter() - .filter_map(|m| m.content.as_deref()) - .map(count_tokens) - .sum(); - rt.messages = - crate::app::runtime::context::shaping::shape_messages( + // Extract config before borrowing session_runtime mutably to avoid + // borrow conflicts. An LLM client is needed for summarization so the + // compacted result preserves meaningful context (goals, decisions, + // files, state) instead of a useless static placeholder. + let api_key = state + .settings + .api_keys + .get(&state.settings.provider) + .cloned() + .unwrap_or_default(); + let model = state.settings.model.clone(); + let base_url = state + .app_config + .providers + .get(&state.settings.provider) + .map(|p| p.api_base.clone()); + let abort_flag = state.abort_flag.clone(); + + // Build the LLM client if we have a configured base_url. + let llm_client = base_url.map(|url| { + let key = if api_key.is_empty() { + state + .app_config + .providers + .get(&state.settings.provider) + .and_then(|cfg| { + cfg.api_key_env + .as_ref() + .and_then(|env| std::env::var(env).ok()) + }) + .or_else(|| { + state + .app_config + .providers + .get(&state.settings.provider) + .and_then(|cfg| cfg.default_api_key.clone()) + }) + .unwrap_or_else(|| crate::service::provider::DEFAULT_API_KEY.to_string()) + } else { + api_key.clone() + }; + crate::service::provider::LlmClient::new(key, model.clone(), Some(url)) + }); + + if llm_client.is_none() { + state.push_toast(Toast::new( + ToastKind::Error, + "Cannot compact: no AI provider configured. Set up a provider in Settings first." + .to_string(), + )); + return; + } + + let (before_tokens, after_tokens, msg_count) = + if let Some(ref mut rt) = state.session_runtime { + let token_estimate: usize = rt + .messages + .iter() + .filter_map(|m| m.content.as_deref()) + .map(count_tokens) + .sum(); + + let before = token_estimate; + rt.messages = crate::app::runtime::context::shaping::shape_messages( &rt.messages, token_estimate, max_wire_tokens, true, - None, - None, + llm_client.as_ref(), + Some(&*abort_flag), ); - state.push_toast(Toast::new( - ToastKind::Success, - "Conversation history compacted.".to_string(), - )); - state.dirty = true; - } + let after: usize = rt + .messages + .iter() + .filter_map(|m| m.content.as_deref()) + .map(count_tokens) + .sum(); + (before, after, rt.messages.len()) + } else { + (0, 0, 0) + }; + + let dropped = before_tokens.saturating_sub(after_tokens); + let msg_label = if before_tokens > 0 { + format!( + "Compacted ({} msgs, ~{}K → ~{}K tokens, dropped ~{}K).", + msg_count, + before_tokens / 1000, + after_tokens / 1000, + dropped / 1000, + ) + } else { + "No active session to compact.".to_string() + }; + state.push_toast(Toast::new(ToastKind::Success, msg_label)); + state.dirty = true; } pub(super) fn handle_lesson_accept(state: &mut AppStateRest, name: String) { diff --git a/crates/zesdex-backend/src/app/runtime/context/shaping.rs b/crates/zesdex-backend/src/app/runtime/context/shaping.rs index 205ad42..61db014 100644 --- a/crates/zesdex-backend/src/app/runtime/context/shaping.rs +++ b/crates/zesdex-backend/src/app/runtime/context/shaping.rs @@ -36,19 +36,206 @@ pub fn should_shape(token_estimate: usize, max_wire_tokens: usize, prev_shaped: /// 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. +/// 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 (spawned in `spawn_turn`), so the blocking -/// summarization call does not freeze the UI. +/// 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. +/// 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::>() + .join("\n\n---\n\n") +} + +/// Extract the content of a previous compaction summary from a message. +fn extract_prior_summary(m: &ChatMessage) -> Option { + 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, @@ -61,10 +248,9 @@ pub fn shape_messages( 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(); + 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() { @@ -73,15 +259,35 @@ pub fn shape_messages( 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); + let mut keep_recent = Vec::new(); + let mut dropped_msgs = Vec::new(); - if current_tokens + msg_tokens <= target_tokens { - current_tokens += msg_tokens; - keep_recent.push(m); - } else { - dropped_msgs.push(m); + 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); + } } } @@ -93,69 +299,65 @@ pub fn shape_messages( } if !dropped_msgs.is_empty() { - let mut summary_text = "[prior conversation compacted]".to_string(); - - if let Some(llm) = client { + // 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. - // The turn loop already runs on a background thread (spawned in - // `spawn_turn`), so this blocking call does not freeze the UI. let aborted = abort_flag.is_some_and(|f| f.load(Ordering::SeqCst)); if !aborted { - let prompt = format!( - "You are a context-preservation summarizer. 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.\n\n\ - History to summarize:\n{}", - dropped_msgs.iter() - .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::>() - .join("\n\n---\n\n") - ); - + let dropped_content = format_dropped_messages(&dropped_msgs); + let prompt = build_summarization_prompt(&dropped_msgs, &dropped_content); 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]"); + + // Try summarization with one retry on failure. + let mut result: Option = None; + let mut last_err: Option = None; + for attempt in 0..2 { + match llm.chat_with_tools_non_streaming(&req_msgs, None) { + 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, + ); + } } } - Err(e) => { - tracing::warn!( - "[context::shaping] LLM summarization failed: {}. \ - Falling back to static placeholder.", - e, - ); + } + + // 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)); } @@ -224,16 +426,30 @@ mod tests { } #[test] - fn shape_messages_without_a_client_falls_back_to_placeholder_summary() { + 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 has_placeholder = result - .iter() - .any(|m| m.content.as_deref() == Some("[prior conversation compacted]")); - assert!(has_placeholder); + 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]