From 5aad7e1eb116d63e65dd9156d50992daad8f7a1a Mon Sep 17 00:00:00 2001 From: asepharyana Date: Fri, 17 Jul 2026 09:40:18 +0700 Subject: [PATCH] refactor: streamline token counting and message shaping logic --- .../src/app/runtime/actions/handlers.rs | 18 ++-- .../src/app/runtime/actions/turn.rs | 13 +-- .../src/app/runtime/context/shaping.rs | 95 +++++++++++++------ 3 files changed, 81 insertions(+), 45 deletions(-) diff --git a/crates/zesdex-backend/src/app/runtime/actions/handlers.rs b/crates/zesdex-backend/src/app/runtime/actions/handlers.rs index fab754b..16bf6ca 100644 --- a/crates/zesdex-backend/src/app/runtime/actions/handlers.rs +++ b/crates/zesdex-backend/src/app/runtime/actions/handlers.rs @@ -2,6 +2,8 @@ //! `apply_action` in the root module. Each handler mutates `AppStateRest` //! in place. +use crate::app::runtime::context::tokens::count_tokens; +use crate::app::runtime::context::window; use crate::app::state::rest::{AppStateRest, ChatMessageDisplay}; use crate::app::state::runtime::TurnEvent; use crate::app::state::types::{Overlay, Toast, ToastKind}; @@ -205,24 +207,15 @@ pub(super) fn handle_abort_turn(state: &mut AppStateRest) { } pub(super) fn handle_compact(state: &mut AppStateRest) { - let max_wire_tokens = state - .app_config - .model_roles - .values() - .find(|role| { - role.provider == state.settings.provider && role.model == state.settings.model - }) - .and_then(|role| role.context_window) - .unwrap_or(state.app_config.default_context_window) as usize; + let max_wire_tokens = window::resolve(&state.app_config, &state.settings); if let Some(ref mut rt) = state.session_runtime { - let total_chars: usize = rt + let token_estimate: usize = rt .messages .iter() .filter_map(|m| m.content.as_deref()) - .map(str::len) + .map(count_tokens) .sum(); - let token_estimate = total_chars / 3; rt.messages = crate::app::runtime::context::shaping::shape_messages( &rt.messages, @@ -230,6 +223,7 @@ pub(super) fn handle_compact(state: &mut AppStateRest) { max_wire_tokens, true, None, + None, ); state.push_toast(Toast::new( ToastKind::Success, diff --git a/crates/zesdex-backend/src/app/runtime/actions/turn.rs b/crates/zesdex-backend/src/app/runtime/actions/turn.rs index fd30c6e..a3216dd 100644 --- a/crates/zesdex-backend/src/app/runtime/actions/turn.rs +++ b/crates/zesdex-backend/src/app/runtime/actions/turn.rs @@ -13,6 +13,7 @@ use sha2::Digest; use zesdex_cms::domain::repository::EditLogRepository; use crate::app::guard::Verdict; +use crate::app::runtime::context::tokens::count_tokens; use crate::app::state::runtime::TurnEvent; use zesdex_cms::domain::repository::MemoryRepository; use crate::dto::chat::message::ChatMessage; @@ -326,12 +327,11 @@ pub(super) fn run_agent_turn( let mut todo_retry_count = 0usize; loop { - let total_chars: usize = msgs + let token_estimate: usize = msgs .iter() .filter_map(|m| m.content.as_deref()) - .map(str::len) + .map(count_tokens) .sum(); - let token_estimate = total_chars / 4; let max_wire_tokens = tc.context_window; // Skip message compaction if abort was requested — the non-streaming @@ -352,6 +352,7 @@ pub(super) fn run_agent_turn( max_wire_tokens, false, Some(&tc.client), + Some(&tc.abort_flag), ); // Dispatch the compacted messages to the main thread so the local session history @@ -488,12 +489,12 @@ pub(super) fn run_agent_turn( let (mut tok_in, mut tok_out) = final_usage.unwrap_or((0, 0)); if tok_in == 0 { - let total_chars: usize = wire_msgs + let total_tokens: usize = wire_msgs .iter() .filter_map(|m| m.content.as_deref()) - .map(str::len) + .map(count_tokens) .sum(); - tok_in = (total_chars / 4).max(1) as u64; + tok_in = total_tokens.max(1) as u64; } if tok_out == 0 { let response_chars = response.content.as_deref().map_or(0, str::len); diff --git a/crates/zesdex-backend/src/app/runtime/context/shaping.rs b/crates/zesdex-backend/src/app/runtime/context/shaping.rs index a6bbff4..205ad42 100644 --- a/crates/zesdex-backend/src/app/runtime/context/shaping.rs +++ b/crates/zesdex-backend/src/app/runtime/context/shaping.rs @@ -3,6 +3,8 @@ //! 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; @@ -37,6 +39,11 @@ pub fn should_shape(token_estimate: usize, max_wire_tokens: usize, prev_shaped: /// `[prior conversation compacted]` (or LLM-generated summary, if /// `client` is `Some`) system message in between. /// +/// 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. +/// /// Why: keeps context-size overhead roughly constant regardless of /// session length. /// @@ -48,6 +55,7 @@ pub fn shape_messages( max_wire_tokens: usize, force: bool, client: Option<&crate::service::provider::LlmClient>, + abort_flag: Option<&AtomicBool>, ) -> Vec { if !force && (token_count <= max_wire_tokens || messages.len() < 5) { return messages.to_vec(); @@ -88,30 +96,63 @@ pub fn shape_messages( 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") - ); + // 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 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]"); + 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: {}. \ + Falling back to static placeholder.", + e, + ); } - } - 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, - ); } } } @@ -150,7 +191,7 @@ mod tests { ChatMessage::user("hi"), ChatMessage::assistant(Some("hello".to_string())), ]; - let result = shape_messages(&messages, 10, 1000, false, None); + let result = shape_messages(&messages, 10, 1000, false, None, None); assert_eq!(result.len(), messages.len()); } @@ -178,7 +219,7 @@ mod tests { for i in 0..20 { messages.push(ChatMessage::user(padded_message(i))); } - let result = shape_messages(&messages, 100_000, 1000, true, None); + let result = shape_messages(&messages, 100_000, 1000, true, None, None); assert_eq!(result[0].content.as_deref(), Some("system prompt")); } @@ -188,7 +229,7 @@ mod tests { for i in 0..20 { messages.push(ChatMessage::user(padded_message(i))); } - let result = shape_messages(&messages, 100_000, 1000, true, None); + 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]")); @@ -201,11 +242,11 @@ mod tests { for i in 0..20 { messages.push(ChatMessage::user(padded_message(i))); } - let result = shape_messages(&messages, 100_000, 1000, true, None); + 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" ); } -} +} \ No newline at end of file