MythEclipse/zesdex
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: enhance message shaping with progressive summarization and improved handling of dropped messages

Browse Source
This commit is contained in:
asepharyana
2026-07-17 10:29:52 +07:00
parent 5aad7e1eb1
commit 796bb09c5b
2 changed files with 383 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/zesdex-backend/src/app/runtime/actions/handlers.rs
+91 -17
View File
@@ -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) {
crates/zesdex-backend/src/app/runtime/context/shaping.rs
+292 -76
View File
@@ -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::<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,
@@ -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::<Vec<_>>()
.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<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) {
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]