diff --git a/src/agent/worker.rs b/src/agent/worker.rs index 5f3221dd5..4279d5cfa 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -725,16 +725,19 @@ impl Worker { self.segments_run .store(segments_run, std::sync::atomic::Ordering::Relaxed); - // Pre-prompt maintenance: dedup stale tool results and check - // context usage *before* each LLM call, not just at segment - // boundaries. Fast models can accumulate large tool results - // within a single segment and exceed the context window before - // we ever reach a checkpoint. - if segments_run > 1 { - dedup_tool_results(&mut history); - self.maybe_compact_history(&mut compacted_history, &mut history) - .await; - } + // Dedup stale tool results and check context usage before + // handing control to the tool loop. This runs on the first + // segment too: a run that finishes inside one segment would + // otherwise never be checked at all, which is how a worker + // reached 269k tokens against a 128k trigger. + // + // It is still only a per-segment check — the loop inside a + // segment can add tens of thousands of tokens per turn without + // yielding — so the request-level ceiling in `SpacebotModel` is + // what actually guarantees the window is respected. + dedup_tool_results(&mut history); + self.maybe_compact_history(&mut compacted_history, &mut history) + .await; match self .hook diff --git a/src/llm/manager.rs b/src/llm/manager.rs index 1c4d59883..3120724c3 100644 --- a/src/llm/manager.rs +++ b/src/llm/manager.rs @@ -44,6 +44,67 @@ pub struct LlmManager { openai_oauth_credentials: RwLock>, /// Cached GitHub Copilot API token (exchanged from PAT, refreshed lazily). copilot_token: RwLock>, + /// What each model's requests are allowed to grow to. + /// + /// Lives here because every `SpacebotModel` already shares this manager, so + /// a ceiling learned by one run applies to the next without threading it + /// through fifteen construction sites. + context_ceilings: ArcSwap, +} + +/// What a request is allowed to grow to, per model. +/// +/// A published context window is not what a backend enforces: the same model +/// answers to a different ceiling depending on which API it is reached through, +/// and that ceiling moves without notice. `default` is the configured fallback; +/// `learned` holds what a provider has demonstrated by refusing a request of +/// known size. +#[derive(Debug, Default, Clone)] +pub struct ContextCeilings { + pub default: Option, + pub learned: HashMap, +} + +impl ContextCeilings { + /// What this model's requests must fit inside, if anything is known. + /// + /// A refusal only ever tightens: it proves the ceiling sits below the size + /// refused and says nothing about whether the configured window was too + /// generous, so the smaller of the two is what a request has to fit. + pub fn ceiling_for(&self, full_model_name: &str) -> Option { + match (self.learned.get(full_model_name).copied(), self.default) { + (Some(learned), Some(default)) => Some(learned.min(default)), + (learned, default) => learned.or(default), + } + } + + /// Fold a rejection of `estimated_tokens` into the ceilings. + /// + /// Returns `None` when nothing was learned: a rejection at or above what is + /// already known says nothing new, so only a smaller one tightens the + /// ceiling. Moving in one direction keeps a single unlucky large request + /// from undoing a limit that was correctly discovered. + pub fn with_overflow(&self, full_model_name: &str, estimated_tokens: usize) -> Option { + // Back off from the refused size rather than sitting on the boundary, + // since the estimate is approximate in both directions. + let ceiling = estimated_tokens.saturating_mul(9) / 10; + if ceiling == 0 { + return None; + } + if self + .ceiling_for(full_model_name) + .is_some_and(|known| known <= ceiling) + { + return None; + } + + let mut learned = self.learned.clone(); + learned.insert(full_model_name.to_string(), ceiling); + Some(Self { + default: self.default, + learned, + }) + } } impl LlmManager { @@ -62,6 +123,7 @@ impl LlmManager { anthropic_oauth_credentials: RwLock::new(None), openai_oauth_credentials: RwLock::new(None), copilot_token: RwLock::new(None), + context_ceilings: ArcSwap::from_pointee(ContextCeilings::default()), }) } @@ -142,9 +204,63 @@ impl LlmManager { anthropic_oauth_credentials: RwLock::new(anthropic_oauth_credentials), openai_oauth_credentials: RwLock::new(openai_oauth_credentials), copilot_token: RwLock::new(copilot_token), + context_ceilings: ArcSwap::from_pointee(ContextCeilings::default()), }) } + /// The configured fallback ceiling, applied to any model with nothing learned. + /// + /// Read-modify-write under `rcu`: a refusal recorded by a request in flight + /// must not be dropped by this write, and vice versa. + pub fn set_default_context_ceiling(&self, tokens: usize) { + self.context_ceilings.rcu(|current| ContextCeilings { + default: Some(tokens), + learned: current.learned.clone(), + }); + } + + /// What this model's requests must fit inside, if anything is known. + pub fn context_ceiling(&self, full_model_name: &str) -> Option { + self.context_ceilings.load().ceiling_for(full_model_name) + } + + /// Record that a request of this size was refused for exceeding the window. + /// + /// The refusal is the only trustworthy measurement available: it proves the + /// ceiling sits below `estimated_tokens`. Following the lowest observed + /// refusal means a backend that silently tightens its limit is tracked + /// rather than fought. + /// Read-modify-write under `rcu`, so two models learning at once cannot + /// drop each other's ceiling and a stale copy cannot widen a tighter one. + /// The closure can run more than once, which is safe: `with_overflow` is a + /// pure function of the state it is handed. + pub fn note_context_overflow(&self, full_model_name: &str, estimated_tokens: usize) { + let mut learned: Option = None; + self.context_ceilings.rcu(|current| { + match current.with_overflow(full_model_name, estimated_tokens) { + Some(updated) => { + learned = updated.ceiling_for(full_model_name); + updated + } + None => { + learned = None; + (**current).clone() + } + } + }); + let Some(ceiling) = learned else { + return; + }; + + tracing::warn!( + model = %full_model_name, + rejected_at = estimated_tokens, + ceiling, + "provider refused a request for exceeding its context window; \ + lowering the ceiling for this model" + ); + } + /// Atomically swap in new provider credentials. pub fn reload_config(&self, config: LlmConfig) { self.config.store(Arc::new(config)); @@ -482,3 +598,118 @@ impl LlmManager { .retain(|_, limited_at| limited_at.elapsed().as_secs() < cooldown_secs); } } + +#[cfg(test)] +mod context_ceiling_tests { + use super::ContextCeilings; + + #[test] + fn nothing_is_enforced_until_a_ceiling_is_known() { + let ceilings = ContextCeilings::default(); + assert_eq!(ceilings.ceiling_for("openai-chatgpt/gpt-5.6-sol"), None); + } + + #[test] + fn the_configured_default_applies_to_every_model() { + let ceilings = ContextCeilings { + default: Some(128_000), + ..Default::default() + }; + assert_eq!( + ceilings.ceiling_for("openai-chatgpt/gpt-5.6-sol"), + Some(128_000) + ); + assert_eq!(ceilings.ceiling_for("anything/else"), Some(128_000)); + } + + /// The case that killed two workers: the backend enforced far less than the + /// model advertises, and the only way to find out was to be refused. + #[test] + fn a_refusal_teaches_the_ceiling_for_that_model_alone() { + let ceilings = ContextCeilings { + default: Some(1_050_000), + ..Default::default() + }; + + let learned = ceilings + .with_overflow("openai-chatgpt/gpt-5.6-sol", 257_963) + .expect("a refusal teaches something"); + + let ceiling = learned + .ceiling_for("openai-chatgpt/gpt-5.6-sol") + .expect("learned"); + assert!( + ceiling < 257_963, + "the ceiling must sit below the size that was refused" + ); + assert_eq!(ceiling, 232_166); + + // Every other model keeps the configured default. + assert_eq!( + learned.ceiling_for("anthropic/claude-sonnet-4"), + Some(1_050_000) + ); + } + + /// A backend that tightens again must be followed down, and one that + /// happens to refuse a larger request must not undo what was learned. + #[test] + fn the_ceiling_only_ever_moves_down() { + let ceilings = ContextCeilings { + default: Some(400_000), + ..Default::default() + }; + + let first = ceilings.with_overflow("m", 300_000).expect("learned"); + let learned = first.ceiling_for("m").expect("learned"); + + assert!( + first.with_overflow("m", 350_000).is_none(), + "a larger refusal says nothing new" + ); + + let tighter = first.with_overflow("m", 200_000).expect("tightened"); + assert!(tighter.ceiling_for("m").expect("learned") < learned); + } + + /// A refusal proves the ceiling sits below the size refused. It proves + /// nothing about a configured window being too small, so it must never + /// raise one — with the shipped default of 128,000, a refusal at 257,963 + /// would otherwise learn 232,166 and start sending far more than the + /// operator asked for. + #[test] + fn a_refusal_cannot_raise_the_configured_ceiling() { + let ceilings = ContextCeilings { + default: Some(128_000), + ..Default::default() + }; + + assert!( + ceilings + .with_overflow("openai-chatgpt/gpt-5.6-sol", 257_963) + .is_none(), + "a refusal above the configured ceiling says nothing new" + ); + + // One below it still tightens, and stays tightened when the default is + // later raised. + let learned = ceilings.with_overflow("m", 100_000).expect("tightened"); + assert_eq!(learned.ceiling_for("m"), Some(90_000)); + + let raised = ContextCeilings { + default: Some(1_050_000), + learned: learned.learned.clone(), + }; + assert_eq!(raised.ceiling_for("m"), Some(90_000)); + assert_eq!(raised.ceiling_for("untouched"), Some(1_050_000)); + } + + #[test] + fn a_nonsense_refusal_is_ignored() { + let ceilings = ContextCeilings { + default: Some(128_000), + ..Default::default() + }; + assert!(ceilings.with_overflow("m", 0).is_none()); + } +} diff --git a/src/llm/model.rs b/src/llm/model.rs index c78c77548..75c76a0a0 100644 --- a/src/llm/model.rs +++ b/src/llm/model.rs @@ -1,5 +1,6 @@ //! SpacebotModel: Custom CompletionModel implementation that routes through LlmManager. +use crate::agent::compactor::{advance_past_stranded_tool_results, estimate_history_tokens}; use crate::config::{ApiType, ProviderConfig}; use crate::llm::manager::LlmManager; use crate::llm::routing::{ @@ -64,6 +65,53 @@ pub struct SpacebotModel { usage_accumulator: Option>>, } +/// Share of the ceiling held back for the model's own response. +const RESPONSE_RESERVE: f32 = 0.15; + +/// Tokens a request spends before any history: system prompt and tool schemas +/// are charged to the same window. +/// What the history has to fit inside, given the ceiling and the fixed cost of +/// the system prompt and tool schemas. +/// +/// The model needs room to answer, and the preamble and tool definitions are +/// charged to the same window as the history. The ceiling is the size of a whole +/// request, so subtracting the overhead here is the only place it is charged. +fn context_budget(ceiling: usize, overhead: usize) -> usize { + let reserve = (ceiling as f32 * RESPONSE_RESERVE) as usize; + ceiling.saturating_sub(reserve).saturating_sub(overhead) +} + +fn request_overhead_tokens(request: &CompletionRequest) -> usize { + let preamble = request.preamble.as_ref().map_or(0, |text| text.len()); + let tools: usize = request + .tools + .iter() + .map(|tool| tool.name.len() + tool.description.len() + tool.parameters.to_string().len()) + .sum(); + (preamble + tools) / 4 +} + +/// Drop the oldest turns until the history fits `budget`, returning how many +/// messages went. +/// +/// Cuts a quarter at a time so a history barely over budget does not lose far +/// more than it needs to, and aligns every cut so a tool result is never left +/// without the call it answers. Returns 0 when no aligned cut can shrink it, +/// which the caller reports rather than sending a request it knows will fail. +fn trim_history_to_budget(history: &mut Vec, budget: usize) -> usize { + let mut dropped = 0usize; + while estimate_history_tokens(history) > budget && history.len() > 2 { + let target = (history.len() / 4).max(1); + let cut = advance_past_stranded_tool_results(history, target, history.len() - 2); + if cut == 0 { + break; + } + history.drain(..cut); + dropped += cut; + } + dropped +} + impl SpacebotModel { pub fn provider(&self) -> &str { &self.provider @@ -192,6 +240,95 @@ impl SpacebotModel { Ok(()) } + /// Trim a request until it fits the ceiling the provider actually enforces. + /// + /// Every history a request can be built from passes through here: the + /// worker's segments, rig's internal tool loop, branches, the cortex. That + /// matters because a budget checked anywhere else can be skipped by a loop + /// that does not yield. A worker reached 269k tokens against a 128k + /// compaction trigger without the trigger ever being evaluated, because the + /// whole run happened inside one segment and the check only ran between + /// segments. This is the point that cannot be bypassed. + /// + /// Cutting here is a backstop, not a replacement for compaction: it drops + /// the oldest turns outright, where compaction summarises them first. It + /// exists so a run degrades instead of dying. + /// Returns the size of the request as sent, which is what a refusal + /// measures: history plus the system prompt and tool schemas charged to the + /// same window. The response reserve is this side's policy and is not part + /// of what the provider receives, so it is not counted here. + fn enforce_context_ceiling(&self, request: &mut CompletionRequest) -> usize { + let overhead = request_overhead_tokens(request); + let history_tokens = |request: &CompletionRequest| { + estimate_history_tokens(&request.chat_history.iter().cloned().collect::>()) + }; + + let Some(ceiling) = self.llm_manager.context_ceiling(&self.full_model_name) else { + return history_tokens(request) + overhead; + }; + + let budget = context_budget(ceiling, overhead); + let before = history_tokens(request); + if budget == 0 { + tracing::warn!( + model = %self.full_model_name, + ceiling, + overhead, + "the system prompt and tool schemas alone fill the context ceiling; \ + sending unchanged so the provider decides" + ); + return before + overhead; + } + if before <= budget { + return before + overhead; + } + + let mut history: Vec = + request.chat_history.iter().cloned().collect(); + let dropped = trim_history_to_budget(&mut history, budget); + + if dropped == 0 { + tracing::error!( + model = %self.full_model_name, + estimated = before, + budget, + "request exceeds the context ceiling and no aligned cut can shrink it" + ); + return before + overhead; + } + + let Ok(chat_history) = OneOrMany::many(history) else { + return before + overhead; + }; + + request.chat_history = chat_history; + let after = history_tokens(request); + + // The trim runs out of room before the budget when only the two + // retained messages are left, and the request goes anyway: the ceiling + // is an estimate, and the provider is the one that decides. + if after > budget { + tracing::warn!( + model = %self.full_model_name, + ceiling, + estimated_after = after, + budget, + "trimmed as far as an aligned cut allows and the request still \ + exceeds the ceiling" + ); + } else { + tracing::warn!( + model = %self.full_model_name, + ceiling, + estimated_before = before, + estimated_after = after, + dropped_messages = dropped, + "trimmed request history to fit the model's context ceiling" + ); + } + after + overhead + } + /// Repair a history a provider has already rejected, for one retry. /// /// The pre-send pass pairs every result, so a mismatch that survives it is @@ -240,7 +377,34 @@ impl SpacebotModel { } /// Direct call to the provider (no fallback logic). + /// + /// The ceiling is enforced here rather than once at the top of `completion` + /// because this is where the model that receives the request is known. A + /// fallback attempt builds its own `SpacebotModel`, so trimming higher up + /// would size one model's request against another model's limit and record + /// its refusal against the wrong name. async fn attempt_completion( + &self, + mut request: CompletionRequest, + ) -> Result, CompletionError> { + let sent_tokens = self.enforce_context_ceiling(&mut request); + let result = self.call_provider(request).await; + + // A rejection is the only trustworthy measurement of where the ceiling + // sits: the published window and the one the backend enforces are + // routinely different, and the difference moves without notice. + if let Err(ref error) = result + && routing::is_context_overflow_error(&error.to_string()) + { + self.llm_manager + .note_context_overflow(&self.full_model_name, sent_tokens); + } + + result + } + + /// Send a prepared request to whichever provider this model belongs to. + async fn call_provider( &self, request: CompletionRequest, ) -> Result, CompletionError> { @@ -762,6 +926,9 @@ impl CompletionModel for SpacebotModel { mut request: CompletionRequest, ) -> Result, CompletionError> { self.repair_request_history(&mut request)?; + // Streaming has no fallback chain, so this model is the one that + // receives the request and the one a refusal belongs to. + let sent_tokens = self.enforce_context_ceiling(&mut request); let mut result = self.dispatch_stream(request.clone()).await; @@ -777,6 +944,15 @@ impl CompletionModel for SpacebotModel { self.record_tool_history_recovery(result.is_ok()); } + // The refusal lands while the stream is opening, so it is measurable + // here for the same reason it is on the non-streaming path. + if let Err(ref error) = result + && routing::is_context_overflow_error(&error.to_string()) + { + self.llm_manager + .note_context_overflow(&self.full_model_name, sent_tokens); + } + result } } @@ -5026,3 +5202,131 @@ mod tests { assert!(msg.contains("invalid schema")); } } + +#[cfg(test)] +mod context_trim_tests { + use super::{context_budget, trim_history_to_budget}; + use crate::agent::compactor::estimate_history_tokens; + use crate::llm::manager::ContextCeilings; + use rig::message::{AssistantContent, Message, UserContent}; + + /// A refusal measures the whole request, so the system prompt and tool + /// schemas are already inside what is learned. Recording the history alone + /// meant the overhead was charged twice — once by the estimate that was + /// never counted, once by the budget — and the usable window shrank on + /// every refusal. + #[test] + fn the_learned_ceiling_and_the_budget_charge_overhead_once() { + let overhead = 20_000; + let history = 240_000; + + let learn = |size: usize| { + ContextCeilings::default() + .with_overflow("m", size) + .expect("a refusal teaches something") + .ceiling_for("m") + .expect("learned") + }; + + let from_whole_request = context_budget(learn(history + overhead), overhead); + let from_history_alone = context_budget(learn(history), overhead); + + assert!( + from_whole_request > from_history_alone, + "measuring only the history gives back a smaller window every time" + ); + // The next request still has to be smaller than the one that was refused. + assert!(from_whole_request + overhead < history + overhead); + } + + /// Overhead alone can fill the window, and there is nothing to trim then. + #[test] + fn a_budget_cannot_go_below_zero() { + assert_eq!(context_budget(10_000, 50_000), 0); + } + + fn assistant_tool_call(id: &str) -> Message { + Message::Assistant { + id: None, + content: rig::OneOrMany::one(AssistantContent::tool_call( + id, + "shell", + serde_json::json!({"command": "cat -n src/agent/worker.rs"}), + )), + } + } + + fn tool_result(id: &str, bytes: usize) -> Message { + Message::User { + content: rig::OneOrMany::one(UserContent::ToolResult(rig::message::ToolResult { + id: id.to_string(), + call_id: None, + content: rig::OneOrMany::one(rig::message::ToolResultContent::text( + "x".repeat(bytes), + )), + })), + } + } + + /// The shape that killed both workers: turns of parallel shell calls, each + /// result at the 50,000-byte cap, run until the history is twice the window. + fn overflowing_history() -> Vec { + let mut history = vec![Message::from("audit the worker backends")]; + for turn in 0..8 { + for call in 0..4 { + let id = format!("call_{turn}_{call}"); + history.push(assistant_tool_call(&id)); + history.push(tool_result(&id, 50_000)); + } + } + history + } + + #[test] + fn a_history_already_under_budget_is_left_alone() { + let mut history = vec![Message::from("hello")]; + assert_eq!(trim_history_to_budget(&mut history, 100_000), 0); + assert_eq!(history.len(), 1); + } + + /// The guarantee the whole change exists for: whatever the loop built, what + /// leaves fits. + #[test] + fn an_overflowing_history_is_brought_under_budget() { + let mut history = overflowing_history(); + let before = estimate_history_tokens(&history); + assert!( + before > 250_000, + "fixture should reproduce the real overflow, got {before}" + ); + + let dropped = trim_history_to_budget(&mut history, 200_000); + + assert!(dropped > 0); + let after = estimate_history_tokens(&history); + assert!(after <= 200_000, "history must fit the budget, got {after}"); + assert!(!history.is_empty()); + } + + /// A tighter ceiling has to cut harder, not give up. + #[test] + fn a_small_budget_still_produces_a_sendable_history() { + let mut history = overflowing_history(); + trim_history_to_budget(&mut history, 30_000); + + assert!(estimate_history_tokens(&history) <= 30_000); + assert!(!history.is_empty()); + } + + /// Trimming must not stand a result up without the call it answers. + #[test] + fn the_retained_head_is_never_a_stranded_result() { + let mut history = overflowing_history(); + trim_history_to_budget(&mut history, 120_000); + + assert!( + !crate::agent::compactor::opens_with_tool_result(&history[0]), + "the retained history must not open on a tool result" + ); + } +} diff --git a/src/main.rs b/src/main.rs index 3f74595d4..850e8f11b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1081,6 +1081,13 @@ async fn run( .with_context(|| "failed to initialize LLM manager")?, ); + // The hard ceiling every request is trimmed to fit. Compaction aims at the + // same number, but it only runs where a loop yields; this is enforced on + // the request itself, so a loop that never yields cannot exceed it. Raising + // `context_window` raises both — set it to what the backend actually + // enforces, which is not always what the model advertises. + llm_manager.set_default_context_ceiling(config.defaults.context_window); + // Shared embedding model (stateless, agent-agnostic) let embedding_cache_dir = config.instance_dir.join("embedding_cache"); let embedding_model = Arc::new( @@ -2127,6 +2134,12 @@ async fn run( { Ok(new_llm) => { let new_llm_manager = Arc::new(new_llm); + // Ceilings live on the manager, so the + // replacement starts with none and every agent + // built after setup would send unbounded. + new_llm_manager.set_default_context_ceiling( + new_config.defaults.context_window, + ); api_state.set_llm_manager(new_llm_manager.clone()).await; // Update agent_humans from the reloaded config // before initialize_agents so agents see the