494 lines
17 KiB
Rust
494 lines
17 KiB
Rust
use serde_json::{json, Map, Value};
|
|
|
|
use crate::types::MessageRequest;
|
|
|
|
/// Incrementally-built JSON request body that caches per-message serialization
|
|
/// and avoids re-serializing the entire message list on every API call.
|
|
///
|
|
/// ## Typical workflow (per agentic-loop iteration)
|
|
/// 1. Build a fresh `MessageRequest` (or reuse the previous one with a new
|
|
/// message appended).
|
|
/// 2. Call `update(&request)` — only new/uncached messages are serialized.
|
|
/// 3. Call `build()` or `build_bytes()` to obtain the final body.
|
|
///
|
|
/// ## Base invalidation
|
|
/// The "base" portion (`model`, `max_tokens`, `system`, `tools`, `tool_choice`,
|
|
/// `stream`, tuning knobs) is cached until a field actually changes. Changes
|
|
/// are detected via a simplified content hash of the non-message fields.
|
|
///
|
|
/// ## Zero-alloc build\_bytes
|
|
/// Messages are cached as pre-serialized `Vec<u8>` so `build_bytes()` can
|
|
/// concatenate them directly into a single buffer without any intermediate
|
|
/// `Value` tree allocation.
|
|
#[derive(Debug, Clone)]
|
|
pub struct IncrementalBody {
|
|
/// Cached serialisation of the non-message fields (model, system, tools, …).
|
|
base: Option<Map<String, Value>>,
|
|
/// Per-message pre-serialised JSON bytes.
|
|
cached_message_bytes: Vec<Vec<u8>>,
|
|
/// Hash of the base-determining fields at the last rebuild.
|
|
base_hash: u64,
|
|
}
|
|
|
|
impl IncrementalBody {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
base: None,
|
|
cached_message_bytes: Vec::new(),
|
|
base_hash: 0,
|
|
}
|
|
}
|
|
|
|
/// Update the cache with a new request.
|
|
///
|
|
/// * If the base (non-message fields) changed → rebuild base.
|
|
/// * If messages grew (delta) → serialise only the new messages.
|
|
/// * If messages shrunk (e.g. after compaction) → truncate internal cache.
|
|
///
|
|
/// When `request.cached_message_values` is non-empty, cached JSON values
|
|
/// from that vector are used for delta messages, skipping re-serialisation.
|
|
pub fn update(&mut self, request: &MessageRequest) {
|
|
let new_hash = hash_base(request);
|
|
|
|
if self.base.is_none() || new_hash != self.base_hash {
|
|
self.base = Some(serialise_base(request));
|
|
self.base_hash = new_hash;
|
|
}
|
|
|
|
let msg_count = request.messages.len();
|
|
|
|
if msg_count > self.cached_message_bytes.len() {
|
|
let base_len = self.cached_message_bytes.len();
|
|
for (i, msg) in request.messages[base_len..]
|
|
.iter()
|
|
.enumerate()
|
|
{
|
|
let abs_idx = base_len + i;
|
|
let bytes: Vec<u8> = request
|
|
.cached_message_values
|
|
.get(abs_idx)
|
|
.and_then(|v| v.clone())
|
|
.map(|val| serde_json::to_vec(&val).unwrap_or_default())
|
|
.unwrap_or_else(|| serde_json::to_vec(msg).unwrap_or_default());
|
|
self.cached_message_bytes.push(bytes);
|
|
}
|
|
} else if msg_count < self.cached_message_bytes.len() {
|
|
self.cached_message_bytes.truncate(msg_count);
|
|
}
|
|
}
|
|
|
|
/// Build the full request body as a JSON `Value`.
|
|
///
|
|
/// Post-processing (image normalisation, system-prompt cache-control,
|
|
/// tools cache-control) must be applied separately if needed.
|
|
pub fn build(&self) -> Value {
|
|
let mut body = self.base.clone().unwrap_or_default();
|
|
body.insert(
|
|
"messages".to_string(),
|
|
Value::Array(
|
|
self.cached_message_bytes
|
|
.iter()
|
|
.map(|b| serde_json::from_slice(b).unwrap_or(Value::Null))
|
|
.collect(),
|
|
),
|
|
);
|
|
Value::Object(body)
|
|
}
|
|
|
|
/// Build the full request body as serialised JSON bytes.
|
|
///
|
|
/// Concatenates pre-serialised base fields and pre-serialised messages
|
|
/// directly into a single buffer — no intermediate `Value` trees are
|
|
/// allocated beyond the base fields that are stored as `Value`.
|
|
pub fn build_bytes(&self) -> Vec<u8> {
|
|
let mut buf = Vec::new();
|
|
buf.push(b'{');
|
|
|
|
let mut written = false;
|
|
if let Some(ref base) = self.base {
|
|
for (i, (key, val)) in base.iter().enumerate() {
|
|
if i > 0 {
|
|
buf.push(b',');
|
|
}
|
|
written = true;
|
|
append_json_string(&mut buf, key);
|
|
buf.push(b':');
|
|
append_json_value(&mut buf, val);
|
|
}
|
|
}
|
|
|
|
if written {
|
|
buf.push(b',');
|
|
}
|
|
buf.extend_from_slice(b"\"messages\":[");
|
|
for (i, msg_bytes) in self.cached_message_bytes.iter().enumerate() {
|
|
if i > 0 {
|
|
buf.push(b',');
|
|
}
|
|
buf.extend_from_slice(msg_bytes);
|
|
}
|
|
buf.push(b']');
|
|
|
|
buf.push(b'}');
|
|
buf
|
|
}
|
|
|
|
/// Clear the cache entirely (forces a full rebuild on next `update`).
|
|
pub fn invalidate(&mut self) {
|
|
self.base = None;
|
|
self.cached_message_bytes.clear();
|
|
self.base_hash = 0;
|
|
}
|
|
}
|
|
|
|
// ── Helpers ───────────────────────────────────────────────────────────────
|
|
|
|
/// Build a `Map` of only the non-message fields from a `MessageRequest`.
|
|
///
|
|
/// Unlike serialising the full `MessageRequest` and removing `"messages"`,
|
|
/// this constructs the map directly from individual fields — never
|
|
/// touching (let alone serialising) the potentially-large message vector.
|
|
fn serialise_base(request: &MessageRequest) -> Map<String, Value> {
|
|
let mut map = Map::new();
|
|
|
|
map.insert("model".into(), Value::String(request.model.clone()));
|
|
map.insert("max_tokens".into(), json!(request.max_tokens));
|
|
|
|
serialise_system_cache_control(&mut map, request.system.as_deref());
|
|
if !request.skip_tools {
|
|
serialise_tools_cache_control(&mut map, &request.tools);
|
|
}
|
|
|
|
if let Some(ref tc) = request.tool_choice {
|
|
map.insert("tool_choice".into(), serde_json::to_value(tc).unwrap_or_default());
|
|
}
|
|
|
|
if request.stream {
|
|
map.insert("stream".into(), Value::Bool(true));
|
|
}
|
|
|
|
if let Some(ref v) = request.temperature {
|
|
map.insert("temperature".into(), json!(v));
|
|
}
|
|
if let Some(ref v) = request.top_p {
|
|
map.insert("top_p".into(), json!(v));
|
|
}
|
|
// frequency_penalty and presence_penalty are not supported by Anthropic's
|
|
// /v1/messages endpoint, so we intentionally omit them here.
|
|
// `stop` is renamed to `stop_sequences` for Anthropic.
|
|
if let Some(ref v) = request.stop {
|
|
if !v.is_empty() {
|
|
map.insert("stop_sequences".into(), serde_json::to_value(v).unwrap_or_default());
|
|
}
|
|
}
|
|
if let Some(ref v) = request.reasoning_effort {
|
|
map.insert("reasoning_effort".into(), Value::String(v.clone()));
|
|
}
|
|
if let Some(ref v) = request.thinking {
|
|
map.insert("thinking".into(), serde_json::to_value(v).unwrap_or_default());
|
|
}
|
|
|
|
map
|
|
}
|
|
|
|
/// Split the flat system-prompt string at the dynamic boundary and emit
|
|
/// the Anthropic block array with `cache_control: ephemeral` on the static
|
|
/// portion. Mirrors `MessageRequest::apply_system_prompt_cache_control`.
|
|
fn serialise_system_cache_control(map: &mut Map<String, Value>, system: Option<&str>) {
|
|
let Some(system_str) = system.filter(|s| !s.is_empty()) else {
|
|
return;
|
|
};
|
|
let boundary = runtime::SYSTEM_PROMPT_DYNAMIC_BOUNDARY;
|
|
let blocks = if let Some(split_pos) = system_str.find(boundary) {
|
|
let static_part = system_str[..split_pos].trim_end();
|
|
let dynamic_part = system_str[split_pos + boundary.len()..].trim_start();
|
|
let mut blocks = Vec::new();
|
|
if !static_part.is_empty() {
|
|
blocks.push(serde_json::json!({
|
|
"type": "text",
|
|
"text": static_part,
|
|
"cache_control": { "type": "ephemeral" }
|
|
}));
|
|
}
|
|
if !dynamic_part.is_empty() {
|
|
// The dynamic portion changes every request, so a cache breakpoint
|
|
// here is useless and fragments the prefix cache. Only the static
|
|
// block above keeps `cache_control`.
|
|
blocks.push(serde_json::json!({
|
|
"type": "text",
|
|
"text": dynamic_part
|
|
}));
|
|
}
|
|
blocks
|
|
} else {
|
|
vec![serde_json::json!({
|
|
"type": "text",
|
|
"text": system_str,
|
|
"cache_control": { "type": "ephemeral" }
|
|
})]
|
|
};
|
|
if !blocks.is_empty() {
|
|
map.insert("system".into(), Value::Array(blocks));
|
|
}
|
|
}
|
|
|
|
/// Add `cache_control: ephemeral` to the last tool definition.
|
|
/// Mirrors `MessageRequest::apply_tools_cache_control`.
|
|
fn serialise_tools_cache_control(map: &mut Map<String, Value>, tools: &Option<Vec<crate::types::ToolDefinition>>) {
|
|
let Some(ref tools) = tools else {
|
|
return;
|
|
};
|
|
if tools.is_empty() {
|
|
return;
|
|
}
|
|
let mut values: Vec<Value> = Vec::with_capacity(tools.len());
|
|
for (i, tool) in tools.iter().enumerate() {
|
|
let mut val = serde_json::to_value(tool).unwrap_or_default();
|
|
if i == tools.len() - 1 {
|
|
if let Some(obj) = val.as_object_mut() {
|
|
obj.insert(
|
|
"cache_control".to_string(),
|
|
serde_json::json!({ "type": "ephemeral" }),
|
|
);
|
|
}
|
|
}
|
|
values.push(val);
|
|
}
|
|
map.insert("tools".into(), Value::Array(values));
|
|
}
|
|
|
|
/// Deterministic hash of the non-message fields so we can detect changes.
|
|
fn hash_base(request: &MessageRequest) -> u64 {
|
|
use std::hash::{Hash, Hasher};
|
|
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
|
request.model.hash(&mut hasher);
|
|
request.max_tokens.hash(&mut hasher);
|
|
request.system.hash(&mut hasher);
|
|
request.stream.hash(&mut hasher);
|
|
|
|
if let Some(ref tools) = request.tools {
|
|
for t in tools {
|
|
t.name.hash(&mut hasher);
|
|
}
|
|
}
|
|
request.tool_choice.hash(&mut hasher);
|
|
request.temperature.map(|v| v.to_bits()).hash(&mut hasher);
|
|
request.top_p.map(|v| v.to_bits()).hash(&mut hasher);
|
|
request.frequency_penalty.map(|v| v.to_bits()).hash(&mut hasher);
|
|
request.presence_penalty.map(|v| v.to_bits()).hash(&mut hasher);
|
|
request.stop.hash(&mut hasher);
|
|
request.reasoning_effort.hash(&mut hasher);
|
|
request.thinking.hash(&mut hasher);
|
|
request.skip_tools.hash(&mut hasher);
|
|
hasher.finish()
|
|
}
|
|
|
|
fn append_json_string(buf: &mut Vec<u8>, s: &str) {
|
|
buf.push(b'"');
|
|
for byte in s.bytes() {
|
|
match byte {
|
|
b'"' => buf.extend_from_slice(b"\\\""),
|
|
b'\\' => buf.extend_from_slice(b"\\\\"),
|
|
b'\n' => buf.extend_from_slice(b"\\n"),
|
|
b'\r' => buf.extend_from_slice(b"\\r"),
|
|
b'\t' => buf.extend_from_slice(b"\\t"),
|
|
0x08 => buf.extend_from_slice(b"\\b"),
|
|
0x0C => buf.extend_from_slice(b"\\f"),
|
|
c if c < 0x20 => {
|
|
write_hex_escape(buf, c);
|
|
}
|
|
c => buf.push(c),
|
|
}
|
|
}
|
|
buf.push(b'"');
|
|
}
|
|
|
|
fn write_hex_escape(buf: &mut Vec<u8>, byte: u8) {
|
|
const HEX: &[u8; 16] = b"0123456789abcdef";
|
|
buf.push(b'\\');
|
|
buf.push(b'u');
|
|
buf.push(b'0');
|
|
buf.push(b'0');
|
|
buf.push(HEX[(byte >> 4) as usize]);
|
|
buf.push(HEX[(byte & 0x0F) as usize]);
|
|
}
|
|
|
|
fn append_json_value(buf: &mut Vec<u8>, val: &Value) {
|
|
match val {
|
|
Value::Null => buf.extend_from_slice(b"null"),
|
|
Value::Bool(true) => buf.extend_from_slice(b"true"),
|
|
Value::Bool(false) => buf.extend_from_slice(b"false"),
|
|
Value::Number(n) => {
|
|
buf.extend_from_slice(n.to_string().as_bytes());
|
|
}
|
|
Value::String(s) => append_json_string(buf, s),
|
|
Value::Array(arr) => {
|
|
buf.push(b'[');
|
|
for (i, v) in arr.iter().enumerate() {
|
|
if i > 0 {
|
|
buf.push(b',');
|
|
}
|
|
append_json_value(buf, v);
|
|
}
|
|
buf.push(b']');
|
|
}
|
|
Value::Object(obj) => {
|
|
buf.push(b'{');
|
|
for (i, (key, val)) in obj.iter().enumerate() {
|
|
if i > 0 {
|
|
buf.push(b',');
|
|
}
|
|
append_json_string(buf, key);
|
|
buf.push(b':');
|
|
append_json_value(buf, val);
|
|
}
|
|
buf.push(b'}');
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::sync::Arc;
|
|
|
|
use crate::types::{InputMessage, ToolDefinition, ToolChoice};
|
|
|
|
use super::*;
|
|
|
|
fn sample_request(msg_count: usize) -> MessageRequest {
|
|
MessageRequest {
|
|
model: "claude-sonnet-4-6".to_string(),
|
|
max_tokens: 1024,
|
|
messages: Arc::new(
|
|
(0..msg_count)
|
|
.map(|i| InputMessage::user_text(format!("message {i}")))
|
|
.collect(),
|
|
),
|
|
system: Some(Arc::from("You are a helpful assistant.")),
|
|
tools: Some(vec![ToolDefinition {
|
|
name: "bash".to_string(),
|
|
description: Some("Run a shell command".to_string()),
|
|
input_schema: serde_json::json!({"type": "object"}),
|
|
}]),
|
|
tool_choice: Some(ToolChoice::Auto),
|
|
stream: true,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn full_build_produces_valid_json() {
|
|
let request = sample_request(3);
|
|
let mut body = IncrementalBody::new();
|
|
body.update(&request);
|
|
|
|
let value = body.build();
|
|
assert_eq!(value["model"], "claude-sonnet-4-6");
|
|
assert_eq!(value["max_tokens"], 1024);
|
|
// System prompt is now wrapped in cache_control array by serialise_base.
|
|
assert_eq!(
|
|
value["system"][0]["text"],
|
|
"You are a helpful assistant."
|
|
);
|
|
assert!(value.get("tools").is_some());
|
|
assert_eq!(
|
|
value["messages"].as_array().map(Vec::len),
|
|
Some(3)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn incremental_update_only_serialises_delta() {
|
|
let mut body = IncrementalBody::new();
|
|
|
|
let req1 = sample_request(2);
|
|
body.update(&req1);
|
|
assert_eq!(body.cached_message_bytes.len(), 2);
|
|
|
|
let req2 = sample_request(5);
|
|
body.update(&req2);
|
|
assert_eq!(body.cached_message_bytes.len(), 5);
|
|
|
|
let value = body.build();
|
|
assert_eq!(
|
|
value["messages"].as_array().map(Vec::len),
|
|
Some(5)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn truncation_handles_compaction() {
|
|
let mut body = IncrementalBody::new();
|
|
body.update(&sample_request(10));
|
|
assert_eq!(body.cached_message_bytes.len(), 10);
|
|
|
|
body.update(&sample_request(4));
|
|
assert_eq!(body.cached_message_bytes.len(), 4);
|
|
|
|
let value = body.build();
|
|
assert_eq!(
|
|
value["messages"].as_array().map(Vec::len),
|
|
Some(4)
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn base_hash_changes_on_model_switch() {
|
|
let mut body = IncrementalBody::new();
|
|
let req1 = sample_request(1);
|
|
|
|
body.update(&req1);
|
|
let hash1 = body.base_hash;
|
|
|
|
let mut req2 = sample_request(1);
|
|
req2.model = "claude-opus-4-6".to_string();
|
|
body.update(&req2);
|
|
|
|
assert_ne!(body.base_hash, hash1, "model change should alter base hash");
|
|
}
|
|
|
|
#[test]
|
|
fn build_bytes_round_trips() {
|
|
let request = sample_request(3);
|
|
let mut body = IncrementalBody::new();
|
|
body.update(&request);
|
|
|
|
let bytes = body.build_bytes();
|
|
let parsed: serde_json::Value =
|
|
serde_json::from_slice(&bytes).expect("build_bytes should be valid JSON");
|
|
|
|
assert_eq!(parsed["model"], "claude-sonnet-4-6");
|
|
assert_eq!(parsed["max_tokens"], 1024);
|
|
// System prompt is now wrapped in cache_control array by serialise_base.
|
|
assert_eq!(
|
|
parsed["system"][0]["text"],
|
|
"You are a helpful assistant."
|
|
);
|
|
assert_eq!(
|
|
parsed["messages"].as_array().map(Vec::len),
|
|
Some(3)
|
|
);
|
|
assert_eq!(
|
|
parsed["messages"][0]["content"][0]["text"],
|
|
"message 0"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn serialise_base_omits_messages() {
|
|
let request = sample_request(100);
|
|
let map = serialise_base(&request);
|
|
assert!(
|
|
!map.contains_key("messages"),
|
|
"serialise_base must not include the messages field"
|
|
);
|
|
assert_eq!(map.get("model").and_then(|v| v.as_str()), Some("claude-sonnet-4-6"));
|
|
assert_eq!(map.get("max_tokens").and_then(|v| v.as_u64()), Some(1024));
|
|
// System is now wrapped in cache_control array rather than flat string.
|
|
assert!(
|
|
map.get("system").and_then(|v| v.as_array()).is_some(),
|
|
"system should be a cache-controlled array"
|
|
);
|
|
}
|
|
}
|