fix(compaction): configurable summarizer tool_choice; wider overflow detection

Port of upstream grok-build 0.2.105 + 0.2.118.

- CompactionToolChoice (auto default | none), resolved env
  KIGI_COMPACTION_TOOL_CHOICE > [features].compaction_tool_choice >
  remote settings. Some OpenAI-compatible servers reject
  tool_choice:"none" with tools attached, failing every long-session
  compaction; auto keeps them working (the summarization prompt already
  forbids tool use). Threaded from agent + subagent spawn through
  CompactionConfig into generate_session_compact; the Responses gate
  counts hosted tools too, so a hosted-only request still carries the
  configured choice.
- is_context_length_error recognizes the "current message ... exceeds
  budget" phrasing, so that overflow is Fatal instead of burning the
  retry budget on a deterministic failure.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KSy9S31pZNWXtRZNVJ1sDf
This commit is contained in:
ZacharyZhang-NYandClaude Fable 5 committed 2026-08-27 10:22:18 -04:00
1 parent 6a4bff0d45
commit 1ae803519e
18 files changed
+190 -13

No files matched your search

@@ -668,6 +668,11 @@ pub struct RemoteSettings {
/// remote settings verbatim-input flag; env (`KIGI_COMPACTION_VERBATIM_INPUT`) and config override it. `None` = default (true).
#[serde(default)]
pub compaction_verbatim_input: Option<bool>,
/// Compaction summarizer `tool_choice` (`auto` | `none`) from remote
/// settings. Env (`KIGI_COMPACTION_TOOL_CHOICE`) and config override it.
/// Some backends reject `tool_choice: "none"` with tools attached.
#[serde(default)]
pub compaction_tool_choice: Option<String>,
/// Master switch for jemalloc heap sampling + threshold dumps.
/// `Some(true)` enables, `Some(false)` kill-switch, `None` = client default off.
#[serde(default)]
@@ -352,6 +352,7 @@ pub fn is_context_length_error(message: &str) -> bool {
|| m.contains("maximum prompt length")
|| m.contains("maximum context length")
|| m.contains("context_length_exceeded")
|| (m.contains("current message") && m.contains("exceeds budget"))
}
/// Decide whether a [`reqwest::Error`] is worth retrying.
@@ -387,6 +388,7 @@ mod tests {
"This model's maximum context length is 200000 tokens",
"invalid_request_error: prompt is too long: 300000 tokens > 200000 maximum",
"error type: context_length_exceeded",
"the current message exceeds budget: 120000 < 250000",
] {
assert!(is_context_length_error(msg), "should match: {msg}");
}
@@ -2171,6 +2171,19 @@ impl Config {
)
.with_segment_detail(self.resolve_compaction_detail())
}
/// Resolve summarizer `tool_choice`: env `KIGI_COMPACTION_TOOL_CHOICE` >
/// config > remote settings > default (`auto`).
pub(crate) fn resolve_compaction_tool_choice(
&self,
) -> crate::util::config::CompactionToolChoice {
crate::util::config::resolve_compaction_tool_choice_from(
env_string(crate::util::config::ENV_COMPACTION_TOOL_CHOICE).as_deref(),
self.features.compaction_tool_choice.as_deref(),
self.remote_settings
.as_ref()
.and_then(|r| r.compaction_tool_choice.as_deref()),
)
}
/// Resolve verbatim-input flag: env `KIGI_COMPACTION_VERBATIM_INPUT` > config > remote settings > default `true`.
pub(crate) fn resolve_compaction_verbatim_input(&self) -> bool {
BoolFlag::env("KIGI_COMPACTION_VERBATIM_INPUT")
@@ -3696,6 +3709,9 @@ pub struct Features {
/// Feed the summarizer the verbatim conversation instead of the lossy rewrite; `None` = defer to env/remote settings/default (true).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub compaction_verbatim_input: Option<bool>,
/// Summarizer `tool_choice` (`auto` | `none`); `None` = defer to env/remote settings/default (`auto`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub compaction_tool_choice: Option<String>,
/// Snapshot a completed subagent's isolated worktree into a durable git ref
/// and delete its directory (resume rehydrates from the ref). This is the
/// per-deployment rollout lever (set in managed_config.toml `[features]`).
@@ -2183,6 +2183,7 @@ impl MvpAgent {
.cfg
.borrow()
.resolve_compaction_verbatim_input();
let compaction_tool_choice = self.cfg.borrow().resolve_compaction_tool_choice();
let two_pass_enabled = self.cfg.borrow().is_two_pass_compaction_enabled();
let auto_update = self.cfg.borrow().cli.auto_update;
let client_type = *self.client_type.borrow();
@@ -2568,6 +2569,7 @@ impl MvpAgent {
system_prompt_label,
compaction_mode,
compaction_verbatim_input,
compaction_tool_choice,
two_pass_enabled,
buffering_settings,
origin_client.clone(),
@@ -1006,6 +1006,7 @@ pub(crate) async fn handle_subagent_request(
kigi_agent::DEFAULT_SYSTEM_PROMPT_LABEL.to_string(),
kigi_chat_state::CompactionMode::Summary,
ctx.resolve_compaction_verbatim_input(),
ctx.resolve_compaction_tool_choice(),
false,
None,
None,
@@ -397,6 +397,21 @@ impl SubagentSpawnContext {
.resolve()
.value
}
/// Subagent summarizer `tool_choice`, mirroring
/// `Config::resolve_compaction_tool_choice` (env > config > remote
/// settings > default `auto`).
pub fn resolve_compaction_tool_choice(&self) -> crate::util::config::CompactionToolChoice {
crate::util::config::resolve_compaction_tool_choice_from(
crate::agent::config::env_string(crate::util::config::ENV_COMPACTION_TOOL_CHOICE)
.as_deref(),
self.agent_config
.as_ref()
.and_then(|c| c.features.compaction_tool_choice.as_deref()),
self.remote_settings
.as_ref()
.and_then(|r| r.compaction_tool_choice.as_deref()),
)
}
/// Whether a completed subagent's worktree is snapshotted into a durable ref
/// and its directory deleted. Resolution mirrors the other subagent gates
/// (env > config > remote settings > default). Default `false` so it ships dark;
@@ -112,6 +112,7 @@ pub(crate) async fn spawn_session_actor(
system_prompt_label: String,
compaction_mode: kigi_chat_state::CompactionMode,
compaction_verbatim_input: bool,
compaction_tool_choice: crate::util::config::CompactionToolChoice,
two_pass_enabled: bool,
buffering_settings: Option<BufferingSettings>,
origin_client: Option<crate::http::OriginClientInfo>,
@@ -1032,6 +1033,7 @@ pub(crate) async fn spawn_session_actor(
previous_model: std::cell::Cell::new(None),
compaction_mode,
verbatim_input: compaction_verbatim_input,
tool_choice: compaction_tool_choice,
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -1535,6 +1537,7 @@ pub(crate) async fn spawn_session_on_thread(
system_prompt_label: String,
compaction_mode: kigi_chat_state::CompactionMode,
compaction_verbatim_input: bool,
compaction_tool_choice: crate::util::config::CompactionToolChoice,
two_pass_enabled: bool,
buffering_settings: Option<BufferingSettings>,
origin_client: Option<crate::http::OriginClientInfo>,
@@ -1679,6 +1682,7 @@ pub(crate) async fn spawn_session_on_thread(
system_prompt_label,
compaction_mode,
compaction_verbatim_input,
compaction_tool_choice,
two_pass_enabled,
buffering_settings,
origin_client,
@@ -157,6 +157,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -611,6 +612,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -871,6 +873,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -1865,6 +1868,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -176,6 +176,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -112,6 +112,7 @@ async fn create_test_actor(
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -539,6 +540,7 @@ async fn create_test_actor_with_memory(
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -1298,6 +1300,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -163,6 +163,7 @@ async fn create_test_actor_with_memory(
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -116,6 +116,7 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -226,6 +226,7 @@ pub(crate) async fn create_test_actor_ex(
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -212,6 +212,7 @@ impl SessionActor {
&sampling_config,
self.inference_idle_timeout,
wall_clock_budget_secs,
self.compaction.tool_choice,
)
.await
{
@@ -983,6 +984,7 @@ impl SessionActor {
sampling_config.clone(),
self.inference_idle_timeout,
wall_clock_budget_secs,
self.compaction.tool_choice,
);
let observer =
crate::session::helpers::full_replace_compaction::ShellFullReplaceObserver::new(
@@ -2198,6 +2200,7 @@ mod inline_auto_compact_flow_tests {
previous_model: std::cell::Cell::new(None),
compaction_mode: kigi_chat_state::CompactionMode::Transcript,
verbatim_input: true,
tool_choice: Default::default(),
prefire: crate::session::compaction_config::PrefireState::default(),
prefix_released: std::sync::atomic::AtomicBool::new(false),
},
@@ -146,6 +146,9 @@ pub struct CompactionConfig {
pub compaction_mode: kigi_chat_state::CompactionMode,
/// When `true`, feed the summarizer the verbatim conversation instead of the lossy rewrite (the retry loop may still fall back).
pub verbatim_input: bool,
/// `tool_choice` for summarizer requests with tools attached (`auto` |
/// `none`); some backends reject `"none"`.
pub tool_choice: crate::util::config::CompactionToolChoice,
/// Prefire two-pass state (background NOTE₁ cache + in-flight guard).
/// `Default` (empty cache, not in-flight).
pub prefire: PrefireState,
@@ -70,6 +70,8 @@ pub(crate) struct ShellCompactionSampler {
/// Wall-clock budget (secs) forwarded to `generate_session_compact` as the
/// reasoning-runaway backstop; `0` disables it.
wall_clock_budget_secs: u64,
/// Summarizer `tool_choice` forwarded to `generate_session_compact`.
tool_choice: crate::util::config::CompactionToolChoice,
/// Full output of the most recent successful sample (for L5 telemetry).
last_success: Mutex<Option<CompactOutput>>,
}
@@ -86,6 +88,7 @@ impl ShellCompactionSampler {
sampling_config: SamplingConfig,
idle_timeout: Duration,
wall_clock_budget_secs: u64,
tool_choice: crate::util::config::CompactionToolChoice,
) -> Self {
Self {
use_short_prompt,
@@ -97,6 +100,7 @@ impl ShellCompactionSampler {
sampling_config,
idle_timeout,
wall_clock_budget_secs,
tool_choice,
last_success: Mutex::new(None),
}
}
@@ -135,6 +139,7 @@ impl CompactionSampler for ShellCompactionSampler {
&self.sampling_config,
self.idle_timeout,
self.wall_clock_budget_secs,
self.tool_choice,
)
.await
{
@@ -339,9 +339,10 @@ where
/// prefix and force a full prefill on the summarizer call — attaching them
/// keeps the request prefix byte-identical to the turn requests so the
/// engine reuses the session's KV cache (the whole point of the verbatim
/// input path). Tool *use* is forbidden via `tool_choice: none` where the
/// backend can express it (ChatCompletions, Responses); the Messages wire
/// enum has no `none`, so that path relies on the prompt instruction alone.
/// input path). Tool *use* is forbidden by the prompt instruction;
/// `tool_choice` says what goes on the wire next to the attached tools
/// (ChatCompletions, Responses — some servers reject `"none"`, hence the
/// `auto` default). The Messages wire enum has no `none` either way.
///
/// Errors carry a [`CompactFailure`] classification so the caller can
/// short-circuit retries on deterministic failures (4xx schema violations,
@@ -356,8 +357,17 @@ pub(crate) async fn generate_session_compact(
sampling_config: &SamplingConfig,
idle_timeout: std::time::Duration,
wall_clock_budget_secs: u64,
tool_choice: crate::util::config::CompactionToolChoice,
) -> Result<CompactOutput, CompactFailure> {
let num_messages = chat_history.len();
let wire_tool_choice = match tool_choice {
crate::util::config::CompactionToolChoice::Auto => ToolChoice::auto(),
crate::util::config::CompactionToolChoice::None => ToolChoice::none(),
};
let conversation_tool_choice = match tool_choice {
crate::util::config::CompactionToolChoice::Auto => ConversationToolChoice::Auto,
crate::util::config::CompactionToolChoice::None => ConversationToolChoice::None,
};
let output = match sampling_config.api_backend {
ApiBackend::ChatCompletions => {
let chat_messages: Vec<ChatRequestMessage> =
@@ -373,7 +383,7 @@ pub(crate) async fn generate_session_compact(
.map(|t| ToolDefinition::function(t.name, t.description, t.parameters))
.collect(),
)
.with_tool_choice(ToolChoice::none());
.with_tool_choice(wire_tool_choice);
}
let sid = session_id.to_string();
message.x_kigi_conv_id = Some(sid.clone());
@@ -466,7 +476,10 @@ pub(crate) async fn generate_session_compact(
ApiBackend::Responses => {
let request = ConversationRequest {
items: chat_history,
tool_choice: (!tools.is_empty()).then_some(ConversationToolChoice::None),
// Hosted tools count: a hosted-only request with the choice
// omitted would leave the hosted tool callable mid-compact.
tool_choice: (!tools.is_empty() || !hosted_tools.is_empty())
.then_some(conversation_tool_choice),
tools,
hosted_tools,
model: Some(sampling_config.model.to_owned()),
@@ -1573,6 +1586,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await
.unwrap_or_else(|_| panic!("compaction must succeed"));
@@ -1661,6 +1675,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
let output = result
@@ -1669,11 +1684,11 @@ mod reasoning_compaction_regression_tests {
let _ = shutdown_tx.send(());
}
/// The compaction request must carry the turn loop's tool definitions
/// (prompt-prefix/KV-cache alignment) with `tool_choice: "none"`, and
/// must omit both keys when no tools are passed (Chat Completions rejects a bare
/// `tool_choice`).
/// (prompt-prefix/KV-cache alignment) with the configured `tool_choice`
/// (`auto` default; `none` when configured), and must omit both keys when
/// no tools are passed (Chat Completions rejects a bare `tool_choice`).
#[tokio::test]
async fn chat_completions_compaction_attaches_tools_with_tool_choice_none() {
async fn chat_completions_compaction_attaches_tools_with_configured_tool_choice() {
use std::sync::{Arc, Mutex};
let captured: Arc<Mutex<Vec<serde_json::Value>>> = Arc::new(Mutex::new(Vec::new()));
let cap = captured.clone();
@@ -1726,12 +1741,13 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await
.unwrap_or_else(|_| panic!("compaction with tools must succeed"));
let client = Client::new(config.clone()).unwrap();
generate_session_compact(
chat_history,
chat_history.clone(),
vec![],
vec![],
client,
@@ -1739,16 +1755,36 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await
.unwrap_or_else(|_| panic!("compaction without tools must succeed"));
let tools = vec![ToolSpec {
name: "read_file".to_string(),
description: Some("Reads a file".to_string()),
parameters: json!({ "type" : "object", "properties" : {} }),
}];
let client = Client::new(config.clone()).unwrap();
generate_session_compact(
chat_history,
tools,
vec![],
client,
acp::SessionId::new("test-session"),
&config,
std::time::Duration::from_secs(30),
0,
crate::util::config::CompactionToolChoice::None,
)
.await
.unwrap_or_else(|_| panic!("compaction with tool_choice none must succeed"));
let bodies = captured.lock().unwrap();
assert_eq!(bodies.len(), 2, "mock must have served both requests");
assert_eq!(bodies.len(), 3, "mock must have served all three requests");
let with_tools = &bodies[0];
assert_eq!(
with_tools["tool_choice"],
json!("none"),
"tool use must be disabled at decode time"
json!("auto"),
"default tool_choice is auto — some backends reject none; the prompt forbids tool use"
);
let sent_tools = with_tools["tools"]
.as_array()
@@ -1764,6 +1800,12 @@ mod reasoning_compaction_regression_tests {
without_tools.get("tool_choice").is_none(),
"tool_choice without tools is rejected by OpenAI-compat backends"
);
let none_configured = &bodies[2];
assert_eq!(
none_configured["tool_choice"],
json!("none"),
"configured none must reach the wire"
);
let _ = shutdown_tx.send(());
}
#[tokio::test]
@@ -1802,6 +1844,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_millis(150),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
match result {
@@ -1869,6 +1912,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_millis(150),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
match result {
@@ -1940,6 +1984,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_millis(150),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
match result {
@@ -2008,6 +2053,7 @@ mod reasoning_compaction_regression_tests {
&config,
std::time::Duration::from_millis(150),
0,
crate::util::config::CompactionToolChoice::Auto,
)
.await;
match result {
@@ -1,6 +1,42 @@
/// Default auto-compact threshold (% of context window) when no source sets it.
pub const DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT: u8 = 85;
/// `tool_choice` sent on compaction summarizer requests when tools are
/// attached (prefix-cache alignment keeps the tool definitions in the
/// request). Default `Auto` because some OpenAI-compatible servers reject
/// `tool_choice: "none"`; the summarization prompt itself forbids tool use.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum CompactionToolChoice {
#[default]
Auto,
None,
}
impl std::str::FromStr for CompactionToolChoice {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.trim().to_ascii_lowercase().as_str() {
"auto" => Ok(Self::Auto),
"none" => Ok(Self::None),
_ => Err(()),
}
}
}
pub(crate) const ENV_COMPACTION_TOOL_CHOICE: &str = "KIGI_COMPACTION_TOOL_CHOICE";
pub(crate) fn resolve_compaction_tool_choice_from(
env: Option<&str>,
config: Option<&str>,
remote: Option<&str>,
) -> CompactionToolChoice {
env.and_then(|s| s.parse().ok())
.or_else(|| config.and_then(|s| s.parse().ok()))
.or_else(|| remote.and_then(|s| s.parse().ok()))
.unwrap_or_default()
}
/// Env-var override for `auto_compact_threshold_percent`. Parsed as `u8`;
/// out-of-range or unparseable values are ignored.
pub(crate) const ENV_AUTO_COMPACT_THRESHOLD_PERCENT: &str = "KIGI_AUTO_COMPACT_THRESHOLD_PERCENT";
@@ -127,6 +163,34 @@ pub fn resolve_compaction_wall_clock_budget_secs(gb_global: Option<u64>) -> u64
resolved
}
#[cfg(test)]
mod compaction_tool_choice_tests {
use super::{CompactionToolChoice, resolve_compaction_tool_choice_from as resolve};
#[test]
fn precedence_env_config_remote_default() {
assert_eq!(resolve(None, None, None), CompactionToolChoice::Auto);
assert_eq!(
resolve(None, None, Some("none")),
CompactionToolChoice::None
);
assert_eq!(
resolve(None, Some("auto"), Some("none")),
CompactionToolChoice::Auto
);
assert_eq!(
resolve(Some("none"), Some("auto"), Some("auto")),
CompactionToolChoice::None
);
// Unparseable tiers fall through instead of masking lower tiers.
assert_eq!(
resolve(Some("bogus"), None, Some("none")),
CompactionToolChoice::None
);
assert_eq!(resolve(Some("NONE "), None, None), CompactionToolChoice::None);
}
}
#[cfg(test)]
mod compaction_wall_clock_budget_tests {
use super::resolve_compaction_wall_clock_budget_secs as resolve;