fix(chat): budget chat history against the model's served context window

Long insight/persona chats on llama-swap failed with "exceeds the
available context size" even after raising Context in the app:

- llama-server fixes its window at launch (-c) and ignores num_ctx, but
  the history budget trusted the requested num_ctx, so raising it kept
  MORE history than the server could hold. LlmClient::context_window()
  now reports the served window (llama-swap: /upstream/<model>/props,
  cached 10 min) and the budget uses min(requested, served).
- apply_context_budget could only drop tool-call pairs, so plain Q&A
  chats never shrank. It now also drops whole earlier exchanges (keeping
  user/assistant alternation) and, last, the current turn's tool results,
  never the in-progress question.
- The budget only ran once, before the tool loop. ContextFitter now
  fits every LLM call (on a copy; the stored transcript is untouched)
  and, on a context-overflow error, shrinks the budget from the server's
  reported token counts and retries once.

GET /insights/models reports each llama-swap model's context_window when
it's known without loading the model (cached, or currently running), for
the app's Context picker.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Cameron CordesandClaude Opus 5.5 committed 2026-10-15 19:21:04 -04:00
1 parent 1d2848db05
commit 305eda5894
8 files changed
+654 -104

No files matched your search

+7
View File
@@ -995,6 +995,8 @@ pub async fn generate_agentic_insight_handler(
/// configured `LLAMA_SWAP_VISION_MODEL` slot id, and `has_tool_calling` is
/// reported as true for every slot (llama-server is launched with `--jinja`
/// by convention — a misconfigured slot surfaces as a chat-call error).
/// `context_window` is each slot's served `-c` window when known without
/// loading the model (cached from a recent chat, or the model is running).
#[get("/insights/models")]
pub async fn get_available_models_handler(
_claims: Claims,
@@ -1005,6 +1007,10 @@ pub async fn get_available_models_handler(
if crate::ai::local_backend_is_llamacpp()
&& let Some(lc) = app_state.llamacpp.as_ref()
{
// Served windows for the context picker; never loads a model.
let windows = lc
.known_context_windows(&app_state.llamacpp_allowed_models)
.await;
let models: Vec<ModelCapabilities> = app_state
.llamacpp_allowed_models
.iter()
@@ -1012,6 +1018,7 @@ pub async fn get_available_models_handler(
name: name.clone(),
has_vision: true,
has_tool_calling: true,
context_window: windows.get(name).copied(),
})
.collect();
let primary = ServerModels {
+455 -87
View File
@@ -21,12 +21,12 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
use uuid::Uuid;
pub const DEFAULT_MAX_ITERATIONS: usize = 6;
/// Assumed context window when the request doesn't specify `num_ctx`.
/// The llama-swap chat slots serve 20k-131k contexts and real conversations
/// rarely pass ~16k tokens, so 32k keeps the truncation pass from gutting
/// history that the server could comfortably hold (which also destroys the
/// server's KV-cache prefix reuse). Override per-deploy with
/// AGENTIC_CHAT_DEFAULT_NUM_CTX if the serving models change shape.
/// Assumed context window when the request doesn't specify `num_ctx` and
/// the backend can't report the window it serves (llama-swap can; see
/// `effective_num_ctx`). Real conversations rarely pass ~16k tokens, so 32k
/// keeps the truncation pass from gutting history the server could hold
/// (which also destroys the server's KV-cache prefix reuse). Override
/// per-deploy with AGENTIC_CHAT_DEFAULT_NUM_CTX.
const DEFAULT_NUM_CTX: i32 = 32768;
/// Headroom reserved for the model's response, deducted from the context
/// budget when deciding whether to truncate the replayed history.
@@ -323,12 +323,10 @@ impl InsightChatService {
None
};
// 6. Apply truncation budget. Drops oldest tool_call+tool pairs
// (preserves system + first user including any images).
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
let truncated = apply_context_budget(&mut messages, budget_bytes);
// 6. Apply truncation budget (see `apply_context_budget`), sized to
// the window the model is actually served with.
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
let truncated = fitter.trim(&mut messages);
if truncated {
span.set_attribute("history_truncated", true);
}
@@ -371,9 +369,8 @@ impl InsightChatService {
iterations_used = iteration + 1;
log::info!("Chat iteration {}/{}", iterations_used, max_iterations);
let (response, prompt_tokens, eval_tokens) = backend
.chat()
.chat_with_tools(messages.clone(), tools.clone())
let (response, prompt_tokens, eval_tokens) = fitter
.chat(backend.chat(), &messages, tools.clone())
.instrument(loop_span.clone())
.await?;
last_prompt_eval_count = prompt_tokens;
@@ -429,9 +426,8 @@ impl InsightChatService {
iterations_used
);
let synthetic_idx = push_synthetic_final_prompt(&mut messages);
let (final_response, prompt_tokens, eval_tokens) = backend
.chat()
.chat_with_tools(messages.clone(), vec![])
let (final_response, prompt_tokens, eval_tokens) = fitter
.chat(backend.chat(), &messages, vec![])
.instrument(loop_span.clone())
.await?;
last_prompt_eval_count = prompt_tokens;
@@ -1021,10 +1017,8 @@ impl InsightChatService {
None
};
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
let truncated = apply_context_budget(&mut messages, budget_bytes);
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
let truncated = fitter.trim(&mut messages);
if truncated {
let _ = entry.push_event(ChatStreamEvent::Truncated).await;
}
@@ -1050,6 +1044,7 @@ impl InsightChatService {
&active_persona,
max_iterations,
&entry,
&mut fitter,
)
.await?;
let AgenticLoopOutcome {
@@ -1255,6 +1250,7 @@ impl InsightChatService {
user_msg.images = Some(vec![img.clone()]);
}
let mut messages = vec![system_msg, user_msg];
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
let outcome = self
.run_streaming_agentic_loop_with_entry(
@@ -1267,6 +1263,7 @@ impl InsightChatService {
&active_persona,
max_iterations,
&entry,
&mut fitter,
)
.await?;
let AgenticLoopOutcome {
@@ -1354,6 +1351,7 @@ impl InsightChatService {
active_persona: &str,
max_iterations: usize,
entry: &Arc<TurnEntry>,
fitter: &mut ContextFitter,
) -> Result<AgenticLoopOutcome> {
crate::ai::insight_chat::run_streaming_agentic_loop_with_entry(
&self.generator,
@@ -1366,6 +1364,7 @@ impl InsightChatService {
active_persona,
max_iterations,
entry,
fitter,
)
.await
}
@@ -1494,10 +1493,8 @@ impl InsightChatService {
};
// Truncate before appending the new user turn.
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
let truncated = apply_context_budget(&mut messages, budget_bytes);
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
let truncated = fitter.trim(&mut messages);
if truncated {
let _ = tx.send(ChatStreamEvent::Truncated).await;
}
@@ -1520,6 +1517,7 @@ impl InsightChatService {
&active_persona,
max_iterations,
&tx,
&mut fitter,
)
.await?;
let AgenticLoopOutcome {
@@ -1740,6 +1738,7 @@ impl InsightChatService {
user_msg.images = Some(vec![img.clone()]);
}
let mut messages = vec![system_msg, user_msg];
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
let outcome = self
.run_streaming_agentic_loop(
@@ -1752,6 +1751,7 @@ impl InsightChatService {
&active_persona,
max_iterations,
&tx,
&mut fitter,
)
.await?;
let AgenticLoopOutcome {
@@ -1833,6 +1833,7 @@ impl InsightChatService {
active_persona: &str,
max_iterations: usize,
tx: &tokio::sync::mpsc::Sender<ChatStreamEvent>,
fitter: &mut ContextFitter,
) -> Result<AgenticLoopOutcome> {
let mut tool_calls_made = 0usize;
let mut iterations_used = 0usize;
@@ -1849,9 +1850,8 @@ impl InsightChatService {
})
.await;
let mut stream = backend
.chat()
.chat_with_tools_stream(messages.clone(), tools.clone())
let mut stream = fitter
.chat_stream(backend.chat(), messages, tools.clone())
.await?;
let mut final_message: Option<ChatMessage> = None;
@@ -1941,10 +1941,7 @@ impl InsightChatService {
// pending_tools at this position (wiping the prior tool
// calls from the final assistant render).
let synthetic_idx = push_synthetic_final_prompt(messages);
let mut stream = backend
.chat()
.chat_with_tools_stream(messages.clone(), vec![])
.await?;
let mut stream = fitter.chat_stream(backend.chat(), messages, vec![]).await?;
let mut final_message: Option<ChatMessage> = None;
while let Some(ev) = stream.next().await {
let ev = ev?;
@@ -2270,6 +2267,145 @@ pub(crate) fn env_default_num_ctx() -> i32 {
.max(RESPONSE_HEADROOM_TOKENS as i32 + 1024)
}
/// Context window (tokens) to budget a chat turn against: the smaller of what
/// the request asked for and what the server actually serves. llama-swap
/// models fix their window at launch (`-c` in config.yaml) and ignore the
/// per-request `num_ctx`, so budgeting against a bigger requested value just
/// keeps more history than the server can hold.
pub(crate) async fn effective_num_ctx(
client: &dyn crate::ai::llm_client::LlmClient,
requested: Option<i32>,
) -> usize {
let requested = requested.filter(|n| *n > 0).map(|n| n as usize);
match (requested, client.context_window().await) {
(Some(r), Some(served)) => r.min(served),
(None, Some(served)) => served,
(Some(r), None) => r,
(None, None) => env_default_num_ctx() as usize,
}
}
/// Keeps each LLM call of a chat turn inside the context budget. Every call
/// goes through [`ContextFitter::chat`] / [`ContextFitter::chat_stream`],
/// which send a trimmed copy of `messages` (the caller's vector — and so what
/// gets persisted — is untouched) and, if the server still rejects the prompt
/// as too long (the byte-per-token estimate is only a heuristic), shrink the
/// budget from the server's own token counts and retry once. The shrunk
/// budget sticks for the rest of the turn.
pub struct ContextFitter {
budget_bytes: usize,
}
impl ContextFitter {
pub(crate) fn from_num_ctx(num_ctx: usize) -> Self {
let budget_tokens = num_ctx.saturating_sub(RESPONSE_HEADROOM_TOKENS);
Self {
budget_bytes: budget_tokens.saturating_mul(BYTES_PER_TOKEN),
}
}
/// Budget for `client`'s model; see [`effective_num_ctx`].
pub(crate) async fn for_client(
client: &dyn crate::ai::llm_client::LlmClient,
requested: Option<i32>,
) -> Self {
Self::from_num_ctx(effective_num_ctx(client, requested).await)
}
/// Trim `messages` in place to the budget. Used once on the replayed
/// history before the turn starts; returns true when anything was dropped.
pub(crate) fn trim(&self, messages: &mut Vec<ChatMessage>) -> bool {
apply_context_budget(messages, self.budget_bytes)
}
fn fit(&self, messages: &[ChatMessage]) -> Vec<ChatMessage> {
let mut fitted = messages.to_vec();
apply_context_budget(&mut fitted, self.budget_bytes);
fitted
}
/// Shrink the budget after a context-overflow error. Returns false (and
/// leaves the budget alone) when `err` isn't one.
fn shrink_after(&mut self, err: &anyhow::Error, sent: &[ChatMessage]) -> bool {
let Some(counts) = parse_context_overflow(&format!("{err:#}")) else {
return false;
};
let sent_bytes = estimate_bytes(sent).min(self.budget_bytes);
let shrunk = match counts {
// Scale what we sent by the server's measured overshoot, with 10%
// slack since trimming lands on message boundaries.
Some((prompt_tokens, n_ctx)) if prompt_tokens > 0 => {
let target = n_ctx.saturating_sub(RESPONSE_HEADROOM_TOKENS);
sent_bytes.saturating_mul(target) / prompt_tokens * 9 / 10
}
_ => sent_bytes / 2,
};
tracing::warn!(
old_budget_bytes = self.budget_bytes,
new_budget_bytes = shrunk,
"chat prompt exceeded the model's context window; trimming and retrying"
);
self.budget_bytes = shrunk;
true
}
pub(crate) async fn chat(
&mut self,
client: &dyn crate::ai::llm_client::LlmClient,
messages: &[ChatMessage],
tools: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
let sent = self.fit(messages);
match client.chat_with_tools(sent.clone(), tools.clone()).await {
Err(e) if self.shrink_after(&e, &sent) => {
client.chat_with_tools(self.fit(messages), tools).await
}
result => result,
}
}
pub(crate) async fn chat_stream(
&mut self,
client: &dyn crate::ai::llm_client::LlmClient,
messages: &[ChatMessage],
tools: Vec<Tool>,
) -> Result<futures::stream::BoxStream<'static, Result<LlmStreamEvent>>> {
let sent = self.fit(messages);
match client
.chat_with_tools_stream(sent.clone(), tools.clone())
.await
{
Err(e) if self.shrink_after(&e, &sent) => {
client
.chat_with_tools_stream(self.fit(messages), tools)
.await
}
result => result,
}
}
}
/// Recognise a "prompt longer than the context window" error from any
/// backend. `None` = some other error. `Some(Some((prompt_tokens, n_ctx)))`
/// when the body carries llama-server's counts
/// (`exceed_context_size_error` with `n_prompt_tokens` / `n_ctx`),
/// `Some(None)` when it's an overflow without usable numbers.
fn parse_context_overflow(err: &str) -> Option<Option<(usize, usize)>> {
let lower = err.to_lowercase();
let is_overflow = lower.contains("exceed_context_size")
|| lower.contains("context_length_exceeded")
|| lower.contains("maximum context length")
|| (lower.contains("exceed") && lower.contains("context"));
if !is_overflow {
return None;
}
let num = |key: &str| -> Option<usize> {
let re = regex::Regex::new(&format!(r#""{key}"\s*:\s*(\d+)"#)).ok()?;
re.captures(err)?.get(1)?.as_str().parse().ok()
};
Some(num("n_prompt_tokens").zip(num("n_ctx")))
}
/// Append a per-turn iteration-budget reminder to the replayed system
/// message so the model knows how many tool-calling rounds this turn gets.
/// Returns the original `content` so the caller can restore it before
@@ -2323,6 +2459,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
active_persona: &str,
max_iterations: usize,
entry: &Arc<TurnEntry>,
fitter: &mut ContextFitter,
) -> Result<AgenticLoopOutcome> {
let mut tool_calls_made = 0usize;
let mut iterations_used = 0usize;
@@ -2350,9 +2487,8 @@ pub async fn run_streaming_agentic_loop_with_entry(
})
.await;
let mut stream = backend
.chat()
.chat_with_tools_stream(messages.clone(), tools.clone())
let mut stream = fitter
.chat_stream(backend.chat(), messages, tools.clone())
.await?;
let mut final_message: Option<ChatMessage> = None;
@@ -2431,10 +2567,7 @@ pub async fn run_streaming_agentic_loop_with_entry(
if final_content.is_empty() {
let synthetic_idx = push_synthetic_final_prompt(messages);
let mut stream = backend
.chat()
.chat_with_tools_stream(messages.clone(), vec![])
.await?;
let mut stream = fitter.chat_stream(backend.chat(), messages, vec![]).await?;
let mut final_message: Option<ChatMessage> = None;
while let Some(ev) = stream.next().await {
let ev = ev?;
@@ -2729,65 +2862,99 @@ pub(crate) fn truncate_tool_result(s: &str) -> (String, bool) {
}
}
/// Trim history to fit within `budget_bytes` of serialized JSON. Preserves
/// the system message and the first user message (with its base64 images
/// intact, since dropping those would invalidate the model's prior visual
/// reasoning). Drops the oldest assistant-tool_call + corresponding
/// tool-result pair on each pass until the budget is met or only the
/// preserved prefix remains.
/// Trim history to fit within `budget_bytes` of serialized JSON. Always
/// preserves the system message and the first user message (with its base64
/// images intact, since dropping those would invalidate the model's prior
/// visual reasoning) plus that first exchange's final reply, i.e. the
/// insight itself.
///
/// Drops, oldest first, in three phases until the budget is met:
/// 1. tool-call scaffolding (an assistant `tool_calls` turn + its tool
/// results) inside *completed* exchanges — cheapest to lose, the final
/// reply that used them stays;
/// 2. whole completed exchanges (user turn through its final reply) after
/// the first. Dropping a whole exchange keeps user/assistant alternation
/// intact, which strict chat templates (Gemma) enforce;
/// 3. tool-call scaffolding in the in-progress exchange (the current turn's
/// own tool results, which grow during the agentic loop).
///
/// An exchange is "completed" when it contains an assistant reply without
/// tool calls. The current turn never is while the loop runs, so the user's
/// question is never dropped — even after `push_synthetic_final_prompt`
/// appends a later user message.
///
/// Returns true when at least one message was dropped.
pub(crate) fn apply_context_budget(messages: &mut Vec<ChatMessage>, budget_bytes: usize) -> bool {
if budget_bytes == 0 {
if budget_bytes == 0 || estimate_bytes(messages) <= budget_bytes {
return false;
}
if estimate_bytes(messages) <= budget_bytes {
return false;
}
// Find the index past the protected prefix: system messages + the first
// user message. Everything after is droppable in pairs.
let first_user_idx = messages.iter().position(|m| m.role == "user");
let preserve_through = match first_user_idx {
Some(i) => i, // keep [0..=i]
None => return false,
};
let mut dropped_any = false;
loop {
if estimate_bytes(messages) <= budget_bytes {
while estimate_bytes(messages) > budget_bytes {
let Some(range) = next_budget_drop(messages) else {
break;
}
// Find the oldest assistant-with-tool_calls strictly after the
// preserved prefix. Drop it together with the following tool turn(s)
// until we hit the next assistant or user turn.
let drop_start = (preserve_through + 1..messages.len()).find(|&i| {
let m = &messages[i];
m.role == "assistant"
&& m.tool_calls
.as_ref()
.map(|c| !c.is_empty())
.unwrap_or(false)
});
let Some(start) = drop_start else { break };
// Determine end: drop the assistant turn plus any contiguous tool
// result turns that follow.
let mut end = start + 1;
while end < messages.len() && messages[end].role == "tool" {
end += 1;
}
// Stop if dropping these would leave the just-appended user turn at
// the end alone with no preceding context — we still want it kept.
if end > messages.len() {
break;
}
messages.drain(start..end);
};
messages.drain(range);
dropped_any = true;
}
dropped_any
}
/// Pick the next span of `messages` for [`apply_context_budget`] to drop,
/// following its phase order. `None` when nothing droppable remains.
fn next_budget_drop(messages: &[ChatMessage]) -> Option<std::ops::Range<usize>> {
let user_idxs: Vec<usize> = messages
.iter()
.enumerate()
.filter(|(_, m)| m.role == "user")
.map(|(i, _)| i)
.collect();
let first_user = *user_idxs.first()?;
// Exchanges: each user turn through the message before the next one.
let exchanges: Vec<std::ops::Range<usize>> = user_idxs
.iter()
.enumerate()
.map(|(n, &u)| u..user_idxs.get(n + 1).copied().unwrap_or(messages.len()))
.collect();
let completed = |r: &std::ops::Range<usize>| {
messages[r.clone()]
.iter()
.any(|m| m.role == "assistant" && !has_tool_calls(m))
};
// Oldest tool-call turn (+ trailing tool results) inside `r`.
let tool_pair_in = |r: &std::ops::Range<usize>| {
let start = (r.start + 1..r.end)
.find(|&i| messages[i].role == "assistant" && has_tool_calls(&messages[i]))?;
let mut end = start + 1;
while end < r.end && messages[end].role == "tool" {
end += 1;
}
Some(start..end)
};
// Phase 1: tool scaffolding in completed exchanges.
if let Some(r) = exchanges
.iter()
.filter(|r| completed(r))
.find_map(tool_pair_in)
{
return Some(r);
}
// Phase 2: whole completed exchanges after the first.
if let Some(r) = exchanges
.iter()
.find(|r| r.start != first_user && completed(r))
{
return Some(r.clone());
}
// Phase 3: tool scaffolding in the in-progress exchange(s).
exchanges.iter().find_map(tool_pair_in)
}
fn has_tool_calls(m: &ChatMessage) -> bool {
m.tool_calls.as_ref().is_some_and(|c| !c.is_empty())
}
/// Estimate the serialized byte size of `messages` for the truncation budget,
/// EXCLUDING inlined base64 image payloads. Images are charged a flat
/// `IMAGE_TOKENS_EACH` instead: their base64 is hundreds of KB of characters
@@ -2913,6 +3080,207 @@ mod tests {
assert!(!dropped);
}
fn roles(msgs: &[ChatMessage]) -> Vec<&str> {
msgs.iter().map(|m| m.role.as_str()).collect()
}
#[test]
fn truncation_drops_old_plain_exchanges_when_no_tool_pairs_left() {
// Long chat of plain Q&A: nothing but whole exchanges to drop. The
// old pass could only drop tool pairs, so this never shrank.
let filler = "x".repeat(2_000);
let mut msgs = vec![
ChatMessage::system("sys"),
ChatMessage::user("first"),
assistant_text("the insight"),
ChatMessage::user(format!("q1 {filler}")),
assistant_text(&format!("a1 {filler}")),
ChatMessage::user(format!("q2 {filler}")),
assistant_text(&format!("a2 {filler}")),
ChatMessage::user("current question"),
];
let budget = estimate_bytes(&msgs) - 1_000;
assert!(apply_context_budget(&mut msgs, budget));
assert!(estimate_bytes(&msgs) <= budget);
// Oldest exchange (q1/a1) went; the insight exchange, q2/a2 and the
// current question stay, and roles still alternate.
assert_eq!(
roles(&msgs),
["system", "user", "assistant", "user", "assistant", "user"]
);
assert_eq!(msgs[2].content, "the insight");
assert!(msgs[3].content.starts_with("q2"));
assert_eq!(msgs[5].content, "current question");
}
#[test]
fn truncation_drops_tool_scaffolding_before_whole_exchanges() {
let mut msgs = vec![
ChatMessage::system("sys"),
ChatMessage::user("first"),
assistant_text("the insight"),
ChatMessage::user("q1"),
assistant_with_tool_call("lookup"),
ChatMessage::tool_result("r".repeat(4_000)),
assistant_text("a1"),
ChatMessage::user("current question"),
];
let budget = estimate_bytes(&msgs) - 1_000;
assert!(apply_context_budget(&mut msgs, budget));
// Only the tool pair was needed; q1/a1 survive.
assert_eq!(
roles(&msgs),
["system", "user", "assistant", "user", "assistant", "user"]
);
assert_eq!(msgs[4].content, "a1");
}
#[test]
fn truncation_never_drops_the_in_progress_question() {
// Mid-loop with the synthetic final prompt appended: the real
// question's exchange has no final reply yet, so phase 2 must skip
// it; only its tool results (phase 3) can go.
let mut msgs = vec![
ChatMessage::system("sys"),
ChatMessage::user("first"),
assistant_text("the insight"),
ChatMessage::user("real question"),
assistant_with_tool_call("lookup"),
ChatMessage::tool_result("r".repeat(8_000)),
ChatMessage::user("synthetic: answer now"),
];
assert!(apply_context_budget(&mut msgs, 600));
assert!(msgs.iter().any(|m| m.content == "real question"));
assert!(msgs.iter().all(|m| m.role != "tool"));
}
#[test]
fn parse_context_overflow_reads_llama_server_counts() {
let err = r#"llama-swap chat request failed: 400 Bad Request — {"error":{"code":400,"message":"the request exceeds the available context size, try increasing it","type":"exceed_context_size_error","n_prompt_tokens":40312,"n_ctx":32768}}"#;
assert_eq!(parse_context_overflow(err), Some(Some((40312, 32768))));
assert_eq!(
parse_context_overflow("400 — This model's maximum context length is 128000 tokens"),
Some(None)
);
assert_eq!(parse_context_overflow("500 Internal Server Error"), None);
}
/// Scripted client for `ContextFitter`: fails the first `fail_first`
/// calls with `error`, records the size of every prompt it's sent.
struct ScriptedClient {
window: Option<usize>,
fail_first: usize,
error: String,
sent: Mutex<Vec<usize>>,
}
#[async_trait::async_trait]
impl crate::ai::llm_client::LlmClient for ScriptedClient {
async fn generate(
&self,
_: &str,
_: Option<&str>,
_: Option<Vec<String>>,
) -> Result<String> {
unimplemented!()
}
async fn chat_with_tools(
&self,
messages: Vec<ChatMessage>,
_: Vec<Tool>,
) -> Result<(ChatMessage, Option<i32>, Option<i32>)> {
let mut sent = self.sent.lock().unwrap();
sent.push(messages.len());
if sent.len() <= self.fail_first {
bail!("{}", self.error);
}
Ok((assistant_text("ok"), None, None))
}
async fn chat_with_tools_stream(
&self,
_: Vec<ChatMessage>,
_: Vec<Tool>,
) -> Result<futures::stream::BoxStream<'static, Result<LlmStreamEvent>>> {
unimplemented!()
}
async fn generate_embeddings(&self, _: &[&str]) -> Result<Vec<Vec<f32>>> {
unimplemented!()
}
async fn describe_image(&self, _: &str) -> Result<String> {
unimplemented!()
}
async fn list_models(&self) -> Result<Vec<crate::ai::llm_client::ModelCapabilities>> {
unimplemented!()
}
async fn model_capabilities(
&self,
_: &str,
) -> Result<crate::ai::llm_client::ModelCapabilities> {
unimplemented!()
}
fn primary_model(&self) -> &str {
"scripted"
}
async fn context_window(&self) -> Option<usize> {
self.window
}
}
fn scripted(window: Option<usize>, fail_first: usize, error: &str) -> ScriptedClient {
ScriptedClient {
window,
fail_first,
error: error.to_string(),
sent: Mutex::new(Vec::new()),
}
}
#[tokio::test]
async fn effective_num_ctx_caps_request_at_served_window() {
let served_32k = scripted(Some(32_768), 0, "");
assert_eq!(effective_num_ctx(&served_32k, Some(65_536)).await, 32_768);
assert_eq!(effective_num_ctx(&served_32k, Some(8_192)).await, 8_192);
assert_eq!(effective_num_ctx(&served_32k, None).await, 32_768);
let unknown = scripted(None, 0, "");
assert_eq!(effective_num_ctx(&unknown, Some(65_536)).await, 65_536);
}
#[tokio::test]
async fn fitter_retries_overflow_with_a_smaller_prompt() {
let filler = "x".repeat(2_000);
let mut msgs = vec![
ChatMessage::system("sys"),
ChatMessage::user("first"),
assistant_text("insight"),
];
for i in 0..10 {
msgs.push(ChatMessage::user(format!("q{i} {filler}")));
msgs.push(assistant_text(&format!("a{i} {filler}")));
}
msgs.push(ChatMessage::user("current"));
let client = scripted(
None,
1,
r#"400 — {"error":{"type":"exceed_context_size_error","n_prompt_tokens":20000,"n_ctx":10000}}"#,
);
let mut fitter = ContextFitter::from_num_ctx(1_000_000);
let (reply, _, _) = fitter.chat(&client, &msgs, vec![]).await.unwrap();
assert_eq!(reply.content, "ok");
let sent = client.sent.lock().unwrap().clone();
assert_eq!(sent.len(), 2, "one retry");
assert_eq!(sent[0], msgs.len(), "first attempt fit the generous budget");
assert!(sent[1] < sent[0], "retry sent a trimmed prompt");
}
#[tokio::test]
async fn fitter_does_not_retry_other_errors() {
let client = scripted(None, 5, "502 Bad Gateway");
let mut fitter = ContextFitter::from_num_ctx(32_768);
let msgs = vec![ChatMessage::system("s"), ChatMessage::user("u")];
assert!(fitter.chat(&client, &msgs, vec![]).await.is_err());
assert_eq!(client.sent.lock().unwrap().len(), 1);
}
#[test]
fn render_tree_path_detects_fork_at_tool_dispatch_child() {
// Regression: a regeneration where the model replies with a tool call
+149 -1
View File
@@ -27,6 +27,7 @@ use opentelemetry::trace::Status;
use reqwest::Client;
use serde::Deserialize;
use serde_json::{Value, json};
use std::collections::HashMap;
use std::time::Duration;
use tracing_opentelemetry::OpenTelemetrySpanExt as _;
@@ -41,6 +42,18 @@ const DEFAULT_VISION_MODEL: &str = "vision";
const DEFAULT_EMBEDDING_MODEL: &str = "embed";
const DEFAULT_TTS_MODEL: &str = "chatterbox";
const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 180;
/// How long a served-context-window lookup stays cached per model.
const CONTEXT_WINDOW_CACHE_TTL: Duration = Duration::from_secs(600);
/// `"<swap root>|<model>"` -> (n_ctx, fetched at). Process-wide because
/// clients are cloned per request.
fn context_window_cache() -> &'static std::sync::Mutex<HashMap<String, (usize, std::time::Instant)>>
{
static CACHE: std::sync::OnceLock<
std::sync::Mutex<HashMap<String, (usize, std::time::Instant)>>,
> = std::sync::OnceLock::new();
CACHE.get_or_init(Default::default)
}
/// OpenAI-compatible client targeting a llama-swap proxy in front of one or
/// more llama-server processes. See the module doc-comment for the slot model.
@@ -145,6 +158,86 @@ impl LlamaCppClient {
// upstream server directly, reached via llama-swap's passthrough at
// `{swap_root}/upstream/<tts_model>/voices`.
/// Per-request context window `model` is actually served with, read from
/// llama-server's `/props` (`default_generation_settings.n_ctx`, the
/// per-slot window) through llama-swap's passthrough. That request loads
/// the model if it isn't already, which is fine: callers ask right before
/// chatting with it. Successful lookups are cached for
/// `CONTEXT_WINDOW_CACHE_TTL` so a config.yaml edit is picked up without
/// a restart; failures aren't cached and return `None`.
pub async fn served_context_window(&self, model: &str) -> Option<usize> {
if let Some(n_ctx) = self.cached_context_window(model) {
return Some(n_ctx);
}
let key = format!("{}|{}", self.swap_root(), model);
let url = format!("{}/upstream/{}/props", self.swap_root(), model);
let props: Value = match self.client.get(&url).send().await {
Ok(resp) if resp.status().is_success() => resp.json().await.ok()?,
Ok(resp) => {
tracing::warn!(model, status = %resp.status(), "llama-swap props lookup failed");
return None;
}
Err(e) => {
tracing::warn!(model, error = %e, "llama-swap props lookup failed");
return None;
}
};
let n_ctx = props
.pointer("/default_generation_settings/n_ctx")
.or_else(|| props.get("n_ctx"))
.and_then(Value::as_u64)
.filter(|n| *n > 0)? as usize;
if let Ok(mut cache) = context_window_cache().lock() {
cache.insert(key, (n_ctx, std::time::Instant::now()));
}
Some(n_ctx)
}
/// Served context windows for `models`, for display, WITHOUT loading
/// anything: a model counts only if its window is cached from an
/// earlier lookup or llama-swap reports it as currently loaded (then
/// `/props` is answered by the running process). Others are absent.
pub async fn known_context_windows(&self, models: &[String]) -> HashMap<String, usize> {
let loaded: std::collections::HashSet<String> = async {
let url = format!("{}/models", self.base_url);
let parsed: Value = self.client.get(&url).send().await.ok()?.json().await.ok()?;
Some(
parsed
.get("data")?
.as_array()?
.iter()
.filter(|m| {
m.pointer("/status/value").and_then(Value::as_str) == Some("loaded")
})
.filter_map(|m| m.get("id")?.as_str().map(str::to_string))
.collect(),
)
}
.await
.unwrap_or_default();
let mut out = HashMap::new();
for model in models {
let cached = self.cached_context_window(model);
let n_ctx = match cached {
Some(n) => Some(n),
None if loaded.contains(model) => self.served_context_window(model).await,
None => None,
};
if let Some(n) = n_ctx {
out.insert(model.clone(), n);
}
}
out
}
fn cached_context_window(&self, model: &str) -> Option<usize> {
let key = format!("{}|{}", self.swap_root(), model);
let cache = context_window_cache().lock().ok()?;
let &(n_ctx, at) = cache.get(&key)?;
(at.elapsed() < CONTEXT_WINDOW_CACHE_TTL).then_some(n_ctx)
}
/// Root of the llama-swap proxy: `base_url` with a trailing `/v1` removed.
/// The `/upstream/...` passthrough lives here, not under `/v1`.
fn swap_root(&self) -> &str {
@@ -472,7 +565,8 @@ impl LlamaCppClient {
}
// num_ctx isn't an OpenAI param; llama-server bakes ctx in at launch
// via -c, so we silently drop the override here. The config.yaml
// entry is the source of truth for context size.
// entry is the source of truth for context size; chat budgets read
// it back via `context_window()`.
let _ = self.num_ctx;
// Reasoning toggle for thinking-capable templates (Qwen3 et al.).
// llama-server forwards chat_template_kwargs into the Jinja render
@@ -939,6 +1033,13 @@ impl LlmClient for LlamaCppClient {
fn primary_model(&self) -> &str {
&self.primary_model
}
/// llama-server fixes its window at launch (`-c`), so ask it via
/// llama-swap's `/upstream/<model>/props`. Cached per model; see
/// [`LlamaCppClient::served_context_window`].
async fn context_window(&self) -> Option<usize> {
self.served_context_window(&self.primary_model).await
}
}
impl LlamaCppClient {
@@ -955,6 +1056,7 @@ impl LlamaCppClient {
name,
has_vision: true,
has_tool_calling: true,
context_window: None,
}
}
}
@@ -1113,6 +1215,52 @@ fn image_to_data_url(img: &str) -> String {
mod tests {
use super::*;
/// One-shot HTTP server answering any request with `body`; returns its
/// base URL and a handle yielding the request line it received.
async fn serve_once(body: &'static str) -> (String, tokio::task::JoinHandle<String>) {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let handle = tokio::spawn(async move {
let (mut sock, _) = listener.accept().await.unwrap();
let mut buf = vec![0u8; 4096];
let n = sock.read(&mut buf).await.unwrap();
let req = String::from_utf8_lossy(&buf[..n])
.lines()
.next()
.unwrap_or("")
.to_string();
let resp = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
sock.write_all(resp.as_bytes()).await.unwrap();
req
});
(format!("http://{addr}/v1"), handle)
}
#[tokio::test]
async fn served_context_window_reads_props_through_upstream() {
let (base, req) =
serve_once(r#"{"default_generation_settings":{"n_ctx":32768},"total_slots":1}"#).await;
let client = LlamaCppClient::new(Some(base), Some("gemma4-31b-qat-uncensored".into()));
assert_eq!(client.context_window().await, Some(32768));
assert_eq!(
req.await.unwrap(),
"GET /upstream/gemma4-31b-qat-uncensored/props HTTP/1.1"
);
// Cached: a second lookup makes no request (the server is gone).
assert_eq!(client.context_window().await, Some(32768));
}
#[tokio::test]
async fn served_context_window_is_none_when_unreachable() {
let client = LlamaCppClient::new(Some("http://127.0.0.1:9/v1".into()), Some("nope".into()));
assert_eq!(client.context_window().await, None);
}
#[test]
fn tool_call_arguments_stringified_on_send() {
let msg = ChatMessage {
+13
View File
@@ -63,6 +63,14 @@ pub trait LlmClient: Send + Sync {
/// Primary model identifier this client was constructed with.
fn primary_model(&self) -> &str;
/// Context window (tokens) the server actually serves for the primary
/// model, when the backend can report one that overrides the request's
/// `num_ctx`. `None` = unknown or request-controlled (Ollama honours
/// `num_ctx` per request); callers fall back to the requested value.
async fn context_window(&self) -> Option<usize> {
None
}
}
/// Events emitted by streaming `chat_with_tools_stream`. A stream is a
@@ -174,6 +182,11 @@ pub struct ModelCapabilities {
pub name: String,
pub has_vision: bool,
pub has_tool_calling: bool,
/// Context window (tokens) the model is served with, when the server
/// fixes it regardless of the request's `num_ctx` (llama-swap) and it's
/// known without loading the model. Omitted otherwise.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub context_window: Option<usize>,
}
/// A single node in the conversation tree. Wraps ChatMessage with structural
+2
View File
@@ -315,6 +315,7 @@ impl OllamaClient {
name: model_name.to_string(),
has_vision,
has_tool_calling,
context_window: None,
})
}
@@ -346,6 +347,7 @@ impl OllamaClient {
name: model_name,
has_vision: false,
has_tool_calling: false,
context_window: None,
});
}
}
+1
View File
@@ -901,6 +901,7 @@ fn parse_model_capabilities(m: &Value) -> ModelCapabilities {
name,
has_vision,
has_tool_calling,
context_window: None,
}
}
+5 -6
View File
@@ -37,8 +37,8 @@ use uuid::Uuid;
use crate::ai::backend::{BackendKind, SamplingOverrides};
use crate::ai::handlers::ReplayQuery;
use crate::ai::insight_chat::{
BYTES_PER_TOKEN, ChatStreamEvent, DEFAULT_MAX_ITERATIONS, ForkInfo, RESPONSE_HEADROOM_TOKENS,
apply_context_budget, env_default_num_ctx, env_max_iterations, render_tree_path,
ChatStreamEvent, ContextFitter, DEFAULT_MAX_ITERATIONS, ForkInfo, env_max_iterations,
render_tree_path,
};
use crate::ai::insight_generator::InsightGenerator;
use crate::ai::llm_client::{ChatHistoryStore, ChatMessage, StoredChatNode};
@@ -842,10 +842,8 @@ impl PersonaChatSession {
// Trim to the model's context window before the turn. Without this a
// rolling persona transcript grows unbounded until the backend
// silently drops the front of it mid-conversation.
let budget_tokens = (req.num_ctx.unwrap_or_else(env_default_num_ctx) as usize)
.saturating_sub(RESPONSE_HEADROOM_TOKENS);
let budget_bytes = budget_tokens.saturating_mul(BYTES_PER_TOKEN);
let truncated = apply_context_budget(&mut messages, budget_bytes);
let mut fitter = ContextFitter::for_client(backend.chat(), req.num_ctx).await;
let truncated = fitter.trim(&mut messages);
if truncated {
let _ = entry.push_event(ChatStreamEvent::Truncated).await;
}
@@ -880,6 +878,7 @@ impl PersonaChatSession {
&persona.persona_id,
max_iterations,
&entry,
&mut fitter,
)
.await
.map_err(|e| anyhow!("persona chat loop failed: {e:?}"))?;