feat: Suppor mid-turn auto compressions so sufficiently large histories don't create deadlocks for max input tokens
This commit is contained in:
9 files changed
+1199
-109
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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<u64>, value: Option<u64>) {
|
||||
if let Some(value) = value {
|
||||
|
||||
+30
-6
@@ -869,17 +869,22 @@ fn reasoning_thinking_block(item: &Value) -> Option<ThinkingBlock> {
|
||||
}
|
||||
|
||||
/// 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<TokenUsage> {
|
||||
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));
|
||||
}
|
||||
|
||||
@@ -71,6 +71,7 @@ pub struct AppConfig {
|
||||
pub compression_threshold: usize,
|
||||
pub compression_keep_last: usize,
|
||||
pub compression_model: Option<String>,
|
||||
pub compression_safety_valve: f32,
|
||||
pub summarization_prompt: Option<String>,
|
||||
pub summary_context_prompt: Option<String>,
|
||||
pub max_tool_result_chars: Option<usize>,
|
||||
@@ -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,
|
||||
|
||||
@@ -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<Session>) -> 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();
|
||||
|
||||
@@ -286,6 +286,7 @@ pub struct Config {
|
||||
pub compression_threshold: usize,
|
||||
pub compression_keep_last: usize,
|
||||
pub compression_model: Option<String>,
|
||||
pub compression_safety_valve: f32,
|
||||
pub summarization_prompt: Option<String>,
|
||||
pub summary_context_prompt: Option<String>,
|
||||
pub max_tool_result_chars: Option<usize>,
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
+374
-32
@@ -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<TokenUsage>,
|
||||
pub last_cost: Option<f64>,
|
||||
|
||||
/// 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<usize>,
|
||||
|
||||
pub tool_scope: ToolScope,
|
||||
|
||||
pub declared_function_names: HashSet<String>,
|
||||
@@ -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<String> {
|
||||
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<F, Fut>(&mut self, fetch: F) -> Result<()>
|
||||
where
|
||||
F: Fn(Input) -> Fut,
|
||||
Fut: Future<Output = Result<String>>,
|
||||
{
|
||||
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<String> {
|
||||
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<String> = 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::<Vec<_>>()
|
||||
.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<Message> = (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 {
|
||||
|
||||
+118
-6
@@ -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<usize>,
|
||||
) -> 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<usize>) -> 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();
|
||||
|
||||
+570
-65
@@ -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<String>, &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)));
|
||||
|
||||
Reference in new issue
Block a user