diff --git a/_docs/prompt-caching.mdx b/_docs/prompt-caching.mdx index c9f9cde..132932d 100644 --- a/_docs/prompt-caching.mdx +++ b/_docs/prompt-caching.mdx @@ -30,14 +30,38 @@ What crabcode does today for each routing path: ## What gets marked (direct Anthropic) -Aligned with OpenCode's `auto` policy: +Hybrid of OpenCode `auto` + Grok Build placement: 1. **Last tool** — tool schemas are large and stable across a tool loop 2. **Last system block** — instructions / project context -3. **Latest user content block** — conversation prefix through the current turn (including tool-result groups) +3. **Transcript tip** — last markable content block (skips `thinking` / `redacted_thinking`) +4. **Previous user tip** — where the prior request ended, so turns past the 20-block lookback still hit cache + +At most **4** breakpoints; the 4th slot stays free when tools or previous-user are missing so gateways can auto-mark. Breakpoints are applied **after** message regrouping so adjacent tool_use / tool_result blocks stay valid. +### xAI Grok Build (cli-chat-proxy) + +In addition to sticky `prompt_cache_key = session_id`, crabcode stamps Grok Build–style affinity headers: + +| Header | Value | +| --- | --- | +| `x-grok-session-id` | Session id (sticky) | +| `x-grok-conv-id` | Same as session for main turns (sticky) | +| `x-grok-req-id` | Unique per model invocation | +| `x-grok-turn-idx` | 0-based user-turn index | +| `x-grok-agent-id` | Process-stable agent id | + +**Parent-cached aux** (e.g. max-steps text-only summary): keeps parent `prompt_cache_key` + session/conv so the conversation prefix can reuse the main turn's KV cache; assigns a fresh `aux-…` req id. + +**Subagents** use the **child** session id on purpose. They have a different system prompt and tool set, so parent-prefix reuse would miss and can pollute sticky routing. Isolation beats a false shared key. + +### Prefix stability (anti-bust) + +- **Empty system/user rows** are dropped before the wire (they pad the sticky prefix). +- **Images** use hysteresis: compact only above a trigger (~6 MiB total data-urls), then reclaim to a lower target (~3 MiB) so later turns stay cache-warm instead of re-evicting every step. + --- ## How to verify it works @@ -53,21 +77,27 @@ Logging writes to `app.log` in the working directory when `--emit-logs` is set. ### Direct Anthropic ``` -[prompt-cache] anthropic input=… output=… cache_read=… cache_creation=… +[prompt-cache] anthropic input=… output=… cache_read=… cache_creation=… total_input=… hit_pct=… ``` ### AI Gateway / OpenAI-compatible ``` -[prompt-cache] openai-compatible prompt=… completion=… cached_tokens=… cache_read=… cache_creation=… +[prompt-cache] openai-compatible prompt=… completion=… cached_tokens=… cache_read=… cache_creation=… hit_pct=… +``` + +### OpenAI / xAI Responses + +``` +[prompt-cache] openai-responses input=… output=… cached_tokens=… hit_pct=… ``` ### Healthy multi-step session | Step | What you want to see | | --- | --- | -| First request | `cache_creation` / write > 0, or full input billed once | -| Later tool steps (same tools + system + prefix) | `cache_read` / `cached_tokens` > 0 and growing | +| First request | `cache_creation` / write > 0, or full input billed once; `hit_pct` near 0 | +| Later tool steps (same tools + system + prefix) | `cache_read` / `cached_tokens` > 0 and `hit_pct` climbing (often 70%+) | Notes: diff --git a/src/agent/subagent.rs b/src/agent/subagent.rs index a09f070..3b3f3ff 100644 --- a/src/agent/subagent.rs +++ b/src/agent/subagent.rs @@ -47,7 +47,24 @@ pub async fn run_subagent( use futures::StreamExt; use std::collections::HashMap; - let session = resolve_subagent_session(&agent, parent_session, sender.as_ref()).await?; + let mut session = resolve_subagent_session(&agent, parent_session, sender.as_ref()).await?; + // Child cache key on purpose: subagents have a different system prompt and tool set, + // so parent-prefix reuse would miss (and risk sticky-routing pollution). See + // SessionAffinity::child_session docs. + session.prompt_cache_key = Some(session_id.clone()); + session.openai_options.prompt_cache_key = Some(session_id.clone()); + if crate::llm::xai_build::is_build_transport(&session.openai_options.additional_headers) { + let affinity = crate::llm::xai_build::SessionAffinity::child_session(&session_id); + crate::llm::xai_build::inject_session_affinity_headers( + &mut session.openai_options.additional_headers, + &affinity, + ); + crate::emit_log!( + "[prompt-cache] xai-build affinity kind=child session_id={} req_id={}", + affinity.session_id, + affinity.req_id + ); + } let scoped_registry = build_scoped_registry(full_registry, &agent).await; diff --git a/src/aisdk/providers/anthropic.rs b/src/aisdk/providers/anthropic.rs index 795ac62..83266f7 100644 --- a/src/aisdk/providers/anthropic.rs +++ b/src/aisdk/providers/anthropic.rs @@ -269,12 +269,24 @@ fn log_anthropic_usage(usage: &serde_json::Value) { return; } + let input_v = input.unwrap_or(0); + let total_input = input_v + .saturating_add(cache_read) + .saturating_add(cache_creation); + let hit_pct = if total_input > 0 { + (cache_read as f64 * 100.0) / total_input as f64 + } else { + 0.0 + }; + crate::emit_log!( - "[prompt-cache] anthropic input={} output={} cache_read={} cache_creation={}", + "[prompt-cache] anthropic input={} output={} cache_read={} cache_creation={} total_input={} hit_pct={:.1}", input.map(|v| v.to_string()).unwrap_or_else(|| "-".into()), output.map(|v| v.to_string()).unwrap_or_else(|| "-".into()), cache_read, - cache_creation + cache_creation, + total_input, + hit_pct ); } @@ -407,9 +419,13 @@ fn anthropic_tool_input_is_empty(value: &serde_json::Value) -> bool { } } -/// Anthropic prompt caching (opencode `auto` parity): -/// mark last tool + last system block + latest user content block. -/// Cap at 4 breakpoints; each is `{"type":"ephemeral"}`. +/// Anthropic prompt caching (Grok Build / OpenCode hybrid): +/// 1. last tool (stable schemas — high value in tool loops) +/// 2. last system block +/// 3. tip of transcript (last markable block; skips thinking) +/// 4. previous user tip (covers turns past the 20-block lookback) +/// Cap at 4 breakpoints; each is `{"type":"ephemeral"}`. The 4th slot stays +/// free when tools or previous-user are missing so gateways can still auto-mark. fn apply_anthropic_prompt_caching(body: &mut serde_json::Value) { let mut remaining = 4usize; @@ -439,27 +455,48 @@ fn apply_anthropic_prompt_caching(body: &mut serde_json::Value) { } } - // 3. Latest user message's last content block (after anthropic_messages - // regrouping so tool_result groups keep the marker). + // 3–4. Transcript tip + previous user tip (skip thinking blocks). + if remaining == 0 { + return; + } + let Some(messages) = body.get_mut("messages").and_then(|v| v.as_array_mut()) else { + return; + }; + + let tip = (0..messages.len()) + .rev() + .find(|&i| mark_message_cache_breakpoint(&mut messages[i])); + if tip.is_some() { + remaining = remaining.saturating_sub(1); + } + + // Where the previous request ended: skip the whole trailing user run after + // the last assistant, then mark that earlier user tip (Grok Build placement). if remaining > 0 { - if let Some(messages) = body.get_mut("messages").and_then(|v| v.as_array_mut()) { - if let Some(user) = messages - .iter_mut() - .rev() - .find(|m| m.get("role").and_then(|r| r.as_str()) == Some("user")) + if let Some(tip) = tip { + if let Some(prev) = messages[..tip] + .iter() + .rposition(|m| m.get("role").and_then(|r| r.as_str()) == Some("assistant")) + .and_then(|assistant| { + messages[..assistant] + .iter() + .rposition(|m| m.get("role").and_then(|r| r.as_str()) == Some("user")) + }) { - mark_latest_user_cache_control(user); + let _ = mark_message_cache_breakpoint(&mut messages[prev]); } } } } -fn mark_latest_user_cache_control(user: &mut serde_json::Value) { - let Some(obj) = user.as_object_mut() else { - return; +/// Marks the last content block that can carry a breakpoint, scanning back past +/// `thinking` / `redacted_thinking` which the API rejects. +fn mark_message_cache_breakpoint(message: &mut serde_json::Value) -> bool { + let Some(obj) = message.as_object_mut() else { + return false; }; - // String content must become a text block to host cache_control. + // Plain string content must become a text block to host cache_control. if let Some(serde_json::Value::String(text)) = obj.get("content").cloned() { obj.insert( "content".to_string(), @@ -469,17 +506,28 @@ fn mark_latest_user_cache_control(user: &mut serde_json::Value) { "cache_control": { "type": "ephemeral" } }]), ); - return; + return true; } - if let Some(blocks) = obj.get_mut("content").and_then(|c| c.as_array_mut()) { - if let Some(last) = blocks.last_mut().and_then(|b| b.as_object_mut()) { - last.insert( - "cache_control".to_string(), - serde_json::json!({ "type": "ephemeral" }), - ); + let Some(blocks) = obj.get_mut("content").and_then(|c| c.as_array_mut()) else { + return false; + }; + + for block in blocks.iter_mut().rev() { + let Some(block_obj) = block.as_object_mut() else { + continue; + }; + let block_type = block_obj.get("type").and_then(|t| t.as_str()).unwrap_or(""); + if block_type == "thinking" || block_type == "redacted_thinking" { + continue; } + block_obj.insert( + "cache_control".to_string(), + serde_json::json!({ "type": "ephemeral" }), + ); + return true; } + false } fn anthropic_user_content(user: &crate::message::UserMessage) -> serde_json::Value { @@ -738,7 +786,7 @@ mod tests { } #[test] - fn prompt_caching_marks_last_tool_system_and_latest_user() { + fn prompt_caching_marks_last_tool_system_tip_and_previous_user() { let mut body = serde_json::json!({ "system": [ { "type": "text", "text": "sys a" }, @@ -767,17 +815,15 @@ mod tests { body["system"][1]["cache_control"], serde_json::json!({ "type": "ephemeral" }) ); - // string user content is wrapped so cache_control can attach + // tip (latest user) wrapped with cache_control assert_eq!( body["messages"][2]["content"][0]["cache_control"], serde_json::json!({ "type": "ephemeral" }) ); - assert!( - body["messages"][0]["content"] - .as_str() - .map(|s| s == "first") - .unwrap_or(false) - || body["messages"][0].get("cache_control").is_none() + // previous user tip also marked (Grok Build placement) + assert_eq!( + body["messages"][0]["content"][0]["cache_control"], + serde_json::json!({ "type": "ephemeral" }) ); } @@ -803,6 +849,29 @@ mod tests { serde_json::json!({ "type": "ephemeral" }) ); } + + #[test] + fn prompt_caching_skips_thinking_blocks_when_marking_tip() { + let mut body = serde_json::json!({ + "messages": [{ + "role": "assistant", + "content": [ + { "type": "text", "text": "answer" }, + { "type": "thinking", "thinking": "secret" }, + ] + }], + }); + + apply_anthropic_prompt_caching(&mut body); + + assert_eq!( + body["messages"][0]["content"][0]["cache_control"], + serde_json::json!({ "type": "ephemeral" }) + ); + assert!(body["messages"][0]["content"][1] + .get("cache_control") + .is_none()); + } } fn anthropic_tool_output_content(tool: &crate::message::ToolOutputMessage) -> serde_json::Value { diff --git a/src/aisdk/providers/compatible.rs b/src/aisdk/providers/compatible.rs index 97d6161..700ca5f 100644 --- a/src/aisdk/providers/compatible.rs +++ b/src/aisdk/providers/compatible.rs @@ -425,15 +425,32 @@ fn log_openai_compatible_usage(usage: &serde_json::Value) { return; } + // Prefer OpenAI-style cached_tokens; fall back to Anthropic-style cache_read. + let effective_cached = if cached > 0 { cached } else { cache_read }; + let prompt_v = prompt.unwrap_or(0); + let hit_pct = if prompt_v > 0 { + (effective_cached as f64 * 100.0) / prompt_v as f64 + } else if effective_cached > 0 || cache_creation > 0 { + let total = effective_cached.saturating_add(cache_creation); + if total > 0 { + (effective_cached as f64 * 100.0) / total as f64 + } else { + 0.0 + } + } else { + 0.0 + }; + crate::emit_log!( - "[prompt-cache] openai-compatible prompt={} completion={} cached_tokens={} cache_read={} cache_creation={}", + "[prompt-cache] openai-compatible prompt={} completion={} cached_tokens={} cache_read={} cache_creation={} hit_pct={:.1}", prompt.map(|v| v.to_string()).unwrap_or_else(|| "-".into()), completion .map(|v| v.to_string()) .unwrap_or_else(|| "-".into()), cached, cache_read, - cache_creation + cache_creation, + hit_pct ); } diff --git a/src/aisdk/providers/openai.rs b/src/aisdk/providers/openai.rs index 759a339..6fa225c 100644 --- a/src/aisdk/providers/openai.rs +++ b/src/aisdk/providers/openai.rs @@ -1522,6 +1522,9 @@ fn response_sse_data_to_chunk(data: &str) -> Option> { return Some(Ok(responses_error_chunk(&value, event_type))); } let resp = &value["response"]; + if let Some(usage) = resp.get("usage") { + log_openai_responses_usage(usage); + } Some(Ok(ChunkType::ResponseCompleted { end_turn: resp.get("end_turn").and_then(|value| value.as_bool()), })) @@ -1544,6 +1547,44 @@ fn response_sse_data_to_chunk(data: &str) -> Option> { } } +/// Log Responses API usage for prompt-cache visibility. +/// Looks for `input_tokens_details.cached_tokens` (OpenAI/xAI shape). +fn log_openai_responses_usage(usage: &serde_json::Value) { + let input = usage + .get("input_tokens") + .or_else(|| usage.get("prompt_tokens")) + .and_then(|v| v.as_u64()); + let output = usage + .get("output_tokens") + .or_else(|| usage.get("completion_tokens")) + .and_then(|v| v.as_u64()); + let cached = usage + .pointer("/input_tokens_details/cached_tokens") + .or_else(|| usage.pointer("/prompt_tokens_details/cached_tokens")) + .and_then(|v| v.as_u64()) + .or_else(|| usage.get("cached_tokens").and_then(|v| v.as_u64())) + .unwrap_or(0); + + if input.is_none() && output.is_none() && cached == 0 { + return; + } + + let input_v = input.unwrap_or(0); + let hit_pct = if input_v > 0 { + (cached as f64 * 100.0) / input_v as f64 + } else { + 0.0 + }; + + crate::emit_log!( + "[prompt-cache] openai-responses input={} output={} cached_tokens={} hit_pct={:.1}", + input.map(|v| v.to_string()).unwrap_or_else(|| "-".into()), + output.map(|v| v.to_string()).unwrap_or_else(|| "-".into()), + cached, + hit_pct + ); +} + fn responses_provider_error_message(value: &serde_json::Value, fallback: &str) -> String { let code = response_error_field(value, "code"); let message = response_error_field(value, "message"); diff --git a/src/aisdk/response.rs b/src/aisdk/response.rs index 191a040..8b87a46 100644 --- a/src/aisdk/response.rs +++ b/src/aisdk/response.rs @@ -17,12 +17,23 @@ const PROVIDER_STEP_MAX_RETRIES: usize = 10; /// Keep the newest N tool outputs intact for the model; older ones are pruned. /// Matches Grok Build / OpenCode mid-session tool-result retention behavior. const KEEP_RECENT_TOOL_OUTPUTS: usize = 6; +/// Tool outputs older than this many user turns are hard-cleared (turn-age gate). +const KEEP_RECENT_USER_TURNS: usize = 2; /// Soft-trim threshold for older-but-still-retained tool outputs (chars). const TOOL_OUTPUT_SOFT_TRIM_CHARS: usize = 4_000; const TOOL_OUTPUT_SOFT_TRIM_HEAD: usize = 1_500; const TOOL_OUTPUT_SOFT_TRIM_TAIL: usize = 1_500; const PRUNED_TOOL_OUTPUT_PLACEHOLDER: &str = "[Old tool result content cleared]"; +/// Image compact hysteresis (Grok Build-inspired): +/// - Gate eviction only when total image payload exceeds the **trigger** +/// - Once firing, reclaim down to the lower **target** so the next few turns +/// stay under the ceiling (avoids re-busting the KV prefix every step) +const IMAGE_COMPACT_TRIGGER_BYTES: usize = 6 * 1024 * 1024; +const IMAGE_COMPACT_RECLAIM_TARGET_BYTES: usize = 3 * 1024 * 1024; +const _: () = assert!(IMAGE_COMPACT_RECLAIM_TARGET_BYTES < IMAGE_COMPACT_TRIGGER_BYTES); +const IMAGE_COMPACT_PLACEHOLDER: &str = "[An earlier image was removed to keep the request within its size limit and is no longer visible. Do not describe or reason about its contents from memory; ask the user to re-share it if you need to see it again.]"; + pub struct StreamTextResponse { pub stream: LanguageModelStream, stop_reason: Arc>>, @@ -139,8 +150,15 @@ pub async fn stream_with_tools( let pruned = prune_stale_tool_outputs_in_place(&mut current_messages); if pruned > 0 { let _ = tx_loop.send(ChunkType::Metadata(format!( - "tool_outputs_pruned count={} keep_recent={}", - pruned, KEEP_RECENT_TOOL_OUTPUTS + "tool_outputs_pruned count={} keep_recent={} keep_user_turns={}", + pruned, KEEP_RECENT_TOOL_OUTPUTS, KEEP_RECENT_USER_TURNS + ))); + } + let images_evicted = compact_images_to_budget_in_place(&mut current_messages); + if images_evicted > 0 { + let _ = tx_loop.send(ChunkType::Metadata(format!( + "images_compacted evicted={} trigger_bytes={} reclaim_target_bytes={}", + images_evicted, IMAGE_COMPACT_TRIGGER_BYTES, IMAGE_COMPACT_RECLAIM_TARGET_BYTES ))); } let _ = tx_loop.send(ChunkType::Metadata(format!( @@ -781,7 +799,8 @@ fn rollback_provider_attempt( /// Strategy (Grok Build / OpenCode inspired): /// - Keep the newest [`KEEP_RECENT_TOOL_OUTPUTS`] results full-size /// - Soft-trim large older results to head+tail -/// - Hard-clear anything older than 2× the keep window to a placeholder +/// - Hard-clear anything older than 2× the keep window **or** older than +/// [`KEEP_RECENT_USER_TURNS`] user turns /// - Drop attached images on pruned outputs (base64 is extremely expensive) fn prune_stale_tool_outputs_in_place(messages: &mut [Message]) -> usize { let tool_output_indices: Vec = messages @@ -794,6 +813,24 @@ fn prune_stale_tool_outputs_in_place(messages: &mut [Message]) -> usize { return 0; } + // Turn age: count User messages from the end. Tools in older turns are + // cheaper to drop even if they still fall inside the rank keep window. + let mut turn_from_end_by_index = vec![0usize; messages.len()]; + { + let mut turn_from_end = 0usize; + let mut seen_user = false; + for i in (0..messages.len()).rev() { + if matches!(&messages[i], Message::User(_)) { + if seen_user { + turn_from_end += 1; + } else { + seen_user = true; + } + } + turn_from_end_by_index[i] = turn_from_end; + } + } + let keep_from = tool_output_indices .len() .saturating_sub(KEEP_RECENT_TOOL_OUTPUTS); @@ -803,17 +840,24 @@ fn prune_stale_tool_outputs_in_place(messages: &mut [Message]) -> usize { let mut pruned = 0usize; for (rank, &idx) in tool_output_indices.iter().enumerate() { - if rank >= keep_from { + let turn_age = turn_from_end_by_index[idx]; + let in_rank_keep = rank >= keep_from; + let in_turn_keep = turn_age < KEEP_RECENT_USER_TURNS; + // Keep full only when recent by tool-count *and* by user-turn. + if in_rank_keep && in_turn_keep { continue; } + let hard_by_rank = rank < hard_clear_before; + let hard_by_turn = !in_turn_keep; + let Message::ToolOutput(output) = &mut messages[idx] else { continue; }; let had_images = !output.images.is_empty(); let original_len = output.output.len(); - if rank < hard_clear_before { + if hard_by_rank || hard_by_turn { if original_len > PRUNED_TOOL_OUTPUT_PLACEHOLDER.len() || had_images { output.output = PRUNED_TOOL_OUTPUT_PLACEHOLDER.to_string(); output.images.clear(); @@ -833,6 +877,88 @@ fn prune_stale_tool_outputs_in_place(messages: &mut [Message]) -> usize { pruned } +/// Evict oldest inline images with hysteresis: +/// fire only above [`IMAGE_COMPACT_TRIGGER_BYTES`], reclaim to +/// [`IMAGE_COMPACT_RECLAIM_TARGET_BYTES`]. +/// Returns how many images were replaced with a text placeholder. +fn compact_images_to_budget_in_place(messages: &mut [Message]) -> usize { + let mut total = total_image_bytes(messages); + // Below trigger: leave every image in place so the KV prefix stays byte-stable. + if total <= IMAGE_COMPACT_TRIGGER_BYTES { + return 0; + } + + let mut evicted = 0usize; + // Collect (message_index, is_user, image_index, bytes) oldest-first. + let mut slots: Vec<(usize, bool, usize, usize)> = Vec::new(); + for (msg_idx, message) in messages.iter().enumerate() { + match message { + Message::User(user) => { + for (img_idx, image) in user.images.iter().enumerate() { + slots.push((msg_idx, true, img_idx, image.data_url.len())); + } + } + Message::ToolOutput(output) => { + for (img_idx, image) in output.images.iter().enumerate() { + slots.push((msg_idx, false, img_idx, image.data_url.len())); + } + } + _ => {} + } + } + + // Evict from oldest message first (already in conversation order). + // Within a message, drop higher image indices first so removals don't shift lower ones. + slots.sort_by(|a, b| a.0.cmp(&b.0).then(b.2.cmp(&a.2))); + + for (msg_idx, is_user, img_idx, bytes) in slots { + // Reclaim past the trigger down to the low-water mark (hysteresis). + if total <= IMAGE_COMPACT_RECLAIM_TARGET_BYTES { + break; + } + let removed = match &mut messages[msg_idx] { + Message::User(user) if is_user && img_idx < user.images.len() => { + user.images.remove(img_idx); + if user.content.trim().is_empty() { + user.content = IMAGE_COMPACT_PLACEHOLDER.to_string(); + } else if !user.content.contains(IMAGE_COMPACT_PLACEHOLDER) { + user.content.push_str("\n\n"); + user.content.push_str(IMAGE_COMPACT_PLACEHOLDER); + } + true + } + Message::ToolOutput(output) if !is_user && img_idx < output.images.len() => { + output.images.remove(img_idx); + if !output.output.contains(IMAGE_COMPACT_PLACEHOLDER) { + if !output.output.is_empty() { + output.output.push_str("\n\n"); + } + output.output.push_str(IMAGE_COMPACT_PLACEHOLDER); + } + true + } + _ => false, + }; + if removed { + total = total.saturating_sub(bytes); + evicted += 1; + } + } + + evicted +} + +fn total_image_bytes(messages: &[Message]) -> usize { + messages + .iter() + .map(|message| match message { + Message::User(user) => user.images.iter().map(|img| img.data_url.len()).sum(), + Message::ToolOutput(output) => output.images.iter().map(|img| img.data_url.len()).sum(), + _ => 0usize, + }) + .sum() +} + fn soft_trim_tool_output(text: &str) -> String { let char_count = text.chars().count(); if char_count <= TOOL_OUTPUT_SOFT_TRIM_CHARS { @@ -1406,9 +1532,10 @@ fn tool_call_key(item: &serde_json::Value, array_index: usize) -> String { #[cfg(test)] mod tests { use super::{ - prune_stale_tool_outputs_in_place, soft_trim_tool_output, stream_with_tools, - ToolCallAccumulator, KEEP_RECENT_TOOL_OUTPUTS, PRUNED_TOOL_OUTPUT_PLACEHOLDER, - TOOL_OUTPUT_SOFT_TRIM_CHARS, + compact_images_to_budget_in_place, prune_stale_tool_outputs_in_place, + soft_trim_tool_output, stream_with_tools, total_image_bytes, ToolCallAccumulator, + IMAGE_COMPACT_PLACEHOLDER, IMAGE_COMPACT_RECLAIM_TARGET_BYTES, IMAGE_COMPACT_TRIGGER_BYTES, + KEEP_RECENT_TOOL_OUTPUTS, PRUNED_TOOL_OUTPUT_PLACEHOLDER, TOOL_OUTPUT_SOFT_TRIM_CHARS, }; use crate::chunk::{ChunkType, FinishReason, MessagePhase}; use crate::message::Message; @@ -1472,6 +1599,125 @@ mod tests { assert_eq!(last.len(), TOOL_OUTPUT_SOFT_TRIM_CHARS + 200); } + #[test] + fn prune_stale_tool_outputs_hard_clears_by_user_turn_age() { + // Few tools overall (would stay inside rank keep window) but spanning + // many user turns — turn-age gate should still hard-clear the oldest. + let mut messages = vec![ + Message::user("u0"), + Message::tool_output( + "c0", + "bash", + "x".repeat(TOOL_OUTPUT_SOFT_TRIM_CHARS + 50), + false, + ), + Message::user("u1"), + Message::tool_output( + "c1", + "bash", + "x".repeat(TOOL_OUTPUT_SOFT_TRIM_CHARS + 50), + false, + ), + Message::user("u2"), + Message::tool_output( + "c2", + "bash", + "x".repeat(TOOL_OUTPUT_SOFT_TRIM_CHARS + 50), + false, + ), + Message::user("u3"), + Message::tool_output( + "c3", + "bash", + "x".repeat(TOOL_OUTPUT_SOFT_TRIM_CHARS + 50), + false, + ), + ]; + + let pruned = prune_stale_tool_outputs_in_place(&mut messages); + assert!(pruned > 0); + + let first_tool = messages + .iter() + .find_map(|m| match m { + Message::ToolOutput(o) if o.call_id == "c0" => Some(o.output.as_str()), + _ => None, + }) + .expect("c0"); + assert_eq!(first_tool, PRUNED_TOOL_OUTPUT_PLACEHOLDER); + + let last_tool = messages + .iter() + .find_map(|m| match m { + Message::ToolOutput(o) if o.call_id == "c3" => Some(o.output.as_str()), + _ => None, + }) + .expect("c3"); + assert_eq!(last_tool.len(), TOOL_OUTPUT_SOFT_TRIM_CHARS + 50); + } + + #[test] + fn compact_images_evicts_oldest_when_over_trigger_and_reclaims_to_target() { + use crate::message::ImageContent; + + // Two oversized images: each ~4 MiB → total ~8 MiB > 6 MiB trigger. + // Hysteresis reclaims to 3 MiB → must drop both (4 MiB still over target). + let big = "x".repeat(4 * 1024 * 1024); + let mut messages = vec![ + Message::user_with_images( + "first", + vec![ImageContent { + data_url: format!("data:image/png;base64,{big}"), + media_type: "image/png".to_string(), + }], + ), + Message::user_with_images( + "second", + vec![ImageContent { + data_url: format!("data:image/png;base64,{big}"), + media_type: "image/png".to_string(), + }], + ), + ]; + + let before = total_image_bytes(&messages); + assert!(before > IMAGE_COMPACT_TRIGGER_BYTES); + + let evicted = compact_images_to_budget_in_place(&mut messages); + assert!(evicted >= 1); + assert!(total_image_bytes(&messages) <= IMAGE_COMPACT_RECLAIM_TARGET_BYTES); + + // Oldest message should lose its image and gain a placeholder note. + match &messages[0] { + Message::User(user) => { + assert!(user.images.is_empty()); + assert!(user.content.contains(IMAGE_COMPACT_PLACEHOLDER)); + } + _ => panic!("expected user"), + } + } + + #[test] + fn compact_images_is_noop_below_trigger() { + use crate::message::ImageContent; + + // ~2 MiB total — under 6 MiB trigger → leave alone (prefix-stable). + let med = "x".repeat(2 * 1024 * 1024); + let mut messages = vec![Message::user_with_images( + "one", + vec![ImageContent { + data_url: format!("data:image/png;base64,{med}"), + media_type: "image/png".to_string(), + }], + )]; + assert!(total_image_bytes(&messages) <= IMAGE_COMPACT_TRIGGER_BYTES); + assert_eq!(compact_images_to_budget_in_place(&mut messages), 0); + match &messages[0] { + Message::User(user) => assert_eq!(user.images.len(), 1), + _ => panic!("expected user"), + } + } + #[derive(Debug, Clone)] struct BlockingTextProvider; diff --git a/src/llm/client.rs b/src/llm/client.rs index 2f5d057..9a83f00 100644 --- a/src/llm/client.rs +++ b/src/llm/client.rs @@ -446,6 +446,8 @@ pub async fn stream_llm_with_cancellation( request_config.supports_image_input, show_vlm_agent_hint, ); + // Stamp Build affinity *after* message conversion so turn_idx matches wire content. + stamp_build_main_turn_affinity(&mut request_config, &session_id, &aisdk_messages); let mut aisdk_tools = convert_to_aisdk_tools( &tool_registry, @@ -544,18 +546,21 @@ pub async fn stream_llm_with_cancellation( let mut follow_up_messages = response.messages().await; follow_up_messages.push(AisdkMessage::assistant(MAX_STEPS_REACHED_PROMPT)); - let summary_message_count = follow_up_messages.len(); + // Parent-cached aux: same conversation prefix + sticky session key, fresh req id. + // Tools are empty (text-only summary) so tool-schema cache may miss; system+history still hit. + let mut summary_config = request_config.clone(); + stamp_build_parent_cached_aux(&mut summary_config, &session_id, &follow_up_messages); let summary_log_context = StreamLogContext::new( "max_steps_summary", - &request_config, - summary_message_count, + &summary_config, + follow_up_messages.len(), 0, None, ); - log_stream_request(summary_log_context, &request_config); + log_stream_request(summary_log_context, &summary_config); let mut summary_response = stream_provider_request( - &request_config, + &summary_config, follow_up_messages, Vec::new(), None, @@ -1238,6 +1243,60 @@ fn send_warning(sender: &crate::llm::ChunkSender, warning: impl Into) { let _ = sender.send(crate::llm::ChunkMessage::Warning(warning.into())); } +/// Stamp sticky Build affinity for a main agent turn (session == conv). +fn stamp_build_main_turn_affinity( + request_config: &mut ProviderRequestConfig, + session_id: &str, + messages: &[AisdkMessage], +) { + if !super::xai_build::is_build_transport(&request_config.openai_options.additional_headers) { + return; + } + let turn_idx = super::xai_build::user_turn_idx_from_aisdk_messages(messages); + let affinity = super::xai_build::SessionAffinity::main_turn(session_id, turn_idx); + super::xai_build::inject_session_affinity_headers( + &mut request_config.openai_options.additional_headers, + &affinity, + ); + crate::emit_log!( + "[prompt-cache] xai-build affinity kind=main session_id={} conv_id={} req_id={} turn_idx={} agent_id={}", + affinity.session_id, + affinity.conv_id, + affinity.req_id, + turn_idx, + affinity.agent_id.as_deref().unwrap_or("-") + ); +} + +/// Stamp parent-cached affinity for continuation aux (e.g. max-steps text summary). +/// +/// Keeps `prompt_cache_key` / session / conv on the parent so the wire prefix can +/// reuse the main turn's cached KV; assigns a fresh `req_id` for telemetry. +fn stamp_build_parent_cached_aux( + request_config: &mut ProviderRequestConfig, + parent_session_id: &str, + messages: &[AisdkMessage], +) { + request_config.openai_options.prompt_cache_key = Some(parent_session_id.to_string()); + if !super::xai_build::is_build_transport(&request_config.openai_options.additional_headers) { + return; + } + let turn_idx = super::xai_build::user_turn_idx_from_aisdk_messages(messages); + let affinity = + super::xai_build::SessionAffinity::parent_cached_aux(parent_session_id, turn_idx); + super::xai_build::inject_session_affinity_headers( + &mut request_config.openai_options.additional_headers, + &affinity, + ); + crate::emit_log!( + "[prompt-cache] xai-build affinity kind=parent_aux session_id={} conv_id={} req_id={} turn_idx={}", + affinity.session_id, + affinity.conv_id, + affinity.req_id, + turn_idx + ); +} + async fn stream_provider_request( config: &ProviderRequestConfig, messages: Vec, @@ -1718,9 +1777,13 @@ fn convert_messages_for_model( match msg.role { crate::session::types::MessageRole::System => { - aisdk_messages.push(AisdkMessage::system( - crate::utils::sanitize::strip_legacy_image_descriptions(&msg.content), - )); + // Skip empty system rows — they pad the sticky prefix and bust cache + // (Grok Build: "resumed sessions no longer pad the sticky prompt with empty rows"). + let content = crate::utils::sanitize::strip_legacy_image_descriptions(&msg.content); + if content.trim().is_empty() { + continue; + } + aisdk_messages.push(AisdkMessage::system(content)); } crate::session::types::MessageRole::User => { let content = crate::utils::sanitize::strip_legacy_image_descriptions(&msg.content); @@ -1763,6 +1826,11 @@ fn convert_messages_for_model( }) .collect::>(); + // Empty user rows without images also pad the sticky prefix. + if content.trim().is_empty() && images.is_empty() { + continue; + } + if images.is_empty() { aisdk_messages.push(AisdkMessage::user(content)); } else { diff --git a/src/llm/xai_build.rs b/src/llm/xai_build.rs index ebfb2ea..4504b1c 100644 --- a/src/llm/xai_build.rs +++ b/src/llm/xai_build.rs @@ -97,6 +97,146 @@ fn request_overrides_with_version( } } +/// True when this request is routed through the Grok Build cli-chat-proxy transport. +pub(crate) fn is_build_transport(headers: &std::collections::HashMap) -> bool { + headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case(TOKEN_AUTH_HEADER) && value == TOKEN_AUTH_VALUE + }) +} + +/// Identity stamped on every cli-chat-proxy request for sticky routing / cache affinity. +/// +/// Mirrors Grok Build's main-turn + side-call header set (minus deployment/user which +/// require managed-enterprise context we may not have). +#[derive(Debug, Clone)] +pub(crate) struct SessionAffinity { + /// Sticky session id (`x-grok-session-id`). + pub session_id: String, + /// Sticky conversation id (`x-grok-conv-id`). Usually equals `session_id`. + /// Side/aux calls that must ride a parent prefix use the **parent** session id here + /// (and as `prompt_cache_key`) even if telemetry wants a different req id. + pub conv_id: String, + /// Unique per model invocation (`x-grok-req-id`). + pub req_id: String, + /// 0-based user-turn index within the conversation (`x-grok-turn-idx`). + pub turn_idx: Option, + /// Process-stable agent id (`x-grok-agent-id`). Defaults to [`process_agent_id`]. + pub agent_id: Option, +} + +impl SessionAffinity { + /// Main-turn / tool-loop affinity: session == conv, fresh req id. + pub fn main_turn(session_id: impl Into, turn_idx: u32) -> Self { + let session_id = session_id.into(); + Self { + conv_id: session_id.clone(), + session_id, + req_id: new_req_id(), + turn_idx: Some(turn_idx), + agent_id: Some(process_agent_id()), + } + } + + /// Side/aux call that reuses a parent conversation's cached prefix. + /// + /// `prompt_cache_key` and sticky headers stay on `parent_session_id`; `req_id` + /// is unique so telemetry can still distinguish the aux call. + pub fn parent_cached_aux(parent_session_id: impl Into, turn_idx: u32) -> Self { + let parent = parent_session_id.into(); + Self { + session_id: parent.clone(), + conv_id: parent, + req_id: format!("aux-{}", new_req_id()), + turn_idx: Some(turn_idx), + agent_id: Some(process_agent_id()), + } + } + + /// Child session (subagent) — intentionally **not** parent-cached: different + /// system prompt / tool set would miss the parent prefix and can pollute routing. + pub fn child_session(child_session_id: impl Into) -> Self { + let session_id = child_session_id.into(); + Self { + conv_id: session_id.clone(), + session_id, + req_id: new_req_id(), + turn_idx: Some(0), + agent_id: Some(process_agent_id()), + } + } +} + +/// Sticky session affinity headers for cli-chat-proxy prompt-cache routing. +/// +/// `prompt_cache_key` remains a separate body field and should also be set to the +/// sticky session/conv id for Responses sticky routing. +pub(crate) fn inject_session_affinity_headers( + headers: &mut std::collections::HashMap, + affinity: &SessionAffinity, +) { + if affinity.session_id.is_empty() { + return; + } + headers.insert("x-grok-session-id".to_string(), affinity.session_id.clone()); + let conv = if affinity.conv_id.is_empty() { + affinity.session_id.as_str() + } else { + affinity.conv_id.as_str() + }; + headers.insert("x-grok-conv-id".to_string(), conv.to_string()); + if !affinity.req_id.is_empty() { + headers.insert("x-grok-req-id".to_string(), affinity.req_id.clone()); + } + if let Some(turn_idx) = affinity.turn_idx { + headers.insert("x-grok-turn-idx".to_string(), turn_idx.to_string()); + } + let agent_id = affinity.agent_id.clone().unwrap_or_else(process_agent_id); + if !agent_id.is_empty() { + headers.insert("x-grok-agent-id".to_string(), agent_id); + } +} + +/// Convenience for callers that only have session + req (tests / simple paths). +pub(crate) fn inject_session_affinity_headers_simple( + headers: &mut std::collections::HashMap, + session_id: &str, + req_id: &str, +) { + inject_session_affinity_headers( + headers, + &SessionAffinity { + session_id: session_id.to_string(), + conv_id: session_id.to_string(), + req_id: req_id.to_string(), + turn_idx: None, + agent_id: Some(process_agent_id()), + }, + ); +} + +pub(crate) fn new_req_id() -> String { + format!("crabcode-{}", cuid2::create_id()) +} + +/// Process-stable agent id (Grok Build `x-grok-agent-id` counterpart). +/// +/// Not persisted across restarts — good enough for proxy affinity within a run. +pub(crate) fn process_agent_id() -> String { + static AGENT_ID: OnceLock = OnceLock::new(); + AGENT_ID + .get_or_init(|| format!("crabcode-agent-{}", cuid2::create_id())) + .clone() +} + +/// Count 0-based user turns in a converted model message list (for `x-grok-turn-idx`). +pub(crate) fn user_turn_idx_from_aisdk_messages(messages: &[crate::aisdk::core::Message]) -> u32 { + messages + .iter() + .filter(|m| matches!(m, crate::aisdk::core::Message::User(_))) + .count() + .saturating_sub(1) as u32 +} + pub(crate) async fn protocol_version() -> String { if let Some(version) = configured_protocol_version() { return version; @@ -225,7 +365,9 @@ mod tests { let headers = HashMap::from([(TOKEN_AUTH_HEADER.to_string(), TOKEN_AUTH_VALUE.to_string())]); assert!(retry_policy_for(&headers).is_some()); + assert!(super::is_build_transport(&headers)); assert!(retry_policy_for(&HashMap::new()).is_none()); + assert!(!super::is_build_transport(&HashMap::new())); } #[test] @@ -258,6 +400,56 @@ mod tests { ); } + #[test] + fn session_affinity_headers_are_sticky_and_req_scoped() { + let mut headers = HashMap::new(); + super::inject_session_affinity_headers_simple(&mut headers, "sess-1", "req-abc"); + assert_eq!( + headers.get("x-grok-session-id").map(String::as_str), + Some("sess-1") + ); + assert_eq!( + headers.get("x-grok-conv-id").map(String::as_str), + Some("sess-1") + ); + assert_eq!( + headers.get("x-grok-req-id").map(String::as_str), + Some("req-abc") + ); + assert!(headers.get("x-grok-agent-id").is_some()); + + // Empty session id is a no-op (never stamp blank affinity). + let mut empty = HashMap::new(); + super::inject_session_affinity_headers_simple(&mut empty, "", "req"); + assert!(empty.is_empty()); + } + + #[test] + fn main_turn_affinity_includes_turn_and_agent() { + let affinity = super::SessionAffinity::main_turn("sess-9", 3); + let mut headers = HashMap::new(); + super::inject_session_affinity_headers(&mut headers, &affinity); + assert_eq!( + headers.get("x-grok-turn-idx").map(String::as_str), + Some("3") + ); + assert_eq!( + headers.get("x-grok-agent-id").map(String::as_str), + Some(super::process_agent_id().as_str()) + ); + assert!(headers + .get("x-grok-req-id") + .is_some_and(|id| id.starts_with("crabcode-"))); + } + + #[test] + fn parent_cached_aux_keeps_parent_ids_and_unique_req() { + let affinity = super::SessionAffinity::parent_cached_aux("parent-sess", 1); + assert_eq!(affinity.session_id, "parent-sess"); + assert_eq!(affinity.conv_id, "parent-sess"); + assert!(affinity.req_id.starts_with("aux-")); + } + #[tokio::test] async fn retry_policy_ignores_non_upgrade_responses() { let policy = XaiBuildRetryPolicy;