use std::{
collections::{HashMap, VecDeque},
num::NonZeroUsize,
sync::{
Arc,
Mutex,
atomic::{AtomicU64, AtomicUsize, Ordering},
},
time::Duration,
};
use axum::{
body::Body,
extract::State,
http::{Method, StatusCode},
response::{IntoResponse, Json, Response},
};
use bytes::Bytes;
use reqwest::Client as HttpClient;
use serde::{Deserialize, Serialize};
use headroom_core::ccr::{
CcrStore,
backends::{in_memory::InMemoryCcrStore, sqlite::SqliteCcrStore},
compute_key,
};
use tokio_util::task::TaskTracker;
use futures::StreamExt;
#[derive(Clone)]
pub struct Secret(pub(crate) String);
impl std::fmt::Debug for Secret {
fn fmt(&self, f:&mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "[REDACTED]") }
}
impl std::fmt::Display for Secret {
fn fmt(&self, f:&mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "[REDACTED]") }
}
impl Secret {
pub fn expose(&self) -> &str { &self.0 }
}
impl From<&str> for Secret {
fn from(s:&str) -> Self { Secret(s.to_string()) }
}
impl From<String> for Secret {
fn from(s:String) -> Self { Secret(s) }
}
use crate::config::{Cli, CompressionConfig, ProxyMode, env_parse_warn};
const CACHE_COMPRESS_THRESHOLD:usize = 8192;
const TOKEN_COMPRESS_THRESHOLD:usize = 1024;
const INLINE_CCR_THRESHOLD:usize = 256;
const CHAT_COMPLETIONS_PATH:&str = "/v1/chat/completions";
const RESPONSE_CACHE_MAX_BODY_BYTES:usize = 1024 * 1024;
const RESPONSE_MAX_BODY_BYTES:usize = 64 * 1024 * 1024;
pub struct ResolvedThresholds {
pub cache:usize,
pub token:usize,
pub inline:usize,
pub code_multiplier:f64,
}
pub fn resolve_thresholds(compression:Option<&CompressionConfig>) -> ResolvedThresholds {
ResolvedThresholds {
cache:env_parse_warn::<usize>("APHRODITE_TOOL_THRESHOLD_CACHE")
.or_else(|| compression.and_then(|c| c.tool_threshold_cache).map(|v| v as usize))
.unwrap_or(CACHE_COMPRESS_THRESHOLD),
token:env_parse_warn::<usize>("APHRODITE_TOOL_THRESHOLD_TOKEN")
.or_else(|| compression.and_then(|c| c.tool_threshold_token).map(|v| v as usize))
.unwrap_or(TOKEN_COMPRESS_THRESHOLD),
inline:env_parse_warn::<usize>("APHRODITE_INLINE_THRESHOLD")
.or_else(|| compression.and_then(|c| c.inline_threshold).map(|v| v as usize))
.unwrap_or(INLINE_CCR_THRESHOLD),
code_multiplier:env_parse_warn::<f64>("APHRODITE_CODE_MULTIPLIER")
.or_else(|| compression.and_then(|c| c.code_multiplier))
.unwrap_or(3.0),
}
}
pub(crate) async fn ccr_get(ccr:&Arc<dyn CcrStore>, hash:&str) -> Option<String> {
let ccr = ccr.clone();
let hash = hash.to_owned();
tokio::task::spawn_blocking(move || ccr.get(&hash)).await.unwrap_or(None)
}
async fn ccr_put(ccr:&Arc<dyn CcrStore>, hash:&str, content:&str) -> bool {
let ccr = ccr.clone();
let hash = hash.to_owned();
let content = content.to_owned();
tokio::task::spawn_blocking(move || ccr.put(&hash, &content))
.await
.unwrap_or(false)
}
async fn ccr_del(ccr:&Arc<dyn CcrStore>, hash:&str) -> bool {
let ccr = ccr.clone();
let hash = hash.to_owned();
tokio::task::spawn_blocking(move || ccr.del(&hash)).await.unwrap_or(false)
}
async fn ccr_len(ccr:&Arc<dyn CcrStore>) -> usize {
let ccr = ccr.clone();
tokio::task::spawn_blocking(move || ccr.len()).await.unwrap_or(0)
}
pub struct AppState {
pub client:HttpClient,
pub api_url:String,
pub model:String,
pub api_key:Secret,
pub ccr:Option<Arc<dyn CcrStore>>,
pub add_markers:bool,
pub mode:ProxyMode,
pub tool_relay:bool,
pub notify_url:Option<String>,
pub notify_key:Option<String>,
pub dev:bool,
pub request_history:std::sync::Mutex<VecDeque<serde_json::Value>>,
pub inline_ccr:std::sync::Mutex<lru::LruCache<String, String>>,
pub latency_buckets:[AtomicU64; 5],
pub total_latency_micros:AtomicU64,
pub last_errors:std::sync::Mutex<VecDeque<String>>,
pub compressions_by_type:std::sync::Mutex<std::collections::HashMap<String, u64>>,
pub requests_total:AtomicU64,
pub requests_compressed:AtomicU64,
pub tokens_saved:AtomicU64,
pub ccr_hits:AtomicU64,
pub ccr_misses:AtomicU64,
pub ccr_created:AtomicU64,
pub tool_relay_calls:AtomicU64,
pub compression_ratio_ema:AtomicU64,
pub response_cache:std::sync::Mutex<lru::LruCache<u64, (std::time::Instant, Vec<u8>)>>,
pub response_cache_ttl:std::time::Duration,
pub cache_hits:AtomicU64,
pub cache_misses:AtomicU64,
pub task_tracker:TaskTracker,
pub fill_pct:AtomicU64,
pub inline_ccr_hits:AtomicU64,
pub inline_ccr_misses:AtomicU64,
pub tool_relay_success:AtomicU64,
pub tool_relay_failure:AtomicU64,
pub notify_success:AtomicU64,
pub notify_failure:AtomicU64,
pub upstream_errors_4xx:AtomicU64,
pub upstream_errors_5xx:AtomicU64,
pub upstream_timeouts:AtomicU64,
pub upstream_connect_errors:AtomicU64,
pub ccr_store_entries:AtomicU64,
pub ccr_store_bytes:AtomicU64,
pub request_body_bytes:AtomicU64,
pub response_body_bytes:AtomicU64,
pub upstream_latency_micros:AtomicU64,
pub upstream_health_cache:std::sync::Mutex<Option<(bool, std::time::Instant)>>,
pub cache_compress_threshold:AtomicUsize,
pub token_compress_threshold:AtomicUsize,
pub inline_ccr_threshold:AtomicUsize,
pub code_multiplier_x100:AtomicU64,
}
impl AppState {
pub fn stats_json(&self) -> serde_json::Value {
serde_json::json!({
"mode": match self.mode {
ProxyMode::Cache => "cache",
ProxyMode::Token => "token",
},
"proxy": "aphrodite",
"ccr_backend": if self.ccr.is_some() { "enabled" } else { "none" },
"tool_relay_enabled": self.tool_relay,
"requests": {
"total": self.requests_total.load(Ordering::Relaxed),
"compressed": self.requests_compressed.load(Ordering::Relaxed),
},
"tokens_saved": self.tokens_saved.load(Ordering::Relaxed),
"ccr": {
"hits": self.ccr_hits.load(Ordering::Relaxed),
"misses": self.ccr_misses.load(Ordering::Relaxed),
"created": self.ccr_created.load(Ordering::Relaxed),
},
"tool_relay_calls": self.tool_relay_calls.load(Ordering::Relaxed),
"cache": {
"hits": self.cache_hits.load(Ordering::Relaxed),
"misses": self.cache_misses.load(Ordering::Relaxed),
},
"latency_buckets_us": [
self.latency_buckets[0].load(Ordering::Relaxed),
self.latency_buckets[1].load(Ordering::Relaxed),
self.latency_buckets[2].load(Ordering::Relaxed),
self.latency_buckets[3].load(Ordering::Relaxed),
self.latency_buckets[4].load(Ordering::Relaxed),
],
"total_latency_micros": self.total_latency_micros.load(Ordering::Relaxed),
"compressions_by_type": self.compressions_by_type.lock().map(|m| m.clone()).unwrap_or_default(),
"compression_ratio_ema": self.compression_ratio_ema.load(Ordering::Relaxed) as f64 / 100.0,
"last_errors": self.last_errors.lock().map(|v| v.iter().rev().take(5).cloned().collect::<Vec<_>>()).unwrap_or_default(),
"request_history": self.request_history.lock().map(|v| v.clone()).unwrap_or_default(),
"inline_ccr": {
"hits": self.inline_ccr_hits.load(Ordering::Relaxed),
"misses": self.inline_ccr_misses.load(Ordering::Relaxed),
},
"tool_relay": {
"total": self.tool_relay_calls.load(Ordering::Relaxed),
"success": self.tool_relay_success.load(Ordering::Relaxed),
"failure": self.tool_relay_failure.load(Ordering::Relaxed),
},
"notify": {
"success": self.notify_success.load(Ordering::Relaxed),
"failure": self.notify_failure.load(Ordering::Relaxed),
},
"upstream_errors": {
"4xx": self.upstream_errors_4xx.load(Ordering::Relaxed),
"5xx": self.upstream_errors_5xx.load(Ordering::Relaxed),
"timeouts": self.upstream_timeouts.load(Ordering::Relaxed),
"connect_errors": self.upstream_connect_errors.load(Ordering::Relaxed),
},
"ccr_store": {
"entries": self.ccr_store_entries.load(Ordering::Relaxed),
"bytes_approx": self.ccr_store_bytes.load(Ordering::Relaxed),
},
"body_bytes": {
"request": self.request_body_bytes.load(Ordering::Relaxed),
"response": self.response_body_bytes.load(Ordering::Relaxed),
},
"upstream_latency_micros": self.upstream_latency_micros.load(Ordering::Relaxed),
})
}
fn compress_threshold(&self) -> usize {
match self.mode {
ProxyMode::Cache => self.cache_compress_threshold.load(Ordering::Relaxed),
ProxyMode::Token => self.token_compress_threshold.load(Ordering::Relaxed),
}
}
fn inline_ccr_threshold(&self) -> usize { self.inline_ccr_threshold.load(Ordering::Relaxed) }
fn code_multiplier(&self) -> f64 { self.code_multiplier_x100.load(Ordering::Relaxed) as f64 / 100.0 }
fn threshold_for(&self, ct:&str) -> usize {
let base = self.compress_threshold();
match ct {
"linter" | "build_output" | "log" => return base,
_ => {},
}
let ratio = self.compression_ratio_ema.load(Ordering::Relaxed) as f64 / 100.0;
let tune = if ratio > 20.0 {
2.0
} else if ratio < 3.0 && ratio > 0.0 {
0.5
} else {
1.0
};
let base = (base as f64 * tune) as usize;
match ct {
"error" => base * 8,
"code_rust" | "code_python" | "code_go" | "code_js" | "code" => (base as f64 * self.code_multiplier()) as usize,
"diff" | "git" => base * 2,
"text" => base * 2,
"tool_output" => base,
"json" => base,
_ => base,
}
}
fn update_compression_ratio(&self, original_len:usize, compressed_len:usize) {
if original_len == 0 || compressed_len == 0 {
return;
}
let ratio = (original_len as f64 / compressed_len as f64 * 100.0) as u64;
let old = self.compression_ratio_ema.load(Ordering::Relaxed);
let new = ((ratio as f64 * 0.2) + (old as f64 * 0.8)) as u64;
self.compression_ratio_ema.store(new, Ordering::Relaxed);
self.compute_fill_pct();
}
fn compute_fill_pct(&self) {
let ratio_ema = self.compression_ratio_ema.load(Ordering::Relaxed);
let pct = if ratio_ema == 0 {
99u64
} else {
let raw = 100u64.saturating_sub(ratio_ema / 20);
raw.clamp(1, 99)
};
self.fill_pct.store(pct * 100, Ordering::Relaxed); }
fn record_latency(&self, d:std::time::Duration) {
let us = d.as_micros() as u64;
let bucket = if us < 1_000 {
0
} else if us < 10_000 {
1
} else if us < 100_000 {
2
} else if us < 1_000_000 {
3
} else {
4
};
self.latency_buckets[bucket].fetch_add(1, Ordering::Relaxed);
self.total_latency_micros.fetch_add(us, Ordering::Relaxed);
}
fn record_error(&self, msg:String) {
if let Ok(mut v) = self.last_errors.lock() {
v.push_back(msg);
if v.len() > 100 {
v.pop_front();
}
}
}
fn record_compression(&self, ct:&str) {
if let Ok(mut m) = self.compressions_by_type.lock() {
*m.entry(ct.to_string()).or_insert(0) += 1;
}
}
fn record_request(&self, id:&str, method:&str, path:&str, status:u16, compressed:bool, elapsed_ms:u128) {
if let Ok(mut hist) = self.request_history.lock() {
hist.push_back(serde_json::json!({
"id": id,
"method": method,
"path": path,
"status": status,
"compressed": compressed,
"elapsed_ms": elapsed_ms,
}));
if hist.len() > 50 {
hist.pop_front();
}
}
}
}
#[derive(Debug, Deserialize)]
pub struct ToolRelayRequest {
pub tool:String,
pub params:serde_json::Value,
pub callback_url:Option<String>,
}
#[derive(Debug, Serialize)]
pub struct ToolRelayResponse {
pub success:bool,
pub result:Option<serde_json::Value>,
pub error:Option<String>,
pub async_call:bool,
}
#[derive(Debug, Deserialize)]
pub struct CcrCreateRequest {
pub content:String,
pub key:Option<String>,
pub ttl_seconds:Option<u64>,
pub tags:Option<Vec<String>>,
}
#[derive(Debug, Serialize)]
pub struct CcrCreateResponse {
pub hash:String,
pub token_savings_ratio:f64,
pub original_size:usize,
pub compressed_size:usize,
pub marker_size:usize,
}
#[derive(Debug, Serialize)]
pub struct CcrNotification {
pub event:String,
pub hash:String,
pub created_at:u64,
pub ttl:u64,
pub tags:Vec<String>,
}
pub async fn build_state(cli:&Cli, compression:Option<&CompressionConfig>) -> anyhow::Result<AppState> {
let client = HttpClient::builder()
.timeout(std::time::Duration::from_secs(cli.timeout))
.connect_timeout(std::time::Duration::from_secs(10))
.pool_max_idle_per_host(100)
.pool_idle_timeout(std::time::Duration::from_secs(90))
.tcp_keepalive(std::time::Duration::from_secs(60))
.build()?;
let ccr:Option<Arc<dyn CcrStore>> = match cli.mode {
ProxyMode::Token if !cli.no_ccr_marker => {
let db_path = cli.ccr_db_path.as_ref().map_or_else(
|| {
dirs::home_dir()
.unwrap_or_else(|| std::path::PathBuf::from("/tmp"))
.join(".hermes")
.join("aphrodite")
.join("ccr.db")
},
|p| p.clone(),
);
if let Some(parent) = db_path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| anyhow::anyhow!("SQLite CCR: cannot create directory {}: {}", parent.display(), e))?;
}
let store = SqliteCcrStore::open(&db_path, cli.ccr_ttl_seconds)
.map_err(|e| anyhow::anyhow!("SQLite CCR: {}", e))?;
Some(Arc::new(store))
},
ProxyMode::Cache => {
let store =
InMemoryCcrStore::with_capacity_and_ttl(10_000, std::time::Duration::from_secs(cli.ccr_ttl_seconds));
Some(Arc::new(store))
},
_ => None,
};
let thresholds = resolve_thresholds(compression);
Ok(AppState {
client,
api_url:cli.api_url.clone(),
model:cli.model.clone(),
api_key:cli.api_key.clone().into(),
ccr,
add_markers:!cli.no_ccr_marker,
mode:cli.mode,
tool_relay:cli.tool_relay,
notify_url:cli.notify_url.clone(),
notify_key:cli.notify_key.clone(),
dev:cli.dev,
latency_buckets:[
AtomicU64::new(0),
AtomicU64::new(0),
AtomicU64::new(0),
AtomicU64::new(0),
AtomicU64::new(0),
],
total_latency_micros:AtomicU64::new(0),
last_errors:Mutex::new(VecDeque::new()),
compressions_by_type:Mutex::new(HashMap::new()),
request_history:Mutex::new(VecDeque::new()),
inline_ccr:Mutex::new(lru::LruCache::new(NonZeroUsize::new(1024).unwrap())),
requests_total:AtomicU64::new(0),
requests_compressed:AtomicU64::new(0),
tokens_saved:AtomicU64::new(0),
ccr_hits:AtomicU64::new(0),
ccr_misses:AtomicU64::new(0),
ccr_created:AtomicU64::new(0),
tool_relay_calls:AtomicU64::new(0),
compression_ratio_ema:AtomicU64::new(200), response_cache:Mutex::new(lru::LruCache::new(NonZeroUsize::new(128).unwrap())),
response_cache_ttl:std::time::Duration::from_secs(cli.ccr_ttl_seconds),
cache_hits:AtomicU64::new(0),
cache_misses:AtomicU64::new(0),
fill_pct:AtomicU64::new(9000), task_tracker:TaskTracker::new(),
inline_ccr_hits:AtomicU64::new(0),
inline_ccr_misses:AtomicU64::new(0),
tool_relay_success:AtomicU64::new(0),
tool_relay_failure:AtomicU64::new(0),
notify_success:AtomicU64::new(0),
notify_failure:AtomicU64::new(0),
upstream_errors_4xx:AtomicU64::new(0),
upstream_errors_5xx:AtomicU64::new(0),
upstream_timeouts:AtomicU64::new(0),
upstream_connect_errors:AtomicU64::new(0),
ccr_store_entries:AtomicU64::new(0),
ccr_store_bytes:AtomicU64::new(0),
request_body_bytes:AtomicU64::new(0),
response_body_bytes:AtomicU64::new(0),
upstream_latency_micros:AtomicU64::new(0),
upstream_health_cache:std::sync::Mutex::new(None),
cache_compress_threshold:AtomicUsize::new(thresholds.cache),
token_compress_threshold:AtomicUsize::new(thresholds.token),
inline_ccr_threshold:AtomicUsize::new(thresholds.inline),
code_multiplier_x100:AtomicU64::new((thresholds.code_multiplier * 100.0) as u64),
})
}
fn cache_key_from_body(body:&[u8], api_key:&str) -> Option<u64> {
let v:serde_json::Value = serde_json::from_slice(body).ok()?;
if v.get("stream").and_then(|s| s.as_bool()).unwrap_or(false) {
return None;
}
v.get("model")?.as_str()?;
v.get("messages")?;
let mut parts:Vec<u8> = Vec::new();
parts.extend_from_slice(api_key.as_bytes());
for (label, val) in [
("model", v.get("model")),
("messages", v.get("messages")),
("tools", v.get("tools")),
("tool_choice", v.get("tool_choice")),
("temperature", v.get("temperature")),
("top_p", v.get("top_p")),
("n", v.get("n")),
("response_format", v.get("response_format")),
] {
parts.push(b':');
parts.extend_from_slice(label.as_bytes());
parts.push(b'=');
if let Some(val) = val {
parts.extend_from_slice(serde_json::to_string(val).ok()?.as_bytes());
}
}
Some(fnv1a_64(&parts))
}
fn response_cache_get(state:&AppState, ck:u64) -> Option<Vec<u8>> {
state.response_cache.lock().ok().and_then(|mut cache| {
let expired = cache
.peek(&ck)
.map(|(inserted_at, _)| inserted_at.elapsed() >= state.response_cache_ttl)
.unwrap_or(false);
if expired {
cache.pop(&ck);
None
} else {
cache.get(&ck).map(|(_, body)| body.clone())
}
})
}
fn copy_upstream_headers(
mut builder:axum::http::response::Builder,
upstream_headers:&reqwest::header::HeaderMap,
) -> axum::http::response::Builder {
const SKIP:&[&str] = &[
"content-length",
"content-type",
"transfer-encoding",
"connection",
"keep-alive",
];
for (name, value) in upstream_headers.iter() {
if SKIP.contains(&name.as_str()) {
continue;
}
if let Ok(v) = axum::http::HeaderValue::from_bytes(value.as_bytes()) {
builder = builder.header(name.as_str(), v);
}
}
builder
}
async fn accumulate_body(
response: reqwest::Response,
max_bytes: usize,
) -> Result<bytes::Bytes, String> {
let mut buf = Vec::new();
let mut stream = response.bytes_stream();
while let Some(chunk) = stream.next().await {
match chunk {
Ok(b) => {
if buf.len() + b.len() > max_bytes {
return Err(format!(
"response body exceeded {} MB limit",
max_bytes / (1024 * 1024)
));
}
buf.extend_from_slice(&b);
},
Err(e) => return Err(format!("body read: {}", e)),
}
}
Ok(bytes::Bytes::from(buf))
}
fn fnv1a_64(bytes:&[u8]) -> u64 {
const FNV_OFFSET:u64 = 14695981039346656037;
const FNV_PRIME:u64 = 1099511628211;
let mut hash = FNV_OFFSET;
for &b in bytes {
hash ^= b as u64;
hash = hash.wrapping_mul(FNV_PRIME);
}
hash
}
pub async fn proxy_handler(
State(state):State<Arc<AppState>>,
method:Method,
path:axum::extract::OriginalUri,
headers:axum::http::HeaderMap,
body:Bytes,
) -> impl IntoResponse {
state.requests_total.fetch_add(1, Ordering::Relaxed);
state.request_body_bytes.fetch_add(body.len() as u64, Ordering::Relaxed);
let t0 = std::time::Instant::now();
let req_id = uuid::Uuid::new_v4().to_string();
let req_id_short = &req_id[..8];
if state.dev {
let mut hdr_log = String::new();
for (k, v) in headers.iter() {
let val = v.to_str().unwrap_or("?");
if k.as_str().to_lowercase() != "authorization" {
hdr_log.push_str(&format!(" {}: {}", k.as_str(), if val.len() > 80 { &val[..80] } else { val }));
} else {
hdr_log.push_str(" authorization: [REDACTED]");
}
hdr_log.push('\n');
}
tracing::info!(
id = %req_id_short,
method = %method,
path = %path.path(),
body_len = body.len(),
headers = %hdr_log,
">>> REQ"
);
}
let deepseek_path_and_query = path
.0
.path_and_query()
.map(|pq| pq.as_str())
.unwrap_or_else(|| path.path())
.trim_start_matches('/');
let url = format!("{}/{}", state.api_url.trim_end_matches('/'), deepseek_path_and_query);
let is_chat_completion = path.path().trim_start_matches('/') == CHAT_COMPLETIONS_PATH.trim_start_matches('/');
let body_vec = body.to_vec();
let cache_key = if is_chat_completion {
cache_key_from_body(&body_vec, state.api_key.expose())
} else {
None
};
if let Some(ck) = cache_key {
let cached_body = response_cache_get(&state, ck);
if let Some(cached_body) = cached_body {
state.cache_hits.fetch_add(1, Ordering::Relaxed);
state.tokens_saved.fetch_add(cached_body.len() as u64, Ordering::Relaxed);
if state.dev {
tracing::info!(
id = %req_id_short,
cached_len = cached_body.len(),
"<<< CACHE HIT"
);
}
state.record_latency(t0.elapsed());
state.record_request(req_id_short, method.as_str(), path.path(), 200, false, t0.elapsed().as_millis());
return Response::builder()
.status(StatusCode::OK)
.header("Content-Type", "application/json; charset=utf-8")
.header("X-Aphrodite-Cache", "HIT")
.header("X-Aphrodite-Fill-Pct", {
let v = state.fill_pct.load(Ordering::Relaxed) as f64 / 100.0;
if v.is_finite() { format!("{:.1}", v) } else { "0.0".to_string() }
})
.body(Body::from(cached_body))
.unwrap();
} else {
state.cache_misses.fetch_add(1, Ordering::Relaxed);
if state.dev {
tracing::info!(
id = %req_id_short,
"<<< CACHE MISS"
);
}
}
}
let mut upstream_result = Err("unreachable".to_string());
let mut final_error_was_timeout = false;
for attempt in 1..=3u32 {
let req = state
.client
.request(method.clone(), &url)
.header("Content-Type", "application/json; charset=utf-8")
.header("Accept", "application/json")
.header("Authorization", format!("Bearer {}", state.api_key.expose()));
let mut req = req;
for (key, val) in headers.iter() {
let k = key.as_str().to_lowercase();
if k != "host"
&& k != "authorization"
&& k != "content-length"
&& k != "content-type"
&& k != "accept"
&& k != "accept-encoding"
&& !k.starts_with("x-aphrodite-")
{
req = req.header(key, val);
}
}
match req.body(body_vec.clone()).send().await {
Ok(r) => {
upstream_result = Ok(r);
break;
},
Err(e) => {
if attempt < 3 && e.is_connect() {
let base_ms = 100 * 2u64.pow(attempt - 1);
let jitter = rand::random::<f64>() * 0.5 + 0.75; let ms = (base_ms as f64 * jitter) as u64;
tracing::warn!(attempt, backoff_ms = ms, "upstream retry after connect error: {}", e);
tokio::time::sleep(std::time::Duration::from_millis(ms)).await;
} else {
final_error_was_timeout = e.is_timeout();
upstream_result = Err(format!("{}", e));
break;
}
},
}
}
match upstream_result {
Ok(response) => {
let status = response.status();
let status_code = status.as_u16();
if status_code >= 500 {
state.upstream_errors_5xx.fetch_add(1, Ordering::Relaxed);
} else if status_code >= 400 {
state.upstream_errors_4xx.fetch_add(1, Ordering::Relaxed);
}
let upstream_headers = response.headers().clone();
let content_type = upstream_headers.get("content-type").cloned();
let is_sse = content_type.as_ref()
.map(|ct| ct.as_bytes().starts_with(b"text/event-stream"))
.unwrap_or(false);
if is_sse {
let stream = response.bytes_stream();
if state.dev {
tracing::info!(id = %req_id_short, status = %status, "<<< STREAM (SSE)");
}
state.record_latency(t0.elapsed());
state.record_request(req_id_short, method.as_str(), path.path(), status.as_u16(), false, t0.elapsed().as_millis());
let mut builder = Response::builder().status(status);
builder = copy_upstream_headers(builder, &upstream_headers);
if let Some(ct) = content_type {
builder = builder.header("Content-Type", ct);
}
builder = builder.header("X-Aphrodite-Streamed", "true");
return builder.body(Body::from_stream(stream)).unwrap();
}
let resp_body = match accumulate_body(response, RESPONSE_MAX_BODY_BYTES).await {
Ok(b) => b,
Err(e) => {
state.record_error(format!("body read: {}", e));
return (
StatusCode::BAD_GATEWAY,
Json(serde_json::json!({"error": format!("body read: {}", e)})),
)
.into_response();
},
};
let upstream_elapsed = t0.elapsed().as_micros() as u64;
state.upstream_latency_micros.fetch_add(upstream_elapsed, Ordering::Relaxed);
state.response_body_bytes.fetch_add(resp_body.len() as u64, Ordering::Relaxed);
let elapsed = t0.elapsed();
if is_chat_completion && state.ccr.is_some() {
let headroom_budget = headers.get("x-headroom-budget").and_then(|v| v.to_str().ok());
if state.dev && headroom_budget.is_some() {
tracing::info!(
id = %req_id_short,
budget = %headroom_budget.unwrap_or(""),
"headroom budget applied to compression threshold"
);
}
if let Some(compressed) = compress_chat_completion(&state, &resp_body, headroom_budget).await {
state.requests_compressed.fetch_add(1, Ordering::Relaxed);
state.record_latency(elapsed);
state.record_request(
req_id_short,
method.as_str(),
path.path(),
status.as_u16(),
true,
elapsed.as_millis(),
);
if state.dev {
let elapsed = t0.elapsed();
let comp_len = serde_json::to_vec(&compressed).map(|v| v.len()).unwrap_or(0);
tracing::info!(
id = %req_id_short,
status = %status,
original_len = resp_body.len(),
compressed_len = comp_len,
ratio = format!("{:.1}x", resp_body.len() as f64 / comp_len.max(1) as f64),
elapsed_ms = elapsed.as_millis(),
"<<< COMPRESSED"
);
}
let body = serde_json::to_vec(&compressed).unwrap_or_else(|_| resp_body.to_vec());
if let Some(ck) = cache_key {
if status.is_success() && body.len() <= RESPONSE_CACHE_MAX_BODY_BYTES {
if let Ok(mut cache) = state.response_cache.lock() {
cache.put(ck, (std::time::Instant::now(), body.clone()));
}
}
}
let mut builder = Response::builder().status(status);
builder = copy_upstream_headers(builder, &upstream_headers);
return builder
.header("Content-Type", "application/json; charset=utf-8")
.header("X-Aphrodite-Compressed", "true")
.header("X-Aphrodite-Cache", "MISS")
.header("X-Aphrodite-Fill-Pct", {
let v = state.fill_pct.load(Ordering::Relaxed) as f64 / 100.0;
if v.is_finite() { format!("{:.1}", v) } else { "0.0".to_string() }
})
.body(Body::from(body))
.unwrap();
}
}
if state.dev {
let elapsed = t0.elapsed();
let body_preview = if resp_body.len() > 500 {
let s = std::str::from_utf8(&resp_body).unwrap_or("?");
let preview:String = s.char_indices().take_while(|(i, _)| *i < 200).map(|(_, c)| c).collect();
format!("{}... ({} total)", preview, resp_body.len())
} else {
std::str::from_utf8(&resp_body).unwrap_or("?").to_string()
};
tracing::info!(
id = %req_id_short,
status = %status,
resp_len = resp_body.len(),
elapsed_ms = elapsed.as_millis(),
body = %body_preview,
"<<< RES"
);
}
state.record_latency(t0.elapsed());
state.record_request(
req_id_short,
method.as_str(),
path.path(),
status.as_u16(),
false,
t0.elapsed().as_millis(),
);
if let Some(ck) = cache_key {
if status.is_success() && resp_body.len() <= RESPONSE_CACHE_MAX_BODY_BYTES {
if let Ok(mut cache) = state.response_cache.lock() {
cache.put(ck, (std::time::Instant::now(), resp_body.to_vec()));
}
}
}
let mut builder = Response::builder().status(status);
builder = copy_upstream_headers(builder, &upstream_headers);
builder = builder.header("X-Aphrodite-Cache", "MISS");
builder = builder.header("X-Aphrodite-Fill-Pct", {
let v = state.fill_pct.load(Ordering::Relaxed) as f64 / 100.0;
if v.is_finite() { format!("{:.1}", v) } else { "0.0".to_string() }
});
if let Some(ct) = content_type {
builder = builder.header("Content-Type", ct);
}
builder.body(Body::from(resp_body)).unwrap()
},
Err(e) => {
if final_error_was_timeout {
state.upstream_timeouts.fetch_add(1, Ordering::Relaxed);
} else {
state.upstream_connect_errors.fetch_add(1, Ordering::Relaxed);
}
state.record_latency(t0.elapsed());
state.record_request(req_id_short, method.as_str(), path.path(), 502, false, t0.elapsed().as_millis());
state.record_error(format!("upstream: {}", e));
if state.dev {
tracing::error!(
id = %req_id_short,
error = %e,
elapsed_ms = t0.elapsed().as_millis(),
"<<< ERR"
);
}
(
StatusCode::BAD_GATEWAY,
Json(serde_json::json!({"error": format!("upstream: {}", e)})),
)
.into_response()
},
}
}
fn proxy_detect_content_type(content:&str) -> &'static str {
let first_line = content.lines().next().unwrap_or("");
if content.starts_with('{') || content.starts_with('[') {
if serde_json::from_str::<serde_json::Value>(content).is_err() {
return "text";
}
if content.contains("exit_code") || content.contains("\"status\"") {
return "tool_output";
}
return "json";
}
if content.lines().count() > 3 {
if content.lines().any(|l| {
let t = l.trim_start();
t.starts_with("fn ")
|| t.starts_with("pub fn ")
|| t.starts_with("async fn ")
|| t.starts_with("pub async fn ")
|| t.starts_with("impl ")
|| t.starts_with("struct ")
|| t.starts_with("pub struct ")
|| t.starts_with("enum ")
|| t.starts_with("pub enum ")
}) && (content.contains("-> ") || content.contains("&") || content.contains("use "))
{
return "code_rust";
}
if content.contains("def ")
&& (content.contains("import ")
|| content.contains("class ")
|| content.contains("from ")
|| content.contains("self."))
{
return "code_python";
}
if (content.contains("func ") || content.contains("package ")) && content.contains("import (") {
return "code_go";
}
if (content.contains("function ") || content.contains("const ") || content.contains("=> "))
&& (content.contains("import ") || content.contains("export "))
{
return "code_js";
}
if content.contains("fn ")
|| content.contains("def ")
|| content.contains("class ")
|| content.contains("import ")
|| content.contains("pub fn")
{
return "code";
}
}
if first_line.contains("error")
|| first_line.contains("Error")
|| first_line.contains("ERROR")
|| first_line.contains("Traceback")
|| first_line.contains("panic")
|| first_line.starts_with("thread '")
{
return "error";
}
if first_line.starts_with("Compiling ")
|| first_line.starts_with(" Compiling ")
|| first_line.contains("Finished")
|| first_line.starts_with("running ")
|| first_line.starts_with("test ")
{
return "build_output";
}
if first_line.starts_with("error[E")
|| first_line.starts_with("error: ")
|| first_line.starts_with("warning[")
|| first_line.starts_with("warning: ")
|| first_line.contains("|") && (first_line.contains("error") || first_line.contains("warning"))
|| first_line.contains("mypy")
|| first_line.contains("clippy")
|| first_line.contains("eslint")
|| first_line.contains("tsc ")
{
return "linter";
}
if first_line.starts_with("diff --git ")
|| first_line.starts_with("@@ -")
|| first_line.starts_with("+++ ")
|| first_line.starts_with("--- ")
{
return "diff";
}
if first_line.starts_with("commit ") || first_line.starts_with("On branch ") {
return "git";
}
if content.lines().any(|l| {
let t = l.trim();
t.starts_with('[')
&& (t.contains("INFO")
|| t.contains("WARN")
|| t.contains("ERROR")
|| t.contains("DEBUG")
|| t.contains("TRACE")
|| t.contains("FATAL")
|| t.contains("PANIC"))
}) || content.lines().any(|l| {
let t = l.trim();
t.starts_with(|c:char| c.is_ascii_digit()) && t.len() > 10 && (t.contains(':') || t.contains('-'))
}) {
return "log";
}
"text"
}
fn generate_metadata(content:&str, ct:&str) -> String {
let line_count = content.lines().count();
let mut parts:Vec<String> = Vec::new();
match ct {
"code_rust" => {
parts.push("lang=rs".to_string());
let fns:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("fn ") || t.starts_with("pub fn ") || t.starts_with("async fn ")
})
.filter_map(|l| {
let t = l.trim_start();
let after_fn = t
.strip_prefix("pub async fn ")
.or_else(|| t.strip_prefix("pub fn "))
.or_else(|| t.strip_prefix("async fn "))
.or_else(|| t.strip_prefix("fn "))?;
after_fn.split(['(', ' ', '<']).next().filter(|s| !s.is_empty())
})
.collect();
if !fns.is_empty() {
parts.push(format!("fns={}", fns.join(",")));
}
let structs:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("struct ") || t.starts_with("pub struct ")
})
.filter_map(|l| {
let t = l.trim_start();
let after = t
.strip_prefix("pub struct ")
.unwrap_or_else(|| t.strip_prefix("struct ").unwrap_or(t));
after.split(['(', ' ', '<', '{']).next().filter(|s| !s.is_empty())
})
.collect();
if !structs.is_empty() {
parts.push(format!("structs={}", structs.join(",")));
}
let impls:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("impl ") || t.starts_with("pub impl ")
})
.filter_map(|l| {
let t = l.trim_start();
let after = t
.strip_prefix("pub impl ")
.unwrap_or_else(|| t.strip_prefix("impl ").unwrap_or(t));
after.split_whitespace().next().map(|w| w.trim_end_matches('<'))
})
.collect();
if !impls.is_empty() {
parts.push(format!("impls={}", impls.join(",")));
}
let traits:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("trait ") || t.starts_with("pub trait ")
})
.filter_map(|l| {
let t = l.trim_start();
let after = t
.strip_prefix("pub trait ")
.unwrap_or_else(|| t.strip_prefix("trait ").unwrap_or(t));
after.split([' ', '<', '{']).next().filter(|s| !s.is_empty())
})
.collect();
if !traits.is_empty() {
parts.push(format!("traits={}", traits.join(",")));
}
parts.push(format!("ln={}", line_count));
},
"code_python" => {
parts.push("lang=py".to_string());
let fns:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("def ") || t.starts_with("async def ")
})
.filter_map(|l| {
let t = l.trim_start();
let after = t
.strip_prefix("async def ")
.unwrap_or_else(|| t.strip_prefix("def ").unwrap_or(t));
after.split(['(', ' ', ':']).next().filter(|s| !s.is_empty())
})
.collect();
if !fns.is_empty() {
parts.push(format!("fns={}", fns.join(",")));
}
let classes:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("class ")
})
.filter_map(|l| {
let t = l.trim_start();
let after = t.strip_prefix("class ")?;
after.split(['(', ' ', ':']).next().filter(|s| !s.is_empty())
})
.collect();
if !classes.is_empty() {
parts.push(format!("classes={}", classes.join(",")));
}
let imports:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("import ") || t.starts_with("from ")
})
.filter_map(|l| {
let t = l.trim_start();
if let Some(rest) = t.strip_prefix("import ") {
rest.split([' ', ',', ';']).next().filter(|s| !s.is_empty())
} else {
t.strip_prefix("from ")?.split(' ').next().filter(|s| !s.is_empty())
}
})
.collect();
if !imports.is_empty() {
parts.push(format!("imports={}", imports.join(",")));
}
let decorators:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with('@')
})
.filter_map(|l| {
let t = l.trim_start();
let name = &t[1..];
name.split(['(', ' ']).next().filter(|s| !s.is_empty())
})
.collect();
if !decorators.is_empty() {
parts.push(format!("decorators={}", decorators.join(",")));
}
parts.push(format!("ln={}", line_count));
},
"code_go" => {
parts.push("lang=go".to_string());
let fns:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("func ")
})
.filter_map(|l| {
let t = l.trim_start();
let after = t.strip_prefix("func ")?;
after.split(['(', ' ']).next().filter(|s| !s.is_empty())
})
.collect();
if !fns.is_empty() {
parts.push(format!("fns={}", fns.join(",")));
}
parts.push(format!("ln={}", line_count));
},
"code_js" => {
parts.push("lang=js".to_string());
let fns:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("function ") || t.starts_with("const ")
})
.filter_map(|l| {
let t = l.trim_start();
if let Some(rest) = t.strip_prefix("function ") {
rest.split(['(', ' ']).next().filter(|s| !s.is_empty())
} else {
t.strip_prefix("const ")?
.split([' ', '=', ':'])
.next()
.filter(|s| !s.is_empty())
}
})
.collect();
if !fns.is_empty() {
parts.push(format!("fns={}", fns.join(",")));
}
parts.push(format!("ln={}", line_count));
},
"code" => {
parts.push("lang=gen".to_string());
let sigs:Vec<&str> = content
.lines()
.filter(|l| {
let t = l.trim_start();
t.starts_with("fn ")
|| t.starts_with("def ")
|| t.starts_with("func ")
|| t.starts_with("function ")
|| t.starts_with("class ")
|| t.starts_with("struct ")
})
.filter_map(|l| {
let t = l.trim_start();
let after = t
.strip_prefix("fn ")
.or_else(|| t.strip_prefix("def "))
.or_else(|| t.strip_prefix("func "))
.or_else(|| t.strip_prefix("function "))
.or_else(|| t.strip_prefix("class "))
.or_else(|| t.strip_prefix("struct "))?;
after.split(['(', ' ']).next()
})
.collect();
if !sigs.is_empty() {
parts.push(format!("sigs={}", sigs.join(",")));
}
parts.push(format!("ln={}", line_count));
},
"error" => {
let mut trace = String::new();
for l in content.lines() {
let t = l.trim();
let ext_pos = t.find(".rs:").or_else(|| t.find(".py:")).or_else(|| t.find(".go:"));
if let Some(pos) = ext_pos {
let mut start = pos.saturating_sub(12);
while start > 0 && !t.is_char_boundary(start) {
start -= 1;
}
let mut end = (pos + 40).min(t.len());
while end < t.len() && !t.is_char_boundary(end) {
end += 1;
}
trace = t[start..end].to_string();
break;
}
}
if !trace.is_empty() {
parts.push(format!("trace={}", trace.replace('|', "/")));
}
let msg = content.lines().find(|l| l.contains("Error:") || l.contains("error[")).map(|l| {
let t = l.trim();
let idx = t.find("Error:").or_else(|| t.find("error[")).unwrap_or(0);
t[idx..].chars().take(80).collect::<String>().replace('|', "/")
});
if let Some(m) = msg {
parts.push(format!("msg={}", m));
} else {
let fl = content.lines().next().unwrap_or("").trim();
if !fl.is_empty() {
parts.push(format!("msg={}", fl.chars().take(80).collect::<String>().replace('|', "/")));
}
}
let err_count = content
.lines()
.filter(|l| l.contains("error") || l.starts_with("thread '"))
.count();
if err_count > 0 {
parts.push(format!("N_errors={}", err_count));
}
},
"diff" => {
let files = content.lines().filter(|l| l.starts_with("diff --git ")).count();
if files > 0 {
parts.push(format!("files={}", files));
}
let adds = content.lines().filter(|l| l.starts_with('+') && !l.starts_with("+++")).count();
let dels = content.lines().filter(|l| l.starts_with('-') && !l.starts_with("---")).count();
if adds > 0 {
parts.push(format!("adds={}", adds));
}
if dels > 0 {
parts.push(format!("dels={}", dels));
}
},
"git" => {
for l in content.lines() {
let t = l.trim();
if let Some(rest) = t.strip_prefix("On branch ") {
parts.push(format!("branch={}", rest.trim().replace('|', "/")));
break;
}
}
if !parts.iter().any(|p| p.starts_with("branch=")) {
for l in content.lines() {
let t = l.trim();
if !t.is_empty() && !t.starts_with("* ") && !t.starts_with(" ") {
parts.push(format!("branch={}", t.chars().take(40).collect::<String>().replace('|', "/")));
break;
}
}
}
let commits = content
.lines()
.filter(|l| l.starts_with("commit ") || l.trim().starts_with("* ") || l.contains("commit"))
.count();
if commits > 0 {
parts.push(format!("commits={}", commits));
}
},
"build_output" => {
if content.contains("error") || content.contains("aborting") {
parts.push("status=FAIL".to_string());
} else {
parts.push("status=OK".to_string());
}
let files = content
.lines()
.filter(|l| l.starts_with("Compiling ") || l.contains(" Compiling "))
.count();
if files > 0 {
parts.push(format!("files={}", files));
}
let first_err = content
.lines()
.find(|l| l.contains("error[") || l.contains("Error:"))
.map(|l| l.trim().chars().take(80).collect::<String>().replace('|', "/"));
if let Some(e) = first_err {
parts.push(format!("first_err={}", e));
}
},
"log" => {
for l in content.lines() {
let t = l.trim();
for level in &["ERROR", "WARN", "WARNING", "INFO", "DEBUG", "TRACE", "FATAL", "PANIC"] {
if t.contains(level) {
parts.push(format!("level={}", level.to_lowercase()));
break;
}
}
if parts.iter().any(|p| p.starts_with("level=")) {
break;
}
}
let last_line = content.lines().last().unwrap_or("").trim().chars().take(60).collect::<String>();
if !last_line.is_empty() {
parts.push(format!("last={}", last_line.replace('|', "/")));
}
parts.push(format!("ln={}", line_count));
},
"linter" => {
let files_linted = content
.lines()
.filter(|l| {
(l.contains(".rs:") || l.contains(".py:") || l.contains(".go:") || l.contains(".ts:"))
&& (l.contains("error") || l.contains("warning"))
})
.count();
if files_linted > 0 {
parts.push(format!("files={}", files_linted));
}
let first_err = content
.lines()
.find(|l| l.contains("error[") || l.contains("Error:") || l.starts_with("error: "))
.map(|l| l.trim().chars().take(80).collect::<String>().replace('|', "/"));
if let Some(e) = first_err {
parts.push(format!("first_err={}", e));
}
parts.push(format!("ln={}", line_count));
},
"json" | "tool_output" => {
let mut keys:Vec<String> = Vec::new();
for l in content.lines() {
let bytes = l.as_bytes();
let mut i = 0;
while i + 3 < bytes.len() {
if bytes[i] == b'"' {
let start = i + 1;
let mut end = start;
while end < bytes.len() && bytes[end] != b'"' {
end += 1;
}
if end < bytes.len() && end + 2 < bytes.len() && bytes[end + 1] == b':' {
let key = &l[start..end];
if !key.starts_with('_') && !keys.contains(&key.to_string()) {
keys.push(key.to_string());
if keys.len() >= 10 {
break;
}
}
}
i = end + 1;
} else {
i += 1;
}
}
if keys.len() >= 10 {
break;
}
}
if !keys.is_empty() {
parts.push(format!("keys={}", keys.join(",")));
}
let entries = content
.lines()
.filter(|l| {
let t = l.trim();
t.starts_with('{') || t.starts_with('"') || t.starts_with('[')
})
.count();
if entries > 1 {
parts.push(format!("entries={}", entries));
}
},
"text" => {
parts.push(format!("ln={}", line_count));
},
_ => {
parts.push(format!("ln={}", line_count));
},
}
let result = parts.join(";").replace('\n', " ").replace('\r', "");
let truncated:String = result.chars().take(400).collect();
truncated.trim_end_matches([';', ' ', ',']).to_string()
}
fn proxy_format_ccr_output(preview:&str, ct:&str, metadata:&str, center:Option<&str>, hash:&str, size:usize) -> String {
let center_seg = center.map(|c| format!(";center={c}")).unwrap_or_default();
format!("{preview}\n[{ct}: {metadata}{center_seg}]\n<<<CCR:{hash}|{ct}|{size}>>>")
}
fn proxy_build_preview(content:&str, ct:&str) -> String {
match ct {
"code_rust" | "code_python" | "code_go" | "code_js" | "code_ts" | "code_sh" | "code" => {
let mut fns:Vec<&str> = Vec::new();
let mut structs:Vec<&str> = Vec::new();
let mut impls:Vec<&str> = Vec::new();
let mut classes:Vec<&str> = Vec::new();
let mut budget:usize = 280;
for line in content.lines() {
if budget == 0 {
break;
}
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
if ct == "code_rust" || ct == "code" {
if trimmed.strip_prefix("fn ").is_some() {
let sig:String = trimmed.chars().take(58).collect();
fns.push(trimmed); budget = budget.saturating_sub(sig.len() + 2);
} else if trimmed.strip_prefix("pub fn ").is_some() {
let sig:String = trimmed.chars().take(58).collect();
fns.push(trimmed);
budget = budget.saturating_sub(sig.len() + 2);
} else if trimmed.starts_with("struct ") || trimmed.starts_with("pub struct ") {
let s:String = trimmed.chars().take(50).collect();
structs.push(trimmed);
budget = budget.saturating_sub(s.len() + 2);
} else if trimmed.starts_with("impl ") {
let s:String = trimmed.chars().take(50).collect();
impls.push(trimmed);
budget = budget.saturating_sub(s.len() + 2);
}
}
if ct == "code_python" || ct == "code" {
if (trimmed.starts_with("def ") || trimmed.starts_with("async def ")) && trimmed.ends_with(':') {
let s:String = trimmed.chars().take(58).collect();
fns.push(trimmed);
budget = budget.saturating_sub(s.len() + 2);
} else if trimmed.starts_with("class ") && trimmed.ends_with(':') {
let s:String = trimmed.chars().take(50).collect();
classes.push(trimmed);
budget = budget.saturating_sub(s.len() + 2);
}
}
if ct == "code_go" && trimmed.starts_with("func ") {
let s:String = trimmed.chars().take(58).collect();
fns.push(trimmed);
budget = budget.saturating_sub(s.len() + 2);
}
}
let mut parts:Vec<String> = Vec::new();
if !fns.is_empty() {
parts.push(format!("{}fns", fns.len()));
}
if !structs.is_empty() {
parts.push(format!("{}structs", structs.len()));
}
if !impls.is_empty() {
parts.push(format!("{}impls", impls.len()));
}
if !classes.is_empty() {
parts.push(format!("{}classes", classes.len()));
}
let summary = if parts.is_empty() { "?".to_string() } else { parts.join("|") };
let sig_previews:Vec<String> = fns.iter().take(2).map(|s| s.chars().take(56).collect::<String>()).collect();
let sig_str = sig_previews.join("; ");
let lines = content.lines().count();
format!("[{ct}:{summary} {sig_str} {lines}L]").chars().take(300).collect()
},
"error" => {
let err_line = content
.lines()
.find(|l| l.contains("Error:") || l.contains("error[") || l.contains("panicked"))
.unwrap_or_else(|| content.lines().next().unwrap_or(""));
err_line.chars().take(300).collect()
},
"diff" => {
let files:Vec<&str> = content.lines().filter(|l| l.starts_with("diff --git ")).take(2).collect();
if files.is_empty() {
content.lines().next().unwrap_or("").chars().take(200).collect()
} else {
files.join("\n").chars().take(300).collect()
}
},
"json" | "tool_output" => {
let first = content.lines().next().unwrap_or("");
let key_count = content.matches("\":").count();
format!("{} … {} keys", first.chars().take(150).collect::<String>(), key_count)
},
"build_output" => {
content
.lines()
.find(|l| l.contains("Compiling") || l.contains("Finished") || l.contains("error"))
.unwrap_or_else(|| content.lines().next().unwrap_or(""))
.chars()
.take(250)
.collect()
},
_ => {
content.lines().next().unwrap_or("").chars().take(250).collect()
},
}
}
fn smart_marker(hash:&str, content:&str, ct:&str, center:Option<&str>) -> String {
let size = content.len();
let metadata = generate_metadata(content, ct);
let preview = proxy_build_preview(content, ct);
proxy_format_ccr_output(&preview, ct, &metadata, center, hash, size)
}
fn cache_marker(hash:&str, content:&str, ct:&str, center:Option<&str>) -> String {
let size = content.len();
let preview:String = content.chars().take(512).collect();
proxy_format_ccr_output(&preview, ct, "", center, hash, size)
}
async fn compress_chat_completion(
state:&AppState,
resp_body:&[u8],
headroom_budget:Option<&str>,
) -> Option<serde_json::Value> {
let mut response:serde_json::Value = serde_json::from_slice(resp_body).ok()?;
let choices = response.get_mut("choices")?.as_array_mut()?;
let base_threshold = state.compress_threshold();
let budget_mult = headroom_budget
.and_then(|b| {
let val:f64 = b.parse().ok()?;
Some((0.50 + (val / 100.0) * 0.50).clamp(0.50, 1.0))
})
.unwrap_or(1.0);
let mut did_compress = false;
for choice in choices {
let message = choice.get_mut("message")?;
if let Some(content_val) = message.get_mut("content") {
if let Some(content) = content_val.as_str() {
let ct = proxy_detect_content_type(content);
let threshold = (state.threshold_for(ct).max(base_threshold) as f64 * budget_mult) as usize;
if content.len() > threshold {
if let Some(ccr) = &state.ccr {
let hash = compute_key(content.as_bytes());
let stored = if ccr_get(ccr, &hash).await.is_some() {
state.ccr_hits.fetch_add(1, Ordering::Relaxed);
true
} else {
state.ccr_misses.fetch_add(1, Ordering::Relaxed);
let ok = ccr_put(ccr, &hash, content).await;
if ok {
state.ccr_created.fetch_add(1, Ordering::Relaxed);
} else {
tracing::error!(hash = %hash, "ccr_put failed - leaving content uncompressed to avoid data loss");
}
ok
};
if stored {
let (compressed, orig_len) = {
let compressed = match state.mode {
ProxyMode::Cache => cache_marker(&hash, content, ct, None),
ProxyMode::Token => smart_marker(&hash, content, ct, None),
};
let len = content.len();
state.record_compression(ct);
(compressed, len)
};
let marker_len = compressed.len();
state
.tokens_saved
.fetch_add(orig_len.saturating_sub(marker_len) as u64, Ordering::Relaxed);
*content_val = serde_json::Value::String(compressed);
did_compress = true;
state.update_compression_ratio(orig_len, marker_len);
}
}
} else if content.len() > state.inline_ccr_threshold() {
let hash = compute_key(content.as_bytes());
if let Ok(mut map) = state.inline_ccr.lock() {
if map.contains(&hash) {
state.inline_ccr_hits.fetch_add(1, Ordering::Relaxed);
} else {
state.inline_ccr_misses.fetch_add(1, Ordering::Relaxed);
map.put(hash, content.to_string());
}
}
}
}
}
if let Some(tool_calls_val) = message.get_mut("tool_calls") {
if let Some(arr) = tool_calls_val.as_array_mut() {
for tc in arr.iter_mut() {
if let Some(func) = tc.get_mut("function") {
if let Some(args) = func.get_mut("arguments") {
if let Some(args_str) = args.as_str() {
let args_owned = args_str.to_string(); let ct = proxy_detect_content_type(&args_owned);
let threshold =
(state.threshold_for(ct).max(base_threshold) as f64 * budget_mult) as usize;
if args_owned.len() > threshold {
if let Some(ccr) = &state.ccr {
let hash = compute_key(args_owned.as_bytes());
let stored = if ccr_get(ccr, &hash).await.is_some() {
state.ccr_hits.fetch_add(1, Ordering::Relaxed);
true
} else {
state.ccr_misses.fetch_add(1, Ordering::Relaxed);
let ok = ccr_put(ccr, &hash, &args_owned).await;
if ok {
state.ccr_created.fetch_add(1, Ordering::Relaxed);
} else {
tracing::error!(hash = %hash, "ccr_put failed - leaving tool_call arguments uncompressed to avoid data loss");
}
ok
};
if stored {
let (compressed, orig_len) = {
let compressed = smart_marker(&hash, &args_owned, ct, None);
let len = args_owned.len();
state.record_compression(ct);
(compressed, len)
};
let marker_len = compressed.len();
state.tokens_saved.fetch_add(
orig_len.saturating_sub(marker_len) as u64,
Ordering::Relaxed,
);
*args = serde_json::Value::String(compressed);
did_compress = true;
state.update_compression_ratio(orig_len, marker_len);
}
}
} else if args_owned.len() > state.inline_ccr_threshold() {
let hash = compute_key(args_owned.as_bytes());
if let Ok(mut map) = state.inline_ccr.lock() {
if map.contains(&hash) {
state.inline_ccr_hits.fetch_add(1, Ordering::Relaxed);
} else {
state.inline_ccr_misses.fetch_add(1, Ordering::Relaxed);
map.put(hash, args_owned);
}
}
}
}
}
}
}
}
}
}
if did_compress { Some(response) } else { None }
}
pub async fn handle_tool_relay(
State(state):State<Arc<AppState>>,
Json(req):Json<ToolRelayRequest>,
) -> impl IntoResponse {
state.tool_relay_calls.fetch_add(1, Ordering::Relaxed);
tracing::info!(tool = %req.tool, "tool_relay");
if req.tool == "aphrodite_retrieve" && req.params.get("hash").and_then(|v| v.as_str()).is_none() {
return (
StatusCode::BAD_REQUEST,
Json(ToolRelayResponse {
success:false,
result:None,
error:Some(
"`hash` is required for 💋/aphrodite_retrieve. Requests with only `query` and no `hash` are \
invalid."
.into(),
),
async_call:false,
}),
)
.into_response();
}
if let Some(cb) = &req.callback_url {
let parsed_url = match url::Url::parse(cb) {
Ok(u) if u.scheme() == "https" => u,
_ => {
tracing::warn!(callback_url = %cb, "tool_relay callback rejected: only https scheme allowed");
return (
StatusCode::BAD_REQUEST,
Json(ToolRelayResponse {
success:false,
result:None,
error:Some("callback_url must use the https scheme".into()),
async_call:false,
}),
)
.into_response();
},
};
let tracker = state.task_tracker.clone();
let state = state.clone();
let tool = req.tool.clone();
let params = req.params.clone();
let cb = parsed_url.to_string();
tracker.spawn(async move {
let result = execute_tool_relay(&state, &tool, ¶ms).await;
if result.is_ok() {
state.tool_relay_success.fetch_add(1, Ordering::Relaxed);
} else {
state.tool_relay_failure.fetch_add(1, Ordering::Relaxed);
}
let _ = state
.client
.post(&cb)
.json(&result)
.timeout(Duration::from_secs(5))
.send()
.await;
});
return Json(ToolRelayResponse { success:true, result:None, error:None, async_call:true }).into_response();
}
match execute_tool_relay(&state, &req.tool, &req.params).await {
Ok(val) => {
state.tool_relay_success.fetch_add(1, Ordering::Relaxed);
Json(ToolRelayResponse { success:true, result:Some(val), error:None, async_call:false }).into_response()
},
Err(e) => {
state.tool_relay_failure.fetch_add(1, Ordering::Relaxed);
Json(ToolRelayResponse { success:false, result:None, error:Some(e), async_call:false }).into_response()
},
}
}
async fn execute_tool_relay(
state:&AppState,
tool:&str,
params:&serde_json::Value,
) -> Result<serde_json::Value, String> {
match tool {
"aphrodite_retrieve" => {
let hash_raw = params.get("hash").and_then(|v| v.as_str()).ok_or("missing hash")?;
let hash = crate::marker::normalize_hash(hash_raw);
if let Ok(mut map) = state.inline_ccr.lock() {
if let Some(content) = map.get(hash) {
state.inline_ccr_hits.fetch_add(1, Ordering::Relaxed);
return Ok(serde_json::json!({"found": true, "content": content.clone()}));
}
}
state.inline_ccr_misses.fetch_add(1, Ordering::Relaxed);
if let Some(ccr) = &state.ccr {
match ccr_get(ccr, hash).await {
Some(content) => Ok(serde_json::json!({"found": true, "content": content})),
None => Ok(serde_json::json!({"found": false})),
}
} else {
Err("CCR not enabled".into())
}
},
"aphrodite_compress" => {
let content = params.get("content").and_then(|v| v.as_str()).ok_or("missing content")?;
let center = params.get("_ccr_center").and_then(|v| v.as_str());
let hash = compute_key(content.as_bytes());
let size = content.len();
if size < state.inline_ccr_threshold() {
if let Ok(mut map) = state.inline_ccr.lock() {
if map.contains(&hash) {
state.inline_ccr_hits.fetch_add(1, Ordering::Relaxed);
} else {
state.inline_ccr_misses.fetch_add(1, Ordering::Relaxed);
map.put(hash.clone(), content.to_string());
}
}
if let Some(ccr) = &state.ccr {
ccr_put(ccr, &hash, content).await;
}
Ok(serde_json::json!({
"compressed": smart_marker(&hash, content, "compress", center),
"hash": hash,
"original_size": size
}))
} else if let Some(ccr) = &state.ccr {
if !ccr_put(ccr, &hash, content).await {
return Err("failed to store content in CCR backend".into());
}
let compressed = smart_marker(&hash, content, "compress", center);
state
.tokens_saved
.fetch_add(size.saturating_sub(compressed.len()) as u64, Ordering::Relaxed);
Ok(serde_json::json!({
"compressed": compressed,
"hash": hash,
"original_size": size
}))
} else {
Err("CCR not enabled".into())
}
},
"aphrodite_list" => {
let entries = match &state.ccr {
Some(ccr) => ccr_len(ccr).await,
None => 0,
};
Ok(serde_json::json!({
"entries": entries,
"backend": match state.mode {
ProxyMode::Cache => "in_memory",
ProxyMode::Token => "sqlite",
},
}))
},
_ => Err(format!("Unknown tool: {}", tool)),
}
}
pub async fn handle_ccr_create(
State(state):State<Arc<AppState>>,
headers:axum::http::HeaderMap,
body:Bytes,
) -> impl IntoResponse {
let content_type = headers.get("content-type").and_then(|v| v.to_str().ok()).unwrap_or("");
if content_type.contains("json") {
match serde_json::from_slice::<CcrCreateRequest>(&body) {
Ok(req) => {
let original_size = req.content.len();
let hash = req.key.unwrap_or_else(|| compute_key(req.content.as_bytes()));
let ccr = match &state.ccr {
Some(ccr) => ccr,
None => {
return (
StatusCode::SERVICE_UNAVAILABLE,
Json(serde_json::json!({"error": "CCR not enabled"})),
)
.into_response();
},
};
if !ccr_put(ccr, &hash, &req.content).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(serde_json::json!({"error": "failed to store content in CCR backend"})),
)
.into_response();
}
state.ccr_created.fetch_add(1, Ordering::Relaxed);
state
.tokens_saved
.fetch_add(original_size.saturating_sub(hash.len()) as u64, Ordering::Relaxed);
state.requests_compressed.fetch_add(1, Ordering::Relaxed);
if let Some(notify_url) = &state.notify_url {
let notification = CcrNotification {
event:"ccr_created".into(),
hash:hash.clone(),
created_at:std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
ttl:req.ttl_seconds.unwrap_or(3600),
tags:req.tags.unwrap_or_default(),
};
let tracker = state.task_tracker.clone();
let client = state.client.clone();
let url = notify_url.clone();
let key = state.notify_key.clone();
let state_clone = state.clone();
tracker.spawn(async move {
let mut req = client.post(&url).json(¬ification);
if let Some(k) = &key {
req = req.header("Authorization", format!("Bearer {k}"));
}
match req.timeout(Duration::from_secs(5)).send().await {
Ok(r) if r.status().is_success() => {
state_clone.notify_success.fetch_add(1, Ordering::Relaxed);
},
_ => {
state_clone.notify_failure.fetch_add(1, Ordering::Relaxed);
},
}
});
}
let compressed_size = hash.len();
Json(CcrCreateResponse {
hash,
token_savings_ratio:if original_size > 0 {
original_size as f64 / compressed_size.max(1) as f64
} else {
1.0
},
original_size,
compressed_size,
marker_size:compressed_size,
})
.into_response()
},
Err(e) => {
(
StatusCode::BAD_REQUEST,
Json(serde_json::json!({"error": format!("invalid JSON: {}", e)})),
)
.into_response()
},
}
} else {
let content = match String::from_utf8(body.to_vec()) {
Ok(c) => c,
Err(_) => {
return (
StatusCode::BAD_REQUEST,
Json(serde_json::json!({"error": "invalid UTF-8 in body"})),
)
.into_response();
},
};
let original_size = content.len();
let hash = compute_key(content.as_bytes());
let ccr = match &state.ccr {
Some(ccr) => ccr,
None => {
return (
StatusCode::SERVICE_UNAVAILABLE,
Json(serde_json::json!({"error": "CCR not enabled"})),
)
.into_response();
},
};
if !ccr_put(ccr, &hash, &content).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(serde_json::json!({"error": "failed to store content in CCR backend"})),
)
.into_response();
}
state.ccr_created.fetch_add(1, Ordering::Relaxed);
state.requests_compressed.fetch_add(1, Ordering::Relaxed);
state
.tokens_saved
.fetch_add(original_size.saturating_sub(hash.len()) as u64, Ordering::Relaxed);
let compressed_size = hash.len();
Json(CcrCreateResponse {
hash,
token_savings_ratio:if original_size > 0 {
original_size as f64 / compressed_size.max(1) as f64
} else {
1.0
},
original_size,
compressed_size,
marker_size:compressed_size,
})
.into_response()
}
}
pub async fn handle_ccr_list(State(state):State<Arc<AppState>>) -> impl IntoResponse {
match &state.ccr {
Some(ccr) => {
let entries = ccr_len(ccr).await;
Json(serde_json::json!({
"entries": entries,
"backend": match state.mode {
ProxyMode::Cache => "in_memory",
ProxyMode::Token => "sqlite",
},
"mode": match state.mode {
ProxyMode::Cache => "cache",
ProxyMode::Token => "token",
},
}))
},
None => Json(serde_json::json!({"entries": 0, "message": "CCR not enabled"})),
}
}
pub async fn handle_ccr_delete(
State(state):State<Arc<AppState>>,
axum::extract::Path(hash):axum::extract::Path<String>,
) -> impl IntoResponse {
match &state.ccr {
Some(ccr) => {
let existed = ccr_del(ccr, &hash).await;
if existed {
(StatusCode::OK, Json(serde_json::json!({"deleted": true, "hash": hash})))
} else {
(
StatusCode::NOT_FOUND,
Json(serde_json::json!({"deleted": false, "hash": hash, "error": "not found"})),
)
}
},
None => {
(
StatusCode::SERVICE_UNAVAILABLE,
Json(serde_json::json!({"error": "CCR not enabled"})),
)
},
}
}
pub async fn handle_ccr_reload(State(state):State<Arc<AppState>>) -> impl IntoResponse {
let config_path = std::env::var("APHRODITE_CONFIG_PATH").unwrap_or_else(|_| "aphrodite.toml".to_string());
match crate::config::MultiConfig::load(&config_path) {
Ok(config) => {
let comp = config.compression.as_ref();
let thresholds = resolve_thresholds(comp);
state.cache_compress_threshold.store(thresholds.cache, Ordering::Relaxed);
state.token_compress_threshold.store(thresholds.token, Ordering::Relaxed);
state.inline_ccr_threshold.store(thresholds.inline, Ordering::Relaxed);
state
.code_multiplier_x100
.store((thresholds.code_multiplier * 100.0) as u64, Ordering::Relaxed);
let body = serde_json::json!({
"reloaded": true,
"applied": true,
"config": config_path,
"compression": {
"tool_threshold_cache": thresholds.cache,
"tool_threshold_token": thresholds.token,
"inline_threshold": thresholds.inline,
"code_multiplier": thresholds.code_multiplier,
},
"parsed_only": {
"auto_expand": comp.and_then(|c| c.auto_expand),
"auto_expand_limit": comp.and_then(|c| c.auto_expand_limit),
"terminal_threshold": comp.and_then(|c| c.terminal_threshold),
"engine_threshold_pct": comp.and_then(|c| c.engine_threshold_pct),
"catalog_mode": comp.and_then(|c| c.catalog_mode.clone()),
}
});
tracing::info!(
%config_path,
cache_threshold = thresholds.cache,
token_threshold = thresholds.token,
inline_threshold = thresholds.inline,
code_multiplier = thresholds.code_multiplier,
"config reloaded - compression thresholds applied"
);
(StatusCode::OK, Json(body)).into_response()
},
Err(e) => {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(serde_json::json!({"error": format!("failed to reload: {e}")})),
)
.into_response()
},
}
}
pub async fn health_check(State(state):State<Arc<AppState>>) -> impl IntoResponse {
let ccr_ok = state.ccr.is_some();
(
StatusCode::OK,
Json(serde_json::json!({
"status": "healthy",
"ccr": ccr_ok,
"mode": match state.mode {
ProxyMode::Cache => "cache",
ProxyMode::Token => "token",
},
"version": env!("CARGO_PKG_VERSION"),
"fill_pct": state.fill_pct.load(Ordering::Relaxed) as f64 / 100.0,
})),
)
.into_response()
}
#[cfg(test)]
pub(crate) mod tests {
use super::*;
#[test]
fn test_compress_threshold_cache() {
use std::{collections::HashMap, sync::Mutex};
let state = AppState {
client:HttpClient::new(),
api_url:"https://upstream-openai.com".into(),
model:"test".into(),
api_key:"test".into(),
ccr:None,
add_markers:false,
mode:ProxyMode::Cache,
tool_relay:false,
notify_url:None,
notify_key:None,
dev:false,
requests_total:AtomicU64::new(0),
requests_compressed:AtomicU64::new(0),
tokens_saved:AtomicU64::new(0),
ccr_hits:AtomicU64::new(0),
ccr_misses:AtomicU64::new(0),
ccr_created:AtomicU64::new(0),
tool_relay_calls:AtomicU64::new(0),
compression_ratio_ema:AtomicU64::new(200), request_history:Mutex::new(VecDeque::new()),
inline_ccr:Mutex::new(lru::LruCache::new(NonZeroUsize::new(1024).unwrap())),
latency_buckets:[
AtomicU64::new(0),
AtomicU64::new(0),
AtomicU64::new(0),
AtomicU64::new(0),
AtomicU64::new(0),
],
total_latency_micros:AtomicU64::new(0),
last_errors:Mutex::new(VecDeque::new()),
compressions_by_type:Mutex::new(HashMap::new()),
response_cache:Mutex::new(lru::LruCache::new(NonZeroUsize::new(128).unwrap())),
response_cache_ttl:std::time::Duration::from_secs(3600),
cache_hits:AtomicU64::new(0),
cache_misses:AtomicU64::new(0),
fill_pct:AtomicU64::new(9000),
task_tracker:TaskTracker::new(),
inline_ccr_hits:AtomicU64::new(0),
inline_ccr_misses:AtomicU64::new(0),
tool_relay_success:AtomicU64::new(0),
tool_relay_failure:AtomicU64::new(0),
notify_success:AtomicU64::new(0),
notify_failure:AtomicU64::new(0),
upstream_errors_4xx:AtomicU64::new(0),
upstream_errors_5xx:AtomicU64::new(0),
upstream_timeouts:AtomicU64::new(0),
upstream_connect_errors:AtomicU64::new(0),
ccr_store_entries:AtomicU64::new(0),
ccr_store_bytes:AtomicU64::new(0),
request_body_bytes:AtomicU64::new(0),
response_body_bytes:AtomicU64::new(0),
upstream_latency_micros:AtomicU64::new(0),
upstream_health_cache:std::sync::Mutex::new(None),
cache_compress_threshold:AtomicUsize::new(CACHE_COMPRESS_THRESHOLD),
token_compress_threshold:AtomicUsize::new(TOKEN_COMPRESS_THRESHOLD),
inline_ccr_threshold:AtomicUsize::new(INLINE_CCR_THRESHOLD),
code_multiplier_x100:AtomicU64::new(300),
};
assert_eq!(state.compress_threshold(), CACHE_COMPRESS_THRESHOLD);
}
#[test]
fn test_compress_threshold_aphrodite() {
let state = AppState { mode:ProxyMode::Token, ..test_state() };
assert_eq!(state.compress_threshold(), TOKEN_COMPRESS_THRESHOLD);
}
#[test]
fn test_resolve_thresholds_toml_overrides_defaults() {
let comp = CompressionConfig {
engine_threshold_pct:None,
engine_protect_first:None,
engine_protect_last:None,
engine_min_msgs:None,
tool_threshold_token:Some(512),
tool_threshold_cache:Some(4096),
terminal_threshold:None,
inline_threshold:Some(2048),
auto_expand:None,
auto_expand_limit:None,
catalog_mode:None,
classifier_poll:None,
code_multiplier:Some(5.0),
};
let t = resolve_thresholds(Some(&comp));
assert_eq!(t.cache, 4096);
assert_eq!(t.token, 512);
assert_eq!(t.inline, 2048);
assert_eq!(t.code_multiplier, 5.0);
}
#[test]
fn test_resolve_thresholds_defaults_when_no_toml() {
let t = resolve_thresholds(None);
assert_eq!(t.cache, CACHE_COMPRESS_THRESHOLD);
assert_eq!(t.token, TOKEN_COMPRESS_THRESHOLD);
assert_eq!(t.inline, INLINE_CCR_THRESHOLD);
assert_eq!(t.code_multiplier, 3.0);
}
#[test]
fn test_handle_ccr_reload_applies_thresholds_to_state() {
let dir = std::env::temp_dir();
let path =
dir.join(format!("aphrodite_reload_test_{}_{}.toml", std::process::id(), fnv1a_64(b"reload-test-salt")));
std::fs::write(
&path,
r#"
[[proxies]]
name = "token"
mode = "token"
[compression]
tool_threshold_token = 999
tool_threshold_cache = 1234
inline_threshold = 77
code_multiplier = 6.5
"#,
)
.unwrap();
std::env::set_var("APHRODITE_CONFIG_PATH", &path);
let state = std::sync::Arc::new(test_state());
let rt = tokio::runtime::Runtime::new().unwrap();
let resp = rt.block_on(handle_ccr_reload(State(state.clone()))).into_response();
std::env::remove_var("APHRODITE_CONFIG_PATH");
let _ = std::fs::remove_file(&path);
assert_eq!(resp.status(), axum::http::StatusCode::OK);
assert_eq!(state.token_compress_threshold.load(Ordering::Relaxed), 999);
assert_eq!(state.cache_compress_threshold.load(Ordering::Relaxed), 1234);
assert_eq!(state.inline_ccr_threshold.load(Ordering::Relaxed), 77);
assert_eq!(state.code_multiplier_x100.load(Ordering::Relaxed), 650);
}
#[test]
fn test_stats_json_modes() {
let cache = test_state();
let stats = cache.stats_json();
assert_eq!(stats["mode"], "cache");
assert_eq!(stats["proxy"], "aphrodite");
let mut aph = test_state();
aph.mode = ProxyMode::Token;
let stats = aph.stats_json();
assert_eq!(stats["mode"], "token");
}
#[test]
fn test_stats_json_tool_relay_enabled_flag_not_shadowed() {
let mut state = test_state();
state.tool_relay = true;
let stats = state.stats_json();
assert_eq!(stats["tool_relay_enabled"], true);
assert!(
stats["tool_relay"].is_object(),
"the calls-stats object must still be present under its own key"
);
assert!(stats["tool_relay"]["total"].is_u64());
}
#[test]
fn test_ccr_create_response_serde_shape() {
let resp = CcrCreateResponse {
hash:"abc123".into(),
token_savings_ratio:2.5,
original_size:100,
compressed_size:40,
marker_size:40,
};
let v = serde_json::to_value(&resp).unwrap();
assert_eq!(v["hash"], "abc123");
assert_eq!(v["original_size"], 100);
assert_eq!(v["compressed_size"], 40);
assert_eq!(v["marker_size"], 40);
assert!((v["token_savings_ratio"].as_f64().unwrap() - 2.5).abs() < 0.01);
assert!(v.get("compression_ratio").is_none());
}
#[test]
fn test_tool_relay_response_sync_serde_shape() {
let resp = ToolRelayResponse {
success:true,
result:Some(serde_json::json!({"found": true})),
error:None,
async_call:false,
};
let v = serde_json::to_value(&resp).unwrap();
assert_eq!(v["success"], true);
assert_eq!(v["async_call"], false);
assert_eq!(v["result"]["found"], true);
assert!(v["error"].is_null());
}
#[test]
fn test_tool_relay_response_async_serde_shape() {
let resp = ToolRelayResponse { success:true, result:None, error:None, async_call:true };
let v = serde_json::to_value(&resp).unwrap();
assert_eq!(v["async_call"], true);
assert!(v["result"].is_null());
}
#[test]
fn test_detect_content_type_json_tool_output() {
assert_eq!(proxy_detect_content_type(r#"{"exit_code": 0, "output": "ok"}"#), "tool_output");
}
#[test]
fn test_detect_content_type_invalid_json_is_text() {
assert_eq!(proxy_detect_content_type("{ not json at all"), "text");
}
#[test]
fn test_detect_content_type_json_array() {
assert_eq!(proxy_detect_content_type(r#"[{"a":1},{"a":2}]"#), "json");
}
#[test]
fn test_detect_content_type_rust_code() {
let src = "use std::fmt;\nfn add(a:i32, b:i32) -> i32 {\n a + b\n}\n";
assert_eq!(proxy_detect_content_type(src), "code_rust");
}
#[test]
fn test_detect_content_type_python_code() {
let src = "import os\nclass Foo:\n def bar(self):\n pass\n";
assert_eq!(proxy_detect_content_type(src), "code_python");
}
#[test]
fn test_detect_content_type_go_code() {
let src = "package main\nimport (\n\t\"fmt\"\n)\nfunc main() {\n\tfmt.Println(\"hi\")\n}\n";
assert_eq!(proxy_detect_content_type(src), "code_go");
}
#[test]
fn test_detect_content_type_js_code() {
let src = "import { foo } from 'bar';\nexport const add = (a, b) => a + b;\nconst x = 1;\nconst y = 2;\n";
assert_eq!(proxy_detect_content_type(src), "code_js");
}
#[test]
fn test_detect_content_type_error_first_line() {
assert_eq!(
proxy_detect_content_type("Traceback (most recent call last):\n File \"x.py\", line 1\nValueError: bad\n"),
"error"
);
}
#[test]
fn test_detect_content_type_diff() {
let d = "diff --git a/src/lib.rs b/src/lib.rs\n--- a/src/lib.rs\n+++ b/src/lib.rs\n@@ -1,3 +1,4 @@\n+added a \
line\n";
assert_eq!(proxy_detect_content_type(d), "diff");
}
#[test]
fn test_detect_content_type_log_lines() {
let log = "starting up\n[INFO] service ready\n[WARN] disk low\n[ERROR] connection lost\n";
assert_eq!(proxy_detect_content_type(log), "log");
}
#[test]
fn test_detect_content_type_empty_is_text() {
assert_eq!(proxy_detect_content_type(""), "text");
}
#[test]
fn test_detect_content_type_plain_text() {
assert_eq!(proxy_detect_content_type("just some plain text\nnothing special\n"), "text");
}
#[test]
fn test_generate_metadata_rust_has_lang_and_fns() {
let src = "fn add(a:i32, b:i32) -> i32 {\n a + b\n}\n";
let meta = generate_metadata(src, "code_rust");
assert!(meta.contains("lang=rs"));
assert!(meta.contains("fns=add"));
}
#[test]
fn test_generate_metadata_escapes_pipes() {
let src = "On branch feature|weird\n";
let meta = generate_metadata(src, "git");
assert!(!meta.contains('|'), "metadata must not contain a raw pipe: {meta}");
}
#[test]
fn test_generate_metadata_max_400_chars() {
let src = (0..100).map(|i| format!("fn f{i}() {{}}")).collect::<Vec<_>>().join("\n");
let meta = generate_metadata(&src, "code_rust");
assert!(meta.chars().count() <= 400, "metadata too long: {} chars", meta.chars().count());
}
#[test]
fn test_generate_metadata_error_branch_no_panic_on_multibyte_utf8() {
let src = "日本語エラー at src/日本.rs:10:5 something\n";
let meta = generate_metadata(src, "error");
assert!(meta.is_empty() || meta.contains("trace=") || meta.contains("msg="));
}
#[test]
fn test_generate_metadata_text_has_line_count() {
let meta = generate_metadata("a\nb\nc\n", "text");
assert_eq!(meta, "ln=3");
}
#[test]
fn test_build_preview_code_has_ct_prefix() {
let src = "fn add(a:i32, b:i32) -> i32 {\n a + b\n}\n";
let preview = proxy_build_preview(src, "code_rust");
assert!(preview.starts_with("[code_rust:"));
}
#[test]
fn test_build_preview_error_has_ct_prefix_via_error_line() {
let src = "some noise\nerror[E0308]: mismatched types\nmore noise\n";
let preview = proxy_build_preview(src, "error");
assert!(preview.contains("error[E0308]"));
}
#[test]
fn test_build_preview_diff_has_ct_prefix() {
let src = "diff --git a/x b/x\n--- a/x\n+++ a/x\n";
let preview = proxy_build_preview(src, "diff");
assert!(preview.starts_with("diff --git"));
}
#[test]
fn test_build_preview_json_has_ct_prefix() {
let src = "{\"a\":1,\"b\":2}\n";
let preview = proxy_build_preview(src, "json");
assert!(preview.contains("keys"));
}
#[test]
fn test_cache_key_from_body_deterministic() {
let body = br#"{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}"#;
let k1 = cache_key_from_body(body, "key-a");
let k2 = cache_key_from_body(body, "key-a");
assert!(k1.is_some());
assert_eq!(k1, k2);
}
#[test]
fn test_cache_key_from_body_differs_by_api_key() {
let body = br#"{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}"#;
let k1 = cache_key_from_body(body, "key-a");
let k2 = cache_key_from_body(body, "key-b");
assert_ne!(k1, k2);
}
#[test]
fn test_cache_key_from_body_none_on_junk() {
assert_eq!(cache_key_from_body(b"not json", "key"), None);
assert_eq!(cache_key_from_body(b"{}", "key"), None); }
#[test]
fn test_cache_key_from_body_differs_by_tools_and_temperature_and_stream() {
let base = br#"{"model":"gpt-4","messages":[{"role":"user","content":"hi"}]}"#;
let with_tools =
br#"{"model":"gpt-4","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function"}]}"#;
let with_temp = br#"{"model":"gpt-4","messages":[{"role":"user","content":"hi"}],"temperature":0.7}"#;
let k_base = cache_key_from_body(base, "key");
let k_tools = cache_key_from_body(with_tools, "key");
let k_temp = cache_key_from_body(with_temp, "key");
assert!(k_base.is_some());
assert_ne!(k_base, k_tools, "differing `tools` must produce a different cache key");
assert_ne!(k_base, k_temp, "differing `temperature` must produce a different cache key");
assert_ne!(k_tools, k_temp);
}
#[test]
fn test_cache_key_from_body_none_when_streaming() {
let streamed = br#"{"model":"gpt-4","messages":[{"role":"user","content":"hi"}],"stream":true}"#;
assert_eq!(cache_key_from_body(streamed, "key"), None);
}
#[test]
fn test_response_cache_get_expires_past_ttl() {
let mut state = test_state();
state.response_cache_ttl = std::time::Duration::from_millis(1);
state.response_cache.lock().unwrap().put(42, (std::time::Instant::now(), b"cached".to_vec()));
std::thread::sleep(std::time::Duration::from_millis(20));
assert_eq!(response_cache_get(&state, 42), None, "expired entry must not be returned");
assert!(
state.response_cache.lock().unwrap().peek(&42).is_none(),
"expired entry must be evicted, not just skipped"
);
}
#[test]
fn test_response_cache_get_hits_within_ttl() {
let mut state = test_state();
state.response_cache_ttl = std::time::Duration::from_secs(3600);
state.response_cache.lock().unwrap().put(7, (std::time::Instant::now(), b"cached".to_vec()));
assert_eq!(response_cache_get(&state, 7), Some(b"cached".to_vec()));
}
#[test]
fn test_fnv1a_64_known_vectors() {
assert_eq!(fnv1a_64(b""), 14695981039346656037);
assert_ne!(fnv1a_64(b"a"), fnv1a_64(b"b"));
}
fn test_state() -> AppState {
use std::{collections::HashMap, sync::Mutex};
AppState {
client:HttpClient::new(),
api_url:"https://upstream-openai.com".into(),
model:"default-model".into(),
api_key:"test".into(),
ccr:None,
add_markers:false,
mode:ProxyMode::Cache,
tool_relay:false,
notify_url:None,
notify_key:None,
dev:false,
requests_total:AtomicU64::new(0),
requests_compressed:AtomicU64::new(0),
tokens_saved:AtomicU64::new(0),
ccr_hits:AtomicU64::new(0),
ccr_misses:AtomicU64::new(0),
ccr_created:AtomicU64::new(0),
tool_relay_calls:AtomicU64::new(0),
compression_ratio_ema:AtomicU64::new(200), request_history:Mutex::new(VecDeque::new()),
inline_ccr:Mutex::new(lru::LruCache::new(NonZeroUsize::new(1024).unwrap())),
latency_buckets:[
AtomicU64::new(0),
AtomicU64::new(0),
AtomicU64::new(0),
AtomicU64::new(0),
AtomicU64::new(0),
],
total_latency_micros:AtomicU64::new(0),
last_errors:Mutex::new(VecDeque::new()),
compressions_by_type:Mutex::new(HashMap::new()),
response_cache:Mutex::new(lru::LruCache::new(NonZeroUsize::new(128).unwrap())),
response_cache_ttl:std::time::Duration::from_secs(3600),
cache_hits:AtomicU64::new(0),
cache_misses:AtomicU64::new(0),
fill_pct:AtomicU64::new(9000),
task_tracker:TaskTracker::new(),
inline_ccr_hits:AtomicU64::new(0),
inline_ccr_misses:AtomicU64::new(0),
tool_relay_success:AtomicU64::new(0),
tool_relay_failure:AtomicU64::new(0),
notify_success:AtomicU64::new(0),
notify_failure:AtomicU64::new(0),
upstream_errors_4xx:AtomicU64::new(0),
upstream_errors_5xx:AtomicU64::new(0),
upstream_timeouts:AtomicU64::new(0),
upstream_connect_errors:AtomicU64::new(0),
ccr_store_entries:AtomicU64::new(0),
ccr_store_bytes:AtomicU64::new(0),
request_body_bytes:AtomicU64::new(0),
response_body_bytes:AtomicU64::new(0),
upstream_latency_micros:AtomicU64::new(0),
upstream_health_cache:std::sync::Mutex::new(None),
cache_compress_threshold:AtomicUsize::new(CACHE_COMPRESS_THRESHOLD),
token_compress_threshold:AtomicUsize::new(TOKEN_COMPRESS_THRESHOLD),
inline_ccr_threshold:AtomicUsize::new(INLINE_CCR_THRESHOLD),
code_multiplier_x100:AtomicU64::new(300),
}
}
pub(crate) fn test_state_with_ccr() -> AppState {
AppState {
ccr:Some(std::sync::Arc::new(InMemoryCcrStore::with_capacity_and_ttl(
1000,
std::time::Duration::from_secs(300),
))),
mode:ProxyMode::Token,
..test_state()
}
}
struct FailingCcrStore;
impl headroom_core::ccr::CcrStore for FailingCcrStore {
fn put(&self, _hash:&str, _payload:&str) -> bool { false }
fn get(&self, _hash:&str) -> Option<String> { None }
fn len(&self) -> usize { 0 }
fn del(&self, _hash:&str) -> bool { false }
}
fn test_state_with_failing_ccr() -> AppState {
AppState {
ccr:Some(std::sync::Arc::new(FailingCcrStore)),
mode:ProxyMode::Token,
..test_state()
}
}
#[test]
fn test_compress_chat_completion_ccr_put_failure_leaves_content_uncompressed() {
let state = test_state_with_failing_ccr();
let content = "fn answer() -> i32 { 42 }\n".repeat(200); let body = chat_completion_body(&content);
let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(compress_chat_completion(&state, &body, None));
assert!(result.is_none(), "a failed ccr_put must not produce a compressed response");
}
#[test]
fn test_handle_ccr_create_503_when_ccr_disabled() {
let mut state = test_state();
state.ccr = None;
let state = std::sync::Arc::new(state);
let rt = tokio::runtime::Runtime::new().unwrap();
let resp = rt
.block_on(handle_ccr_create(
State(state),
axum::http::HeaderMap::new(),
Bytes::from_static(b"hello world"),
))
.into_response();
assert_eq!(resp.status(), axum::http::StatusCode::SERVICE_UNAVAILABLE);
}
#[test]
fn test_handle_ccr_create_500_when_put_fails() {
let state = std::sync::Arc::new(test_state_with_failing_ccr());
let rt = tokio::runtime::Runtime::new().unwrap();
let resp = rt
.block_on(handle_ccr_create(
State(state),
axum::http::HeaderMap::new(),
Bytes::from_static(b"hello world"),
))
.into_response();
assert_eq!(resp.status(), axum::http::StatusCode::INTERNAL_SERVER_ERROR);
}
fn chat_completion_body(content:&str) -> Vec<u8> {
serde_json::json!({
"choices": [{
"message": {"role": "assistant", "content": content}
}]
})
.to_string()
.into_bytes()
}
#[test]
fn test_compress_chat_completion_above_threshold_produces_marker() {
let content = "the quick brown fox jumps over the lazy dog. ".repeat(200);
let body = chat_completion_body(&content);
let state = test_state_with_ccr();
let result = tokio::runtime::Runtime::new()
.unwrap()
.block_on(compress_chat_completion(&state, &body, None));
let response = result.expect("content above threshold must be compressed");
let new_content = response["choices"][0]["message"]["content"].as_str().unwrap();
assert!(new_content.contains("<<<CCR:"), "expected a CCR marker, got: {new_content}");
let ccr = state.ccr.as_ref().unwrap().clone();
let hash = compute_key(content.as_bytes());
let stored = tokio::runtime::Runtime::new().unwrap().block_on(ccr_get(&ccr, &hash));
assert_eq!(stored.as_deref(), Some(content.as_str()));
}
#[test]
fn test_compress_chat_completion_below_threshold_is_none() {
let content = "short reply";
let body = chat_completion_body(content);
let state = test_state_with_ccr();
let result = tokio::runtime::Runtime::new()
.unwrap()
.block_on(compress_chat_completion(&state, &body, None));
assert!(result.is_none(), "short content must not be compressed: {result:?}");
}
#[test]
fn test_compress_chat_completion_tool_call_arguments_compressed_independently() {
let big_args = serde_json::json!({"data": "x".repeat(4000)}).to_string();
let body = serde_json::json!({
"choices": [{
"message": {
"role": "assistant",
"content": "ok",
"tool_calls": [{
"function": {"name": "f", "arguments": big_args}
}]
}
}]
})
.to_string()
.into_bytes();
let state = test_state_with_ccr();
let result = tokio::runtime::Runtime::new()
.unwrap()
.block_on(compress_chat_completion(&state, &body, None));
let response = result.expect("large tool-call arguments must be compressed");
assert_eq!(response["choices"][0]["message"]["content"], "ok");
let new_args = response["choices"][0]["message"]["tool_calls"][0]["function"]["arguments"]
.as_str()
.unwrap();
assert!(
new_args.contains("<<<CCR:"),
"expected tool-call arguments to be compressed: {new_args}"
);
}
#[test]
fn test_compress_chat_completion_budget_header_lowers_effective_threshold() {
let content = "line of moderate length text content here.\n".repeat(40); let body = chat_completion_body(&content);
let state_no_budget = test_state_with_ccr();
let result_no_budget =
tokio::runtime::Runtime::new()
.unwrap()
.block_on(compress_chat_completion(&state_no_budget, &body, None));
let state_low_budget = test_state_with_ccr();
let result_low_budget = tokio::runtime::Runtime::new().unwrap().block_on(compress_chat_completion(
&state_low_budget,
&body,
Some("0"),
));
if result_no_budget.is_some() {
assert!(
result_low_budget.is_some(),
"lower budget must compress at least as much as no budget"
);
}
}
#[test]
fn test_execute_tool_relay_retrieve_missing_hash_param() {
let state = test_state_with_ccr();
let result = tokio::runtime::Runtime::new().unwrap().block_on(execute_tool_relay(
&state,
"aphrodite_retrieve",
&serde_json::json!({}),
));
assert_eq!(result, Err("missing hash".to_string()));
}
#[test]
fn test_execute_tool_relay_unknown_tool_is_err() {
let state = test_state_with_ccr();
let result = tokio::runtime::Runtime::new().unwrap().block_on(execute_tool_relay(
&state,
"not_a_real_tool",
&serde_json::json!({}),
));
assert!(result.is_err());
}
#[test]
fn test_execute_tool_relay_compress_small_content_stores_inline_and_returns_marker() {
let state = test_state_with_ccr();
let content = "tiny"; let result = tokio::runtime::Runtime::new().unwrap().block_on(execute_tool_relay(
&state,
"aphrodite_compress",
&serde_json::json!({"content": content}),
));
let v = result.expect("compress must succeed");
assert!(v["compressed"].as_str().unwrap().contains("<<<CCR:"));
assert_eq!(v["original_size"], content.len());
}
#[test]
fn test_execute_tool_relay_compress_small_content_also_stores_durably() {
let state = test_state_with_ccr();
let content = "tiny"; let rt = tokio::runtime::Runtime::new().unwrap();
let result = rt.block_on(execute_tool_relay(&state, "aphrodite_compress", &serde_json::json!({"content": content})));
let v = result.expect("compress must succeed");
let hash = v["hash"].as_str().unwrap().to_string();
let ccr = state.ccr.as_ref().unwrap();
let durable = rt.block_on(ccr_get(ccr, &hash));
assert_eq!(durable.as_deref(), Some(content), "tiny content must also be durable, not inline-only");
}
#[test]
fn test_execute_tool_relay_retrieve_finds_inline_entry() {
let state = test_state_with_ccr();
let content = "tiny";
let compressed = tokio::runtime::Runtime::new()
.unwrap()
.block_on(execute_tool_relay(
&state,
"aphrodite_compress",
&serde_json::json!({"content": content}),
))
.unwrap();
let hash = compressed["hash"].as_str().unwrap().to_string();
let retrieved = tokio::runtime::Runtime::new()
.unwrap()
.block_on(execute_tool_relay(
&state,
"aphrodite_retrieve",
&serde_json::json!({"hash": hash}),
))
.unwrap();
assert_eq!(retrieved["found"], true);
assert_eq!(retrieved["content"], content);
}
#[test]
fn test_execute_tool_relay_retrieve_normalizes_pipe_suffixed_and_whitespace_hash() {
let state = test_state_with_ccr();
let content = "tiny";
let rt = tokio::runtime::Runtime::new().unwrap();
let compressed = rt
.block_on(execute_tool_relay(
&state,
"aphrodite_compress",
&serde_json::json!({"content": content}),
))
.unwrap();
let hash = compressed["hash"].as_str().unwrap().to_string();
for hash_arg in [hash.clone(), format!("{hash}|tool|1024"), format!(" {hash} ")] {
let retrieved = rt
.block_on(execute_tool_relay(
&state,
"aphrodite_retrieve",
&serde_json::json!({"hash": hash_arg}),
))
.unwrap();
assert_eq!(retrieved["found"], true, "hash arg {hash_arg:?} must resolve: {retrieved:?}");
assert_eq!(retrieved["content"], content);
}
}
#[test]
fn regression_07_tokens_saved_increments_on_compress() {
let content = "the quick brown fox jumps over the lazy dog. ".repeat(200);
let body = chat_completion_body(&content);
let state = test_state_with_ccr();
assert_eq!(state.tokens_saved.load(Ordering::Relaxed), 0);
let result = tokio::runtime::Runtime::new()
.unwrap()
.block_on(compress_chat_completion(&state, &body, None));
assert!(result.is_some(), "content above threshold must compress");
assert!(
state.tokens_saved.load(Ordering::Relaxed) > 0,
"tokens_saved must be incremented by the real compression path"
);
}
#[test]
fn regression_11_below_threshold_skips_compression_and_counter() {
let content = "short reply below any threshold";
let body = chat_completion_body(content);
let state = test_state_with_ccr();
let result = tokio::runtime::Runtime::new()
.unwrap()
.block_on(compress_chat_completion(&state, &body, None));
assert!(result.is_none(), "below-threshold content must not compress");
assert_eq!(
state.tokens_saved.load(Ordering::Relaxed),
0,
"no savings should be recorded when nothing was compressed"
);
}
#[test]
fn regression_13_marker_terminator_never_truncated() {
let hash = "abc123def456abc123def456abc123def456";
let huge_preview = "x".repeat(10_000);
let huge_metadata = "y".repeat(10_000);
let out = proxy_format_ccr_output(&huge_preview, "text", &huge_metadata, None, hash, 123456);
let expected_terminator = format!("<<<CCR:{hash}|text|123456>>>");
assert!(
out.contains(&expected_terminator),
"marker terminator must always be complete and unsliced, regardless of preview/metadata length"
);
}
use proptest::{prop_assert, proptest};
proptest! {
#[test]
fn prop_classifier_and_metadata_never_panic(s in ".*") {
let ct = proxy_detect_content_type(&s);
let meta = generate_metadata(&s, ct);
prop_assert!(!meta.contains('|'));
prop_assert!(!meta.contains('\n'));
prop_assert!(meta.chars().count() <= 400);
}
}
}