diff --git a/assets/config-template.yaml b/assets/config-template.yaml index d66c694..62a83a5 100644 --- a/assets/config-template.yaml +++ b/assets/config-template.yaml @@ -186,6 +186,7 @@ save_session: null # Controls the persistence of the session. If t compression_threshold: 4000 # Compress the session when the token count reaches or exceeds this threshold compression_keep_last: 0 # Number of most-recent messages to keep visible after compression (0 = compress all messages) compression_model: null # Model to use for session compression (e.g. 'openai:gpt-4o-mini'); defaults to the session's current model +compression_safety_valve: 0.8 # Force compression when the last real request's prompt size (or the local estimate) reaches this fraction of the model's max_input_tokens, even if compression_threshold would not fire (0 = disabled) summarization_prompt: null # The text prompt used for creating a concise summary of session messages. If null, uses built-in default summary_context_prompt: null # The text prompt used for including the summary of the entire session as context to the model. If null, uses built-in default max_tool_result_chars: null # Cap on tool result characters forwarded to the model per call (null = no cap) diff --git a/src/client/common.rs b/src/client/common.rs index 3ab7585..888a45c 100644 --- a/src/client/common.rs +++ b/src/client/common.rs @@ -420,6 +420,15 @@ impl TokenUsage { && self.cache_read_input_tokens.is_none() } + /// Prompt-side tokens the API reported for one request: fresh input plus + /// cache writes and cache reads. The buckets are disjoint (see + /// [`Self::cost_usd`]), so the terms sum without adjustment. + pub fn total_prompt_tokens(&self) -> u64 { + self.input_tokens.unwrap_or(0) + + self.cache_creation_input_tokens.unwrap_or(0) + + self.cache_read_input_tokens.unwrap_or(0) + } + pub fn accumulate(&mut self, other: &TokenUsage) { fn add(acc: &mut Option, value: Option) { if let Some(value) = value { diff --git a/src/client/openai.rs b/src/client/openai.rs index bf78a4a..7439be6 100644 --- a/src/client/openai.rs +++ b/src/client/openai.rs @@ -869,17 +869,22 @@ fn reasoning_thinking_block(item: &Value) -> Option { } /// OpenAI reports cached input under `input_tokens_details` and has no -/// cache-creation concept, so that field stays `None`. +/// cache-creation concept, so that field stays `None`. The reported +/// `input_tokens` total already includes cached reads, so cached tokens are +/// subtracted here to keep the buckets disjoint like other providers. fn openai_parse_responses_usage(usage: &Value) -> Option { if !usage.is_object() { return None; } + let cached_tokens = usage["input_tokens_details"]["cached_tokens"].as_u64(); Some(TokenUsage { - input_tokens: usage["input_tokens"].as_u64(), + input_tokens: usage["input_tokens"] + .as_u64() + .map(|v| v.saturating_sub(cached_tokens.unwrap_or(0))), output_tokens: usage["output_tokens"].as_u64(), cache_creation_input_tokens: None, - cache_read_input_tokens: usage["input_tokens_details"]["cached_tokens"].as_u64(), + cache_read_input_tokens: cached_tokens, }) } @@ -1267,12 +1272,31 @@ mod tests { let output = openai_extract_responses(&data).unwrap(); let usage = output.usage.unwrap(); - assert_eq!(usage.input_tokens, Some(100)); + assert_eq!(usage.input_tokens, Some(40)); assert_eq!(usage.output_tokens, Some(20)); assert_eq!(usage.cache_creation_input_tokens, None); assert_eq!(usage.cache_read_input_tokens, Some(60)); } + #[test] + fn extract_responses_usage_without_cached_details_keeps_input_intact() { + let data = json!({ + "output": [ + { "type": "message", "content": [{ "type": "output_text", "text": "answer" }] }, + ], + "usage": { + "input_tokens": 100, + "output_tokens": 20, + } + }); + + let output = openai_extract_responses(&data).unwrap(); + + let usage = output.usage.unwrap(); + assert_eq!(usage.input_tokens, Some(100)); + assert_eq!(usage.cache_read_input_tokens, None); + } + #[test] fn extract_responses_without_usage_leaves_it_none() { let data = json!({ @@ -1357,7 +1381,7 @@ mod tests { assert!(done); let (_, _, _, usage) = handler.take(); let usage = usage.unwrap(); - assert_eq!(usage.input_tokens, Some(100)); + assert_eq!(usage.input_tokens, Some(40)); assert_eq!(usage.output_tokens, Some(20)); assert_eq!(usage.cache_creation_input_tokens, None); assert_eq!(usage.cache_read_input_tokens, Some(60)); @@ -1467,7 +1491,7 @@ mod tests { let (text, _, _, usage) = handler.take(); assert_eq!(text, "partial answer"); let usage = usage.expect("usage should be recorded for truncated turns"); - assert_eq!(usage.input_tokens, Some(50)); + assert_eq!(usage.input_tokens, Some(25)); assert_eq!(usage.output_tokens, Some(10)); assert_eq!(usage.cache_read_input_tokens, Some(25)); } diff --git a/src/config/app_config.rs b/src/config/app_config.rs index 6fe6ddd..368d418 100644 --- a/src/config/app_config.rs +++ b/src/config/app_config.rs @@ -71,6 +71,7 @@ pub struct AppConfig { pub compression_threshold: usize, pub compression_keep_last: usize, pub compression_model: Option, + pub compression_safety_valve: f32, pub summarization_prompt: Option, pub summary_context_prompt: Option, pub max_tool_result_chars: Option, @@ -164,6 +165,7 @@ impl Default for AppConfig { compression_threshold: 4000, compression_keep_last: 0, compression_model: None, + compression_safety_valve: 0.8, summarization_prompt: None, summary_context_prompt: None, max_tool_result_chars: None, @@ -270,6 +272,7 @@ impl AppConfig { compression_threshold: config.compression_threshold, compression_keep_last: config.compression_keep_last, compression_model: config.compression_model, + compression_safety_valve: config.compression_safety_valve, summarization_prompt: config.summarization_prompt, summary_context_prompt: config.summary_context_prompt, max_tool_result_chars: config.max_tool_result_chars, diff --git a/src/config/input.rs b/src/config/input.rs index eb1fedc..fafc03d 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -68,6 +68,34 @@ impl Input { }) } + /// An input carrying only `text`: no session history, no RAG, and no + /// function declarations. For auxiliary LLM calls (such as session + /// summarization) whose requests must not grow with the session. + pub fn detached_text(ctx: &RequestContext, text: &str) -> Self { + let mut role = Role::new("", ""); + role.set_model(ctx.current_model().clone()); + Self { + app_config: Arc::clone(&ctx.app.config), + stream_enabled: ctx.app.config.stream, + session: None, + rag: None, + functions: None, + text: text.to_string(), + raw: (text.to_string(), vec![]), + patched_text: None, + last_reply: None, + continue_output: None, + regenerate: false, + medias: Default::default(), + data_urls: Default::default(), + tool_calls: None, + role, + rag_name: None, + with_session: false, + with_agent: false, + } + } + pub async fn from_files( ctx: &RequestContext, raw_text: &str, @@ -329,6 +357,14 @@ impl Input { } } + /// Re-captures `ctx.session` in place of the snapshot taken at + /// construction. Session compression mutates the ctx's session, not the + /// snapshot, so an input that outlives a compression must refresh before + /// building its next request or it re-sends the pre-compression history. + pub fn refresh_session(&mut self, ctx: &RequestContext) { + self.session = ctx.session.clone(); + } + pub fn session_mut<'a>(&self, session: &'a mut Option) -> Option<&'a mut Session> { if self.with_session { session.as_mut() @@ -660,6 +696,60 @@ mod tests { assert!(!with_session); } + #[test] + fn detached_text_input_carries_no_session_history_or_functions() { + let mut ctx = create_test_ctx(); + ctx.session = Some(Session::default()); + let earlier = Input::from_str(&ctx, "earlier question", None).unwrap(); + ctx.session + .as_mut() + .unwrap() + .add_message(&earlier, "earlier answer") + .unwrap(); + + let detached = Input::detached_text(&ctx, "summarize the above"); + let messages = detached.build_messages().unwrap(); + assert_eq!(messages.len(), 1, "must contain only the detached text"); + assert!(messages[0].role.is_user()); + assert_eq!(messages[0].content.to_text(), "summarize the above"); + assert!(detached.declared_function_names().is_empty()); + assert!(!detached.with_session()); + assert!(!detached.with_agent()); + } + + #[test] + fn refresh_session_rebuilds_from_compressed_session() { + let mut ctx = create_test_ctx(); + ctx.session = Some(Session::default()); + for i in 0..4 { + let question = + Input::from_str(&ctx, &format!("question {i}: {}", "x".repeat(500)), None).unwrap(); + ctx.session + .as_mut() + .unwrap() + .add_message(&question, &"y".repeat(500)) + .unwrap(); + } + + let mut input = Input::from_str(&ctx, "retry me", None).unwrap(); + let stale = serde_json::to_string(&input.build_messages().unwrap()).unwrap(); + + ctx.session.as_mut().unwrap().compress("recap".into(), 1); + + assert_eq!( + serde_json::to_string(&input.build_messages().unwrap()).unwrap(), + stale, + "the snapshot captured at construction must be untouched by compression" + ); + + input.refresh_session(&ctx); + let refreshed = serde_json::to_string(&input.build_messages().unwrap()).unwrap(); + assert!( + refreshed.len() < stale.len(), + "a refreshed input must build a strictly smaller request after compression" + ); + } + #[test] fn resolve_paths_detects_last_reply_syntax() { let loaders = HashMap::new(); diff --git a/src/config/mod.rs b/src/config/mod.rs index 0491e4e..5f3511d 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -286,6 +286,7 @@ pub struct Config { pub compression_threshold: usize, pub compression_keep_last: usize, pub compression_model: Option, + pub compression_safety_valve: f32, pub summarization_prompt: Option, pub summary_context_prompt: Option, pub max_tool_result_chars: Option, @@ -376,6 +377,7 @@ impl Default for Config { compression_threshold: 4000, compression_keep_last: 0, compression_model: None, + compression_safety_valve: 0.8, summarization_prompt: None, summary_context_prompt: None, max_tool_result_chars: None, @@ -1379,6 +1381,7 @@ clients: assert!(cfg.mcp_server_support); assert_eq!(cfg.compression_threshold, 4000); assert!(cfg.compression_model.is_none()); + assert_eq!(cfg.compression_safety_valve, 0.8); assert_eq!(cfg.rag_top_k, 5); assert!(cfg.save_shell_history); assert_eq!(cfg.keybindings, "emacs"); @@ -1419,6 +1422,7 @@ clients: assert!(cfg.mapping_tools.is_empty()); assert!(cfg.document_loaders.is_empty()); assert_eq!(cfg.compression_threshold, 4000); + assert_eq!(cfg.compression_safety_valve, 0.8); assert!(cfg.theme.is_none()); assert_eq!(cfg.clients.len(), 1); } diff --git a/src/config/request_context.rs b/src/config/request_context.rs index d974b16..57b727d 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -16,7 +16,9 @@ use super::{ list_agents_with_descriptions, memory, paths, }; use super::{MessageContentToolCalls, prompts}; -use crate::client::{Model, ModelType, TokenUsage, list_models}; +use crate::client::{ + Message, MessageContent, MessageRole, Model, ModelType, TokenUsage, list_models, +}; use crate::function::{ FunctionDeclaration, Functions, ToolCallTracker, ToolResult, agents::AGENT_FUNCTION_PREFIX, @@ -51,7 +53,7 @@ use super::memory::{ }; use crate::graph; use crate::hooks::{self, HookEvent}; -use anyhow::{Context, Error, Result, bail}; +use anyhow::{Context, Error, Result, anyhow, bail}; use colored::Colorize; use gman::providers::SupportedProvider; use hooks::{McpServerHooks, RagSyncHooks}; @@ -68,7 +70,7 @@ use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; -use std::{env, fs}; +use std::{env, fs, slice}; pub(crate) fn expand_enabled_mcp_server_ids( app: &AppConfig, @@ -316,6 +318,13 @@ pub struct RequestContext { pub last_token_usage: Option, pub last_cost: Option, + /// Prompt-side tokens from the most recent request that reported usage. + /// Unlike `last_token_usage`, this is never clobbered by requests that + /// report nothing (e.g. a stream rejected before `message_start`), so the + /// compression safety valve always sees the last real prompt + /// measurement. + pub last_prompt_token_usage: Option, + pub tool_scope: ToolScope, pub declared_function_names: HashSet, @@ -385,6 +394,7 @@ impl RequestContext { last_message: None, last_token_usage: None, last_cost: None, + last_prompt_token_usage: None, tool_scope: ToolScope::default(), declared_function_names: Default::default(), node_job_scope: None, @@ -457,6 +467,7 @@ impl RequestContext { last_message: None, last_token_usage: None, last_cost: None, + last_prompt_token_usage: None, tool_scope: ToolScope { functions, mcp_runtime, @@ -522,6 +533,7 @@ impl RequestContext { last_message: self.last_message.clone(), last_token_usage: self.last_token_usage.clone(), last_cost: self.last_cost, + last_prompt_token_usage: self.last_prompt_token_usage, tool_scope: self.tool_scope.clone(), declared_function_names: self.declared_function_names.clone(), node_job_scope: None, @@ -577,6 +589,7 @@ impl RequestContext { last_message: None, last_token_usage: None, last_cost: None, + last_prompt_token_usage: None, tool_scope: ToolScope { functions: Functions::default(), mcp_runtime: McpRuntime::default(), @@ -2294,6 +2307,14 @@ impl RequestContext { session.accumulate_cost(cost); } } + // Partial usage from an aborted stream lands here on purpose: + // it is a genuine API-reported prompt measurement, not noise. + // Only usage-less recordings (failed requests) leave the + // previous measurement in place. + let prompt_tokens = usage.total_prompt_tokens(); + if prompt_tokens > 0 { + self.last_prompt_token_usage = Some(prompt_tokens as usize); + } } self.last_token_usage = usage; self.last_cost = cost; @@ -5030,7 +5051,7 @@ impl RequestContext { } else { role_name = self.role.as_ref().map(|v| v.name().to_string()); } - let name = role_name.ok_or_else(|| anyhow::anyhow!("No role"))?; + let name = role_name.ok_or_else(|| anyhow!("No role"))?; self.upsert_role(app, &name)?; self.use_role(app, &name, abort_signal).await } @@ -5243,14 +5264,31 @@ impl RequestContext { } pub async fn compress_session(&mut self) -> Result<()> { - match self.session.as_ref() { + async fn timed_fetch(input: Input) -> Result { + tokio::time::timeout(Duration::from_secs(120), input.fetch_chat_text()) + .await + .map_err(|_| anyhow!("Compression LLM call timed out after 120 s"))? + } + + self.compress_session_impl(timed_fetch).await + } + + /// The summarization flow with the LLM call injectable, so tests can + /// capture and answer each chunk request without a network round trip. + async fn compress_session_impl(&mut self, fetch: F) -> Result<()> + where + F: Fn(Input) -> Fut, + Fut: Future>, + { + let session_window = match self.session.as_ref() { Some(session) => { if !session.has_user_messages() { bail!("No need to compress since there are no messages in the session") } + session.model().max_input_tokens() } None => bail!("No session"), - } + }; let prompt = self .app @@ -5259,44 +5297,79 @@ impl RequestContext { .clone() .unwrap_or_else(|| SUMMARIZATION_PROMPT.into()); - async fn timed_fetch(input: &Input) -> Result { - tokio::time::timeout(Duration::from_secs(120), input.fetch_chat_text()) - .await - .map_err(|_| anyhow::anyhow!("Compression LLM call timed out after 120 s"))? - } - - let mut input = Input::from_str(self, &prompt, None)?; - let compression_model_id = self + let compression_model = match self .compression_model() - .filter(|id| *id != self.current_model().id()); - let summary = match compression_model_id { + .filter(|id| *id != self.current_model().id()) + { Some(model_id) => { match Model::retrieve_model(self.app.config.as_ref(), &model_id, ModelType::Chat) { - Ok(model) => { - input.set_role_model(model); - match timed_fetch(&input).await { - Ok(summary) => summary, - Err(err) => { - eprintln!( - "Warning: compression model '{model_id}' failed: {err}; falling back to the current model '{}' for session compression", - self.current_model().id() - ); - let input = Input::from_str(self, &prompt, None)?; - timed_fetch(&input).await? - } - } - } + Ok(model) => Some(model), Err(err) => { eprintln!( "Warning: compression model '{model_id}' could not be used ({err}); falling back to the current model '{}' for session compression", self.current_model().id() ); - timed_fetch(&input).await? + None } } } - None => timed_fetch(&input).await?, + None => None, }; + + // Cap each chunk well below the summarizer's window so every request + // (prior summary + chunk + prompt) fits, even for sessions that have + // grown past the window itself. + let window = compression_model + .as_ref() + .and_then(|model| model.max_input_tokens()) + .or(session_window) + .unwrap_or(SUMMARIZATION_WINDOW_FALLBACK_TOKENS); + let chunk_budget = (window as f32 * SUMMARIZATION_CHUNK_BUDGET_RATIO) as usize; + + let keep_last = self.compression_keep_last(); + let chunks: Vec = self + .session + .as_ref() + .map(|session| { + // `Session::compress` archives every message when `keep_last` + // covers the whole list, so summarize them all in that case. + let mut foldable = session.foldable_messages(keep_last); + if foldable.is_empty() { + foldable = session.foldable_messages(0); + } + slice_summarization_chunks(session.model(), foldable, chunk_budget) + .into_iter() + .map(render_summarization_chunk) + .collect() + }) + .unwrap_or_default(); + + // Fold forward, oldest chunk first: each call carries the summary of + // everything before it, and the summarization prompt instructs the + // model to merge that prior summary into the new one. + let mut summary = String::new(); + for chunk in &chunks { + let request = compose_summarization_request(&prompt, &summary, chunk); + summary = match compression_model.as_ref() { + Some(model) => { + let mut input = Input::detached_text(self, &request); + input.set_role_model(model.clone()); + match fetch(input).await { + Ok(summary) => summary, + Err(err) => { + eprintln!( + "Warning: compression model '{}' failed: {err}; falling back to the current model '{}' for session compression", + model.id(), + self.current_model().id() + ); + fetch(Input::detached_text(self, &request)).await? + } + } + } + None => fetch(Input::detached_text(self, &request)).await?, + }; + } + let summary_context_prompt = self .app .config @@ -5541,6 +5614,73 @@ fn fork_base_name(name: &str) -> &str { name } +/// Assumed summarizer window when neither the compression model nor the +/// session model declares `max_input_tokens`. +const SUMMARIZATION_WINDOW_FALLBACK_TOKENS: usize = 200_000; +/// Share of the summarizer's window a single history chunk may occupy; the +/// rest is headroom for the summarization prompt, the folded-forward prior +/// summary, and the response. +const SUMMARIZATION_CHUNK_BUDGET_RATIO: f32 = 0.6; + +/// Splits `messages` into consecutive runs whose estimated tokens each stay +/// within `budget`. A message that alone exceeds the budget forms its own +/// chunk: messages are the smallest unit summarization can fold. +fn slice_summarization_chunks<'a>( + model: &Model, + messages: &'a [Message], + budget: usize, +) -> Vec<&'a [Message]> { + let mut chunks = vec![]; + let mut start = 0; + let mut tokens = 0; + for (i, message) in messages.iter().enumerate() { + let message_tokens = model.total_tokens(slice::from_ref(message)); + if i > start && tokens + message_tokens > budget { + chunks.push(&messages[start..i]); + start = i; + tokens = 0; + } + tokens += message_tokens; + } + if start < messages.len() { + chunks.push(&messages[start..]); + } + chunks +} + +/// Renders messages as a role-prefixed transcript for the summarizer. Tool +/// transcripts are serialized whole rather than dropped: their results often +/// carry the facts the summary must preserve. +fn render_summarization_chunk(messages: &[Message]) -> String { + messages + .iter() + .map(|message| { + let role = match message.role { + MessageRole::System => "SYSTEM", + MessageRole::Assistant => "ASSISTANT", + MessageRole::User => "USER", + MessageRole::Tool => "TOOL", + }; + let content = match &message.content { + MessageContent::ToolCalls(tool_calls) => { + serde_json::to_string(tool_calls).unwrap_or_else(|_| tool_calls.text.clone()) + } + content => content.to_text(), + }; + format!("{role}: {content}") + }) + .collect::>() + .join("\n\n") +} + +fn compose_summarization_request(prompt: &str, prior_summary: &str, chunk: &str) -> String { + if prior_summary.is_empty() { + format!("{chunk}\n\n{prompt}") + } else { + format!("{SUMMARY_CONTEXT_PROMPT}{prior_summary}\n\n{chunk}\n\n{prompt}") + } +} + #[cfg(test)] mod tests { use super::super::mcp_factory::McpFactory; @@ -5755,6 +5895,30 @@ mod tests { assert!(!vars.contains_key("total_cost")); } + #[test] + fn record_token_usage_keeps_last_prompt_tokens_across_failed_requests() { + let mut ctx = create_test_ctx(); + ctx.record_token_usage( + Some(TokenUsage { + input_tokens: Some(100), + output_tokens: Some(5), + cache_creation_input_tokens: Some(40), + cache_read_input_tokens: Some(60), + }), + &priced_model(), + ); + + assert_eq!(ctx.last_prompt_token_usage, Some(200)); + + // A request that fails before the provider reports usage records + // `None`: the display field is clobbered by design, but the safety + // valve's measurement must survive. + ctx.record_token_usage(None, &priced_model()); + + assert!(ctx.last_token_usage.is_none()); + assert_eq!(ctx.last_prompt_token_usage, Some(200)); + } + #[test] fn generate_prompt_context_omits_costs_that_render_as_zero() { let mut ctx = create_test_ctx(); @@ -6028,6 +6192,184 @@ mod tests { ); } + fn windowed_model(max_input_tokens: usize) -> Model { + let mut data = ModelData::new("test-window"); + data.max_input_tokens = Some(max_input_tokens); + Model::from_config("openai", &[data]).remove(0) + } + + fn text_message(role: MessageRole, text: &str) -> Message { + Message::new(role, MessageContent::Text(text.to_string())) + } + + #[test] + fn summarization_chunks_respect_budget_and_preserve_order() { + let model = Model::default(); + // 1000 ascii chars => ~250 estimated tokens per message. + let messages: Vec = (0..10) + .map(|i| { + let role = if i % 2 == 0 { + MessageRole::User + } else { + MessageRole::Assistant + }; + text_message(role, &format!("msg-{i:02} {}", "x".repeat(993))) + }) + .collect(); + + let budget = 600; + let chunks = slice_summarization_chunks(&model, &messages, budget); + + assert!(chunks.len() > 1, "a history over budget must be split"); + for chunk in &chunks { + assert!( + model.total_tokens(chunk) <= budget, + "each chunk must fit the budget" + ); + } + let flattened: Vec<&Message> = chunks.iter().flat_map(|chunk| chunk.iter()).collect(); + assert_eq!( + flattened.len(), + messages.len(), + "chunks must cover every message" + ); + for (original, chunked) in messages.iter().zip(flattened) { + assert_eq!(original.content.to_text(), chunked.content.to_text()); + } + + assert_eq!( + slice_summarization_chunks(&model, &messages, 10_000).len(), + 1, + "a history within budget must stay a single chunk" + ); + } + + #[test] + fn summarization_chunk_slicing_isolates_oversized_message() { + let model = Model::default(); + let messages = vec![ + text_message(MessageRole::User, "small question"), + text_message(MessageRole::Assistant, &"y".repeat(8000)), + text_message(MessageRole::User, "another small question"), + ]; + let chunks = slice_summarization_chunks(&model, &messages, 100); + assert_eq!(chunks.len(), 3); + assert_eq!( + chunks[1].len(), + 1, + "an over-budget message forms its own chunk" + ); + } + + #[test] + fn render_summarization_chunk_keeps_tool_transcripts() { + let tool_calls = MessageContentToolCalls::new( + vec![ToolResult::new( + ToolCall::new( + "fs_cat".to_string(), + json!({"path": "notes.txt"}), + Some("call-1".to_string()), + ), + json!("the file contents"), + )], + "reading the file".to_string(), + ); + let messages = vec![ + text_message(MessageRole::User, "read notes.txt"), + Message::new(MessageRole::Tool, MessageContent::ToolCalls(tool_calls)), + ]; + + let rendered = render_summarization_chunk(&messages); + + assert!(rendered.contains("USER: read notes.txt")); + assert!(rendered.contains("TOOL: ")); + assert!(rendered.contains("fs_cat")); + assert!(rendered.contains("the file contents")); + } + + #[test] + fn compress_session_folds_forward_across_chunks() { + let mut app = AppState::test_default(); + app.config = Arc::new(AppConfig { + compression_keep_last: 2, + ..AppConfig::default() + }); + let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Cmd); + let mut session = Session::default(); + // A 2_000-token window gives a 1_200-token chunk budget; the session + // built below estimates past the window itself. + session.set_model(windowed_model(2_000)); + ctx.session = Some(session); + for i in 0..10 { + let text = format!("msg-{i:02} {}", "x".repeat(793)); + let input = Input::from_str(&ctx, &text, Some(Role::new("", ""))).unwrap(); + ctx.session + .as_mut() + .unwrap() + .add_message(&input, &format!("reply-{i:02}")) + .unwrap(); + } + assert!( + ctx.session.as_ref().unwrap().tokens() > 2_000, + "the session must not fit the summarizer window in one request" + ); + + let requests = std::cell::RefCell::new(Vec::new()); + let requests_ref = &requests; + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap() + .block_on(ctx.compress_session_impl(move |input: Input| async move { + // The real request pipeline, guard_max_input_tokens included. + let data = input.prepare_completion_data(input.role().model(), false)?; + requests_ref + .borrow_mut() + .push((input.text(), data.messages, data.functions)); + Ok(format!("folded-summary-{}", requests_ref.borrow().len())) + })) + .unwrap(); + + let requests = requests.into_inner(); + assert_eq!(requests.len(), 2, "the foldable middle needs two chunks"); + + let (first_text, first_messages, first_functions) = &requests[0]; + assert!(first_text.contains("msg-00")); + assert!(!first_text.contains(SUMMARY_CONTEXT_PROMPT)); + assert_eq!(first_messages.len(), 1, "detached: no session history"); + assert!(first_messages[0].role.is_user()); + assert!( + first_functions.is_none(), + "detached: no function declarations" + ); + + let (second_text, second_messages, second_functions) = &requests[1]; + assert!(second_text.contains(SUMMARY_CONTEXT_PROMPT)); + assert!(second_text.contains("folded-summary-1")); + assert!(second_text.contains("msg-05")); + assert!( + !second_text.contains("msg-00"), + "the first chunk reaches the second call only through its summary" + ); + assert!( + !second_text.contains("msg-09"), + "the keep_last tail is never summarized" + ); + assert_eq!(second_messages.len(), 1, "detached: no session history"); + assert!(second_functions.is_none()); + + let session = ctx.session.as_ref().unwrap(); + assert_eq!(session.messages().len(), 3, "summary + keep_last tail"); + assert!(session.messages()[0].role.is_system()); + assert!( + session.messages()[0] + .content + .to_text() + .contains("folded-summary-2"), + "the recap must carry the last folded summary" + ); + } + #[test] fn should_inject_skill_instructions_requires_function_calling() { let app = AppConfig { diff --git a/src/config/session.rs b/src/config/session.rs index abef564..784104a 100644 --- a/src/config/session.rs +++ b/src/config/session.rs @@ -719,17 +719,50 @@ impl Session { self.reclaimable_tokens(keep_last) >= (threshold / COMPRESSION_MIN_RECLAIM_DIVISOR).max(1) } - /// Estimated tokens of the messages that compression would actually fold - /// away: everything except the leading system message (whose content is - /// carried into the summary message) and the `keep_last` tail (kept verbatim). - fn reclaimable_tokens(&self, keep_last: usize) -> usize { + /// Whether the prompt is close enough to the model's hard input limit + /// that compression must run regardless of how `compression_threshold` + /// is configured. Sizes the prompt as the larger of the API-reported + /// measurement from the last request that reported usage and the local + /// estimate, because the estimate can run well below the provider's + /// real count. + pub fn safety_valve_triggered( + &self, + valve: f32, + keep_last: usize, + last_prompt_tokens: Option, + ) -> bool { + if self.compressing || !valve.is_finite() || valve <= 0.0 { + return false; + } + let Some(max_input_tokens) = self.model().max_input_tokens() else { + return false; + }; + // With nothing to fold away, compressing is a no-op; firing anyway + // would spin compress attempts while a single in-flight turn holds + // all the tokens. + if self.reclaimable_tokens(keep_last) == 0 { + return false; + } + last_prompt_tokens.unwrap_or(0).max(self.tokens()) + >= (valve * max_input_tokens as f32) as usize + } + + /// The messages that compression would actually fold away: everything + /// except the leading system message (whose content is carried into the + /// summary message) and the `keep_last` tail (kept verbatim). + pub fn foldable_messages(&self, keep_last: usize) -> &[Message] { let start = usize::from( self.messages .first() .is_some_and(|v| v.role == MessageRole::System), ); let end = self.messages.len().saturating_sub(keep_last).max(start); - self.model().total_tokens(&self.messages[start..end]) + &self.messages[start..end] + } + + /// Estimated tokens of the messages compression would fold away. + fn reclaimable_tokens(&self, keep_last: usize) -> usize { + self.model().total_tokens(self.foldable_messages(keep_last)) } pub fn compressing(&self) -> bool { @@ -1109,7 +1142,7 @@ impl AutoName { mod tests { use super::*; use crate::client::{ - Message, MessageContent, MessageContentToolCalls, MessageRole, Model, TokenUsage, + Message, MessageContent, MessageContentToolCalls, MessageRole, Model, ModelData, TokenUsage, }; use crate::config::{AppConfig, AppState, RequestContext, WorkingMode}; use crate::function::{Functions, ToolCall, ToolResult}; @@ -1578,6 +1611,85 @@ mod tests { assert!(session.needs_compression(100, 1)); } + /// A session with a system message, a large user message, and a small + /// assistant reply, on a model with the given input limit. With + /// `keep_last = 1` the user message is reclaimable; with `keep_last = 2` + /// nothing is. + fn valve_session(max_input_tokens: Option) -> Session { + let mut session = Session::default(); + let mut data = ModelData::new("valve-model"); + data.max_input_tokens = max_input_tokens; + session.set_model(Model::from_config("provider", &[data]).remove(0)); + session.messages.push(Message::new( + MessageRole::System, + MessageContent::Text("x".repeat(400)), + )); + session.messages.push(Message::new( + MessageRole::User, + MessageContent::Text("y".repeat(4000)), + )); + session.messages.push(Message::new( + MessageRole::Assistant, + MessageContent::Text("ok".to_string()), + )); + session.update_tokens(); + session + } + + #[test] + fn session_safety_valve_fires_from_estimate() { + let session = valve_session(Some(1000)); + + assert!(session.tokens() >= 800); + assert!(session.safety_valve_triggered(0.8, 1, None)); + } + + #[test] + fn session_safety_valve_fires_from_real_usage_over_estimate() { + let session = valve_session(Some(1_000_000)); + + assert!(!session.safety_valve_triggered(0.8, 1, None)); + assert!(!session.safety_valve_triggered(0.8, 1, Some(700_000))); + assert!(session.safety_valve_triggered(0.8, 1, Some(800_000))); + } + + #[test] + fn session_safety_valve_disabled_when_zero() { + let session = valve_session(Some(1000)); + + assert!(!session.safety_valve_triggered(0.0, 1, Some(1_000_000))); + } + + #[test] + fn session_safety_valve_disabled_when_nan() { + let session = valve_session(Some(1000)); + + assert!(!session.safety_valve_triggered(f32::NAN, 1, Some(1_000_000))); + } + + #[test] + fn session_safety_valve_respects_compressing() { + let mut session = valve_session(Some(1000)); + session.set_compressing(true); + + assert!(!session.safety_valve_triggered(0.8, 1, Some(1_000_000))); + } + + #[test] + fn session_safety_valve_requires_reclaimable_tokens() { + let session = valve_session(Some(1000)); + + assert!(!session.safety_valve_triggered(0.8, 2, Some(1_000_000))); + assert!(session.safety_valve_triggered(0.8, 1, Some(1_000_000))); + } + + #[test] + fn session_safety_valve_requires_model_max_input_tokens() { + let session = valve_session(None); + + assert!(!session.safety_valve_triggered(0.8, 1, Some(1_000_000))); + } + #[test] fn session_set_compressing_flag() { let mut session = Session::default(); diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 9751a5a..2807d58 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -1515,6 +1515,8 @@ async fn ask_inner( ) -> Result<()> { let mut input = input; let mut with_embeddings = with_embeddings; + // One compress-and-retry per fresh input; tool-loop iterations share it. + let mut retried_after_compress = false; loop { if input.is_empty() { return Ok(()); @@ -1530,6 +1532,10 @@ async fn ask_inner( let app = Arc::clone(&ctx.app.config); + if maybe_auto_compress(ctx, app.as_ref()).await.succeeded { + input.refresh_session(ctx); + } + if graph::active_agent_graph_name(ctx).is_some() { ctx.before_chat_completion(&input)?; let output = @@ -1560,6 +1566,18 @@ async fn ask_inner( match result { Ok(v) => v, Err(err) => { + if compress_for_error_retry( + ctx, + app.as_ref(), + retried_after_compress, + &abort_signal, + ) + .await + { + input.refresh_session(ctx); + retried_after_compress = true; + continue; + } ctx.on_chat_completion_error(app.as_ref(), &input); return Err(err); } @@ -1574,6 +1592,7 @@ async fn ask_inner( match check_pending_tasks_guardrail(ctx) { GuardrailAction::Inject(prompt) => { input = Input::from_str(ctx, &prompt, None)?; + retried_after_compress = false; continue; } GuardrailAction::ForceTerminate(ids) => { @@ -1617,6 +1636,7 @@ async fn ask_inner( format!("{prompt}\n\n{todo_state}") }; input = Input::from_str(ctx, &full_prompt, None)?; + retried_after_compress = false; continue; } @@ -1637,63 +1657,45 @@ async fn ask_inner( } } - let keep_last = ctx.compression_keep_last(); - let needs_compression = ctx - .session - .as_ref() - .is_some_and(|s| s.needs_compression(app.compression_threshold, keep_last)); + let agent_can_continue_after_compress = + should_continue(ctx) && !abort_signal.aborted_ctrlc(); - if needs_compression { - let agent_can_continue_after_compress = - should_continue(ctx) && !abort_signal.aborted_ctrlc(); + // Auto-continue keys off the compression CONDITION, not the outcome: + // a transiently failed compression retries at the next turn end, and + // gating on success would stall an unattended agent at the pause + // banner with incomplete todos. + let compression = maybe_auto_compress(ctx, app.as_ref()).await; + if compression.fired && agent_can_continue_after_compress { + let full_prompt = { + let config = ctx.auto_continue_config(); + let todo_state = ctx.todo_list.render_for_model(); + let remaining = ctx.todo_list.incomplete_count(); + ctx.increment_auto_continue_count(); + let count = ctx.auto_continue_count; + let max = config.max_continues; - if let Some(session) = ctx.session.as_mut() { - session.set_compressing(true); - } + let prompt = config + .continuation_prompt + .as_deref() + .unwrap_or(DEFAULT_CONTINUATION_PROMPT); - let color = if app.light_theme() { - nu_ansi_term::Color::LightGray - } else { - nu_ansi_term::Color::DarkGray - }; - eprintln!("\nšŸ“¢ {}", color.italic().paint("Compressing the session."),); - - auto_compress_session(ctx).await; - if let Some(session) = ctx.session.as_mut() { - session.set_compressing(false); - } - - if agent_can_continue_after_compress { - let full_prompt = { - let config = ctx.auto_continue_config(); - let todo_state = ctx.todo_list.render_for_model(); - let remaining = ctx.todo_list.incomplete_count(); - ctx.increment_auto_continue_count(); - let count = ctx.auto_continue_count; - let max = config.max_continues; - - let prompt = config - .continuation_prompt - .as_deref() - .unwrap_or(DEFAULT_CONTINUATION_PROMPT); - - let color = if app.light_theme() { - nu_ansi_term::Color::LightGray - } else { - nu_ansi_term::Color::DarkGray - }; - eprintln!( - "\nšŸ“‹ {}", - color.italic().paint(format!( - "Auto-continuing after compression ({count}/{max}): {remaining} incomplete todo(s) remain" - )) - ); - - format!("{prompt}\n\n{todo_state}") + let color = if app.light_theme() { + nu_ansi_term::Color::LightGray + } else { + nu_ansi_term::Color::DarkGray }; - input = Input::from_str(ctx, &full_prompt, None)?; - continue; - } + eprintln!( + "\nšŸ“‹ {}", + color.italic().paint(format!( + "Auto-continuing after compression ({count}/{max}): {remaining} incomplete todo(s) remain" + )) + ); + + format!("{prompt}\n\n{todo_state}") + }; + input = Input::from_str(ctx, &full_prompt, None)?; + retried_after_compress = false; + continue; } print_turn_divider(); @@ -1709,15 +1711,87 @@ fn print_turn_divider() { println!(); } -async fn auto_compress_session(ctx: &mut RequestContext) { +struct AutoCompressOutcome { + fired: bool, + succeeded: bool, +} + +async fn maybe_auto_compress(ctx: &mut RequestContext, app: &AppConfig) -> AutoCompressOutcome { + let keep_last = ctx.compression_keep_last(); + let should_compress = ctx.session.as_ref().is_some_and(|session| { + session.needs_compression(app.compression_threshold, keep_last) + || session.safety_valve_triggered( + app.compression_safety_valve, + keep_last, + ctx.last_prompt_token_usage, + ) + }); + if !should_compress { + return AutoCompressOutcome { + fired: false, + succeeded: false, + }; + } + + if let Some(session) = ctx.session.as_mut() { + session.set_compressing(true); + } + + let color = if app.light_theme() { + nu_ansi_term::Color::LightGray + } else { + nu_ansi_term::Color::DarkGray + }; + eprintln!("\nšŸ“¢ {}", color.italic().paint("Compressing the session."),); + + let compressed = auto_compress_session(ctx).await; + if let Some(session) = ctx.session.as_mut() { + session.set_compressing(false); + } + if compressed { + ctx.last_prompt_token_usage = None; + } + AutoCompressOutcome { + fired: true, + succeeded: compressed, + } +} + +async fn compress_for_error_retry( + ctx: &mut RequestContext, + app: &AppConfig, + already_retried: bool, + abort_signal: &AbortSignal, +) -> bool { + if already_retried || abort_signal.aborted() { + return false; + } + let keep_last = ctx.compression_keep_last(); + let valve_triggered = ctx.session.as_ref().is_some_and(|session| { + session.safety_valve_triggered( + app.compression_safety_valve, + keep_last, + ctx.last_prompt_token_usage, + ) + }); + valve_triggered && maybe_auto_compress(ctx, app).await.succeeded +} + +async fn auto_compress_session(ctx: &mut RequestContext) -> bool { match ctx.compress_session().await { - Ok(()) => hooks::fire( - HookEvent::SessionCompressed, - ctx, - &hooks::role_extras(ctx), - None, - ), - Err(err) => warn!("Failed to compress the session: {err}"), + Ok(()) => { + hooks::fire( + HookEvent::SessionCompressed, + ctx, + &hooks::role_extras(ctx), + None, + ); + true + } + Err(err) => { + warn!("Failed to compress the session: {err}"); + false + } } } @@ -2036,10 +2110,12 @@ pub fn split_args_text(line: &str, is_win: bool) -> (Vec, &str) { #[cfg(test)] mod tests { use super::*; - use crate::client::{ClientConfig, Model}; + use crate::client::{ClientConfig, MessageContent, Model, ModelData, ToolCall}; use crate::config::{AppState, Role, RoleLike, Session, TEMP_ROLE_NAME, WorkingMode}; + use crate::function::ToolResult; use crate::hooks::{HookDef, HooksMap, test_sink}; use anyhow::anyhow; + use serde_json::json; use serial_test::serial; use std::future::Future; @@ -2357,6 +2433,13 @@ mod tests { /// resolvable, and the session already holds a user exchange so /// `compress_session` has something to summarize. fn compressible_ctx(marker: &str) -> RequestContext { + compressible_ctx_with(marker, AppConfig::default()) + } + + /// `compressible_ctx` with caller-tuned compression settings, so tests + /// can drive `needs_compression` and the safety valve; `dry_run` and the + /// default client are enforced here to keep compression offline. + fn compressible_ctx_with(marker: &str, config: AppConfig) -> RequestContext { let mut hooks_map = HooksMap::default(); hooks_map.insert( "session.compressed".to_string(), @@ -2370,7 +2453,7 @@ mod tests { hooks: hooks_map, dry_run: true, clients: vec![ClientConfig::default()], - ..Default::default() + ..config }); let mut ctx = RequestContext::new(Arc::new(app), WorkingMode::Repl); let mut session = Session::default(); @@ -2382,6 +2465,7 @@ mod tests { .unwrap() .add_message(&input, "hi") .unwrap(); + ctx.session.as_mut().unwrap().update_tokens(); ctx } @@ -2432,7 +2516,7 @@ mod tests { let marker = "auto-compress-ok-f8r"; run_async(async { let mut ctx = compressible_ctx(marker); - auto_compress_session(&mut ctx).await; + assert!(auto_compress_session(&mut ctx).await); }); assert_eq!( @@ -2494,7 +2578,7 @@ mod tests { // An empty session makes compress_session fail; the seam must // swallow the error (no panic, no propagation) and fire nothing. ctx.session = Some(Session::default()); - auto_compress_session(&mut ctx).await; + assert!(!auto_compress_session(&mut ctx).await); }); assert_eq!( @@ -2504,6 +2588,427 @@ mod tests { ); } + #[test] + #[serial] + fn maybe_auto_compress_skips_when_neither_threshold_nor_valve_fires() { + let _sink = test_sink::install(); + let marker = "maybe-compress-skip-q7m"; + run_async(async { + let mut ctx = compressible_ctx_with( + marker, + AppConfig { + compression_keep_last: 1, + ..Default::default() + }, + ); + let app = Arc::clone(&ctx.app.config); + + assert!(!maybe_auto_compress(&mut ctx, app.as_ref()).await.fired); + }); + + assert_eq!(session_compressed_count(marker), 0); + } + + #[test] + #[serial] + fn maybe_auto_compress_fires_on_threshold_each_loop_iteration() { + // Exercises the shared seam the ask loop calls at the top of every + // iteration (tool-loop, guardrail-inject, and auto-continue turns + // included) as well as at turn end. + let _sink = test_sink::install(); + let marker = "maybe-compress-threshold-n2x"; + run_async(async { + let mut ctx = compressible_ctx_with( + marker, + AppConfig { + compression_threshold: 1, + compression_keep_last: 1, + ..Default::default() + }, + ); + let app = Arc::clone(&ctx.app.config); + + assert!(maybe_auto_compress(&mut ctx, app.as_ref()).await.succeeded); + assert!( + !ctx.is_compressing_session(), + "the compressing flag must be released after compression" + ); + }); + + assert_eq!(session_compressed_count(marker), 1); + } + + #[test] + #[serial] + fn maybe_auto_compress_reports_no_success_when_compression_fails() { + let _sink = test_sink::install(); + let marker = "maybe-compress-fail-g4z"; + run_async(async { + let mut ctx = compressible_ctx_with( + marker, + AppConfig { + compression_threshold: 1, + compression_keep_last: 0, + ..Default::default() + }, + ); + // A session whose only reclaimable message is non-user: the + // threshold fires, but `compress_session` bails on a session + // without user messages. + { + let input = Input::from_str(&ctx, "more context", None).unwrap(); + let session = ctx.session.as_mut().unwrap(); + session.add_message(&input, &"y".repeat(400)).unwrap(); + session.compress("recap".into(), 1); + } + let app = Arc::clone(&ctx.app.config); + + let outcome = maybe_auto_compress(&mut ctx, app.as_ref()).await; + assert!(outcome.fired, "the compression condition must fire"); + assert!( + !outcome.succeeded, + "a failed compression must not report success" + ); + assert!( + !ctx.is_compressing_session(), + "the compressing flag must be released after a failed compression" + ); + }); + + assert_eq!(session_compressed_count(marker), 0); + } + + #[test] + #[serial] + fn failed_compression_still_reports_fired_with_kept_tail() { + // Pins the turn-end auto-continue gate: with keep_last > 0, a + // compression whose condition fired but whose summarization FAILED + // must still report `fired`, so an unattended agent auto-continues + // and retries the compression at the next turn end instead of + // stalling at the pause banner. + let _sink = test_sink::install(); + let marker = "maybe-compress-fired-fail-t8r"; + run_async(async { + let mut ctx = compressible_ctx_with( + marker, + AppConfig { + compression_threshold: 1, + compression_keep_last: 1, + ..Default::default() + }, + ); + // A session whose foldable middle is a tool transcript and whose + // kept tail is an assistant reply: the threshold fires, but + // `compress_session` bails on a session without user messages. + { + let input = Input::from_str(&ctx, "more context", None) + .unwrap() + .merge_tool_results( + "calling a tool to gather context".into(), + vec![ToolResult::new(ToolCall::default(), json!("ok"))], + ); + let session = ctx.session.as_mut().unwrap(); + session.add_message(&input, &"y".repeat(400)).unwrap(); + session.compress("recap".into(), 2); + } + let app = Arc::clone(&ctx.app.config); + + let outcome = maybe_auto_compress(&mut ctx, app.as_ref()).await; + assert!(outcome.fired, "the compression condition must fire"); + assert!( + !outcome.succeeded, + "a failed compression must not report success" + ); + }); + + assert_eq!(session_compressed_count(marker), 0); + } + + #[test] + #[serial] + fn loop_top_compression_refresh_rebuilds_from_compressed_session() { + let _sink = test_sink::install(); + let marker = "loop-top-refresh-h5w"; + run_async(async { + let mut ctx = compressible_ctx_with( + marker, + AppConfig { + compression_threshold: 1, + compression_keep_last: 1, + ..Default::default() + }, + ); + let app = Arc::clone(&ctx.app.config); + let mut input = Input::from_str(&ctx, "next turn", None).unwrap(); + let stale = serde_json::to_string(&input.build_messages().unwrap()).unwrap(); + + assert!(maybe_auto_compress(&mut ctx, app.as_ref()).await.succeeded); + input.refresh_session(&ctx); + + let refreshed = serde_json::to_string(&input.build_messages().unwrap()).unwrap(); + assert_ne!( + refreshed, stale, + "the request built after a loop-top compression must reflect the compressed session" + ); + }); + + assert_eq!(session_compressed_count(marker), 1); + } + + #[test] + #[serial] + fn maybe_auto_compress_fires_on_safety_valve_despite_huge_threshold() { + let _sink = test_sink::install(); + let marker = "maybe-compress-valve-v5c"; + run_async(async { + let mut ctx = compressible_ctx_with( + marker, + AppConfig { + compression_threshold: 1_000_000, + compression_keep_last: 1, + compression_safety_valve: 0.8, + ..Default::default() + }, + ); + // A distinct name: `set_model` is a no-op for an unchanged model + // id, and the ctx's default session model declares no limit. + let mut data = ModelData::new("test-compress-model-limited"); + data.max_input_tokens = Some(100_000); + let model = Model::from_config("openai", &[data]).remove(0); + ctx.session.as_mut().unwrap().set_model(model); + // The last successful request measured a prompt near the limit + // even though the local estimate is tiny. + ctx.last_prompt_token_usage = Some(90_000); + let app = Arc::clone(&ctx.app.config); + + assert!(maybe_auto_compress(&mut ctx, app.as_ref()).await.succeeded); + }); + + assert_eq!(session_compressed_count(marker), 1); + } + + /// A ctx where only the safety valve can trigger compression: the last + /// successful request measured a prompt near the session model's limit + /// while the threshold is out of reach. + fn valve_tripped_ctx(marker: &str) -> RequestContext { + let mut ctx = compressible_ctx_with( + marker, + AppConfig { + compression_threshold: 1_000_000, + compression_keep_last: 1, + compression_safety_valve: 0.8, + ..Default::default() + }, + ); + let mut data = ModelData::new("test-compress-model-limited"); + data.max_input_tokens = Some(100_000); + let model = Model::from_config("openai", &[data]).remove(0); + ctx.session.as_mut().unwrap().set_model(model); + ctx.last_prompt_token_usage = Some(90_000); + ctx + } + + #[test] + #[serial] + fn successful_compression_clears_stale_prompt_measurement() { + let _sink = test_sink::install(); + let marker = "maybe-compress-stale-p6b"; + run_async(async { + let mut ctx = valve_tripped_ctx(marker); + let app = Arc::clone(&ctx.app.config); + + assert!(maybe_auto_compress(&mut ctx, app.as_ref()).await.succeeded); + assert_eq!( + ctx.last_prompt_token_usage, None, + "a successful compression must drop the pre-compression prompt measurement" + ); + + // A subsequent usage-less turn makes tokens reclaimable again; + // the valve must not re-fire from the vanished prompt's + // measurement. + { + let input = Input::from_str(&ctx, "next turn", None).unwrap(); + let session = ctx.session.as_mut().unwrap(); + session.add_message(&input, "reply").unwrap(); + } + assert!(!maybe_auto_compress(&mut ctx, app.as_ref()).await.fired); + }); + + assert_eq!(session_compressed_count(marker), 1); + } + + /// Counts `on_chat_completion_error` checkpoints in the session; the + /// literal must match `INTERRUPTED_RESPONSE_TEXT`, whose presence is + /// pinned by the give-up tests below. + fn interrupted_checkpoint_count(ctx: &RequestContext) -> usize { + ctx.session.as_ref().map_or(0, |session| { + session + .messages() + .iter() + .filter(|message| { + matches!(&message.content, MessageContent::Text(text) + if text.ends_with("[Response interrupted due to error]")) + }) + .count() + }) + } + + #[test] + #[serial] + fn error_retry_compresses_once_without_checkpointing() { + let _sink = test_sink::install(); + let marker = "error-retry-valve-k1s"; + run_async(async { + let mut ctx = valve_tripped_ctx(marker); + let app = Arc::clone(&ctx.app.config); + let abort_signal = create_abort_signal(); + + assert!(compress_for_error_retry(&mut ctx, app.as_ref(), false, &abort_signal).await); + assert!( + !ctx.is_compressing_session(), + "the compressing flag must be released after compression" + ); + assert_eq!( + interrupted_checkpoint_count(&ctx), + 0, + "the retry path must not checkpoint the in-flight turn" + ); + }); + + assert_eq!(session_compressed_count(marker), 1); + } + + #[test] + #[serial] + fn error_retry_gives_up_on_second_failure_with_single_checkpoint() { + let _sink = test_sink::install(); + let marker = "error-retry-giveup-p4t"; + run_async(async { + let mut ctx = valve_tripped_ctx(marker); + let app = Arc::clone(&ctx.app.config); + let abort_signal = create_abort_signal(); + let input = Input::from_str(&ctx, "still too big", None).unwrap(); + + assert!(compress_for_error_retry(&mut ctx, app.as_ref(), false, &abort_signal).await); + // The retried request failed too: no second compression, and the + // give-up path checkpoints the interrupted turn exactly once. + assert!(!compress_for_error_retry(&mut ctx, app.as_ref(), true, &abort_signal).await); + ctx.on_chat_completion_error(app.as_ref(), &input); + assert_eq!(interrupted_checkpoint_count(&ctx), 1); + }); + + assert_eq!(session_compressed_count(marker), 1); + } + + #[test] + #[serial] + fn error_retry_declines_when_valve_is_quiet() { + let _sink = test_sink::install(); + let marker = "error-retry-quiet-m8d"; + run_async(async { + let mut ctx = compressible_ctx_with( + marker, + AppConfig { + compression_keep_last: 1, + ..Default::default() + }, + ); + let app = Arc::clone(&ctx.app.config); + let abort_signal = create_abort_signal(); + let input = Input::from_str(&ctx, "transient failure", None).unwrap(); + + assert!(!compress_for_error_retry(&mut ctx, app.as_ref(), false, &abort_signal).await); + // Valve-false errors keep today's behavior: checkpoint and + // surface the error immediately, no compression, no retry. + ctx.on_chat_completion_error(app.as_ref(), &input); + assert_eq!(interrupted_checkpoint_count(&ctx), 1); + }); + + assert_eq!(session_compressed_count(marker), 0); + } + + #[test] + #[serial] + fn error_retry_flag_blocks_repeat_and_rearms_on_fresh_input() { + let _sink = test_sink::install(); + let marker = "error-retry-flag-c6y"; + run_async(async { + let mut ctx = valve_tripped_ctx(marker); + let app = Arc::clone(&ctx.app.config); + let abort_signal = create_abort_signal(); + + // Tool-loop iterations carry the flag forward: a set flag wins + // over a still-firing valve. + assert!(!compress_for_error_retry(&mut ctx, app.as_ref(), true, &abort_signal).await); + assert_eq!(session_compressed_count(marker), 0); + // A fresh input rebuild clears the flag and re-arms the retry. + assert!(compress_for_error_retry(&mut ctx, app.as_ref(), false, &abort_signal).await); + }); + + assert_eq!(session_compressed_count(marker), 1); + } + + #[test] + #[serial] + fn error_retry_refreshed_input_rebuilds_from_compressed_session() { + let _sink = test_sink::install(); + let marker = "error-retry-refresh-j7q"; + run_async(async { + let mut ctx = valve_tripped_ctx(marker); + let app = Arc::clone(&ctx.app.config); + let abort_signal = create_abort_signal(); + let mut input = Input::from_str(&ctx, "still too big", None).unwrap(); + let stale = serde_json::to_string(&input.build_messages().unwrap()).unwrap(); + + assert!(compress_for_error_retry(&mut ctx, app.as_ref(), false, &abort_signal).await); + // The snapshot captured at construction is untouched by the + // compression; retrying without a refresh would re-send the + // failed prompt byte for byte. + assert_eq!( + serde_json::to_string(&input.build_messages().unwrap()).unwrap(), + stale + ); + + input.refresh_session(&ctx); + let refreshed = serde_json::to_string(&input.build_messages().unwrap()).unwrap(); + assert_ne!(refreshed, stale); + let rebuilt = Input::from_str(&ctx, "still too big", None).unwrap(); + assert_eq!( + refreshed, + serde_json::to_string(&rebuilt.build_messages().unwrap()).unwrap(), + "a refreshed input must build the same request as one created after compression" + ); + }); + + assert_eq!(session_compressed_count(marker), 1); + } + + #[test] + #[serial] + fn error_retry_declines_when_user_aborted() { + let _sink = test_sink::install(); + let marker = "error-retry-abort-b2n"; + run_async(async { + let mut ctx = valve_tripped_ctx(marker); + let app = Arc::clone(&ctx.app.config); + let input = Input::from_str(&ctx, "killed mid-request", None).unwrap(); + + let ctrlc = create_abort_signal(); + ctrlc.set_ctrlc(); + assert!(!compress_for_error_retry(&mut ctx, app.as_ref(), false, &ctrlc).await); + let ctrld = create_abort_signal(); + ctrld.set_ctrld(); + assert!(!compress_for_error_retry(&mut ctx, app.as_ref(), false, &ctrld).await); + + // A user-killed request keeps today's behavior: checkpoint and + // surface the error, no compression, no re-send. + ctx.on_chat_completion_error(app.as_ref(), &input); + assert_eq!(interrupted_checkpoint_count(&ctx), 1); + }); + + assert_eq!(session_compressed_count(marker), 0); + } + #[test] fn test_process_command_line() { assert_eq!(parse_command(" ."), Some((".", None)));