diff --git a/crates/proxy/src/cache/memory.rs b/crates/proxy/src/cache/memory.rs index 3a9f71d..6f4b246 100644 --- a/crates/proxy/src/cache/memory.rs +++ b/crates/proxy/src/cache/memory.rs @@ -2,14 +2,41 @@ //! //! moka provides a concurrent, lock-free cache with TTL-based expiration //! and bounded capacity (LRU eviction when full). +//! +//! Per-entry TTL is enforced via moka's `Expiry` trait. Each `CacheEntry` +//! carries an optional `ttl_secs` override; when absent, the cache-level +//! default applies. use super::{CacheBackend, CacheConfig, CacheEntry}; +use moka::Expiry; use std::time::Duration; +/// Per-entry expiry policy. Reads `CacheEntry::ttl_secs` to decide lifetime; +/// falls back to `default_ttl` when the entry has no override. +struct EntryExpiry { + default_ttl: Duration, +} + +impl Expiry for EntryExpiry { + fn expire_after_create( + &self, + _key: &String, + value: &CacheEntry, + _current_time: std::time::Instant, + ) -> Option { + let ttl = match value.ttl_secs { + Some(secs) => Duration::from_secs(secs), + None => self.default_ttl, + }; + Some(ttl) + } +} + /// In-memory cache backed by moka::future::Cache. /// /// Configured with a default TTL and max entry count. Per-request TTL -/// overrides are applied at insert time via moka's `insert_with_expiry` API. +/// overrides are enforced via the `EntryExpiry` implementation of moka's +/// `Expiry` trait. pub struct MemoryCache { inner: moka::future::Cache, /// Default TTL applied when the request does not specify cache_ttl_secs. @@ -19,9 +46,10 @@ pub struct MemoryCache { impl MemoryCache { /// Create a new in-memory cache from the provided configuration. pub fn new(config: &CacheConfig) -> Self { + let default_ttl = Duration::from_secs(config.ttl_secs); let inner = moka::future::Cache::builder() .max_capacity(config.max_entries) - .time_to_live(Duration::from_secs(config.ttl_secs)) + .expire_after(EntryExpiry { default_ttl }) .build(); Self { inner, @@ -35,18 +63,10 @@ impl CacheBackend for MemoryCache { self.inner.get(key).await } - async fn put(&self, key: &str, entry: CacheEntry, ttl_secs: u64) { - // moka does not support per-entry TTL at insert time via the standard - // insert API. We use the expiry API by configuring the cache with the - // default TTL. For per-request TTL, we use policy-level TTL which - // applies to all entries. This is acceptable because most requests - // use the default. A future enhancement could use moka's Expiry trait. - // - // For now, if ttl_secs differs from the default, we still insert - // (the cache-level TTL applies). This is a pragmatic tradeoff: the - // entry may live slightly longer or shorter than requested, but cache - // correctness is not compromised (stale data is acceptable in caching). - let _ = ttl_secs; // Acknowledged but not separately enforced per-entry. + async fn put(&self, key: &str, entry: CacheEntry, _ttl_secs: u64) { + // Per-entry TTL is now handled by EntryExpiry reading entry.ttl_secs. + // The _ttl_secs parameter from CacheBackend::put is unused; the entry + // itself carries the authoritative TTL override. self.inner.insert(key.to_string(), entry).await; } } @@ -70,6 +90,16 @@ mod tests { response_body: Bytes::from(body.to_string()), model: "test-model".to_string(), created_at: Instant::now(), + ttl_secs: None, + } + } + + fn test_entry_with_ttl(body: &str, ttl: u64) -> CacheEntry { + CacheEntry { + response_body: Bytes::from(body.to_string()), + model: "test-model".to_string(), + created_at: Instant::now(), + ttl_secs: Some(ttl), } } @@ -128,4 +158,54 @@ mod tests { // grow unbounded beyond the configured max_capacity. assert!(cache.inner.entry_count() <= 3); } + + #[tokio::test] + async fn per_entry_ttl_shorter_than_default() { + // Default TTL is 10s, but entry requests 1s. + let config = CacheConfig { + ttl_secs: 10, + max_entries: 100, + redis_url: None, + }; + let cache = MemoryCache::new(&config); + let entry = test_entry_with_ttl("short-lived", 1); + cache.put("test:short", entry, 1).await; + + // Present immediately + assert!(cache.get("test:short").await.is_some()); + + // Expired after 1.5s (entry TTL = 1s) + tokio::time::sleep(Duration::from_millis(1500)).await; + assert!( + cache.get("test:short").await.is_none(), + "entry with 1s TTL should be expired after 1.5s despite 10s default" + ); + } + + #[tokio::test] + async fn per_entry_ttl_longer_than_default() { + // Default TTL is 1s, but entry requests 3s. + let config = CacheConfig { + ttl_secs: 1, + max_entries: 100, + redis_url: None, + }; + let cache = MemoryCache::new(&config); + let entry = test_entry_with_ttl("long-lived", 3); + cache.put("test:long", entry, 3).await; + + // Still alive after 1.5s (past the default 1s) + tokio::time::sleep(Duration::from_millis(1500)).await; + assert!( + cache.get("test:long").await.is_some(), + "entry with 3s TTL should survive past the 1s default" + ); + + // Expired after 3.5s + tokio::time::sleep(Duration::from_millis(2000)).await; + assert!( + cache.get("test:long").await.is_none(), + "entry with 3s TTL should be expired after 3.5s" + ); + } } diff --git a/crates/proxy/src/cache/mod.rs b/crates/proxy/src/cache/mod.rs index 780f409..646b4ac 100644 --- a/crates/proxy/src/cache/mod.rs +++ b/crates/proxy/src/cache/mod.rs @@ -9,6 +9,11 @@ //! - `oai:` for /v1/chat/completions pub mod memory; +/// Redis L2 cache backend (requires `redis` feature). +pub mod redis; +/// Semantic cache backed by Qdrant vector store (requires `qdrant` feature). +#[cfg(feature = "qdrant")] +pub mod semantic; use bytes::Bytes; use sha2::{Digest, Sha256}; @@ -27,6 +32,9 @@ pub struct CacheEntry { pub model: String, /// When this entry was created (wall-clock, not persisted to Redis). pub created_at: Instant, + /// Per-entry TTL override in seconds. When set, moka's Expiry trait + /// uses this instead of the cache-level default. + pub ttl_secs: Option, } /// Namespace prefix for cache keys, preventing cross-endpoint collisions. diff --git a/crates/proxy/src/cache/redis.rs b/crates/proxy/src/cache/redis.rs index fd523d2..0285f36 100644 --- a/crates/proxy/src/cache/redis.rs +++ b/crates/proxy/src/cache/redis.rs @@ -1,12 +1,99 @@ -//! Redis cache backend (future implementation). +//! Redis L2 cache backend. //! -//! This module is a placeholder for an L2 cache backed by Redis. -//! It will be implemented behind a `redis` feature flag when the -//! `redis` crate dependency is added. +//! Provides a `RedisCache` that implements `CacheBackend` for use as an +//! L2 cache behind the in-memory moka cache. Feature-gated behind `redis`. //! -//! Design notes: -//! - Will use `redis::aio::ConnectionManager` for async, pooled connections. -//! - CacheEntry will be serialized via serde_json for storage. -//! - Graceful fallback: if Redis is unreachable, log error and continue -//! (memory cache still serves as L1). -//! - TTL is set via Redis SETEX/PSETEX, giving true per-entry expiration. +//! Graceful fallback: if Redis is unreachable, operations return None/no-op +//! and log a warning. The in-memory cache still serves as L1. + +#[cfg(feature = "redis")] +use redis::aio::ConnectionManager; + +#[cfg(feature = "redis")] +use super::CacheEntry; + +/// Redis-backed response cache using SETEX for per-entry TTL. +#[cfg(feature = "redis")] +pub struct RedisCache { + conn: ConnectionManager, + /// Key prefix to namespace cache entries. + prefix: String, +} + +#[cfg(feature = "redis")] +impl RedisCache { + /// Create a new Redis cache from an existing connection manager. + pub fn new(conn: ConnectionManager) -> Self { + Self { + conn, + prefix: "anyllm:cache:".to_string(), + } + } + + /// Connect to Redis and create a cache. + pub async fn connect(redis_url: &str) -> Result { + let client = redis::Client::open(redis_url)?; + let conn = ConnectionManager::new(client).await?; + Ok(Self::new(conn)) + } + + fn redis_key(&self, key: &str) -> String { + format!("{}{}", self.prefix, key) + } + + /// Get a cached entry from Redis. + pub async fn get(&self, key: &str) -> Option { + let redis_key = self.redis_key(key); + let mut conn = self.conn.clone(); + let result: Result, redis::RedisError> = redis::cmd("GET") + .arg(&redis_key) + .query_async(&mut conn) + .await; + match result { + Ok(Some(json)) => serde_json::from_str::(&json) + .ok() + .map(|v| CacheEntry { + response_body: bytes::Bytes::from(v.response_body), + model: v.model, + created_at: std::time::Instant::now(), + ttl_secs: None, // Redis manages its own TTL via SETEX + }), + Ok(None) => None, + Err(e) => { + tracing::warn!(error = %e, "Redis cache GET failed"); + None + } + } + } + + /// Store a cache entry in Redis with the given TTL. + pub async fn put(&self, key: &str, entry: &CacheEntry, ttl_secs: u64) { + let redis_key = self.redis_key(key); + let value = RedisCacheValue { + response_body: String::from_utf8_lossy(&entry.response_body).to_string(), + model: entry.model.clone(), + }; + let json = match serde_json::to_string(&value) { + Ok(j) => j, + Err(_) => return, + }; + let mut conn = self.conn.clone(); + let result: Result<(), redis::RedisError> = redis::cmd("SETEX") + .arg(&redis_key) + .arg(ttl_secs) + .arg(&json) + .query_async(&mut conn) + .await; + if let Err(e) = result { + tracing::warn!(error = %e, "Redis cache SETEX failed"); + } + } +} + +/// Serializable value stored in Redis. +#[cfg(feature = "redis")] +#[derive(serde::Serialize, serde::Deserialize)] +struct RedisCacheValue { + response_body: String, + model: String, +} diff --git a/crates/proxy/src/cache/semantic.rs b/crates/proxy/src/cache/semantic.rs index e4b806c..5c1dded 100644 --- a/crates/proxy/src/cache/semantic.rs +++ b/crates/proxy/src/cache/semantic.rs @@ -3,33 +3,35 @@ //! Requires `--features qdrant` and `QDRANT_URL` env var. //! When `QDRANT_URL` is not set, `SemanticCache::new()` returns `None` //! and the proxy falls back to exact-match caching only. +//! +//! The caller is responsible for generating embeddings (via the backend's +//! embedding endpoint). This module handles only vector store operations. -use qdrant_client::Qdrant; +use qdrant_client::qdrant::{ + CreateCollectionBuilder, Distance, PointStruct, SearchPointsBuilder, UpsertPointsBuilder, + VectorParamsBuilder, +}; +use qdrant_client::{Payload, Qdrant}; +use std::sync::atomic::{AtomicBool, Ordering}; -/// Semantic cache that embeds prompts and searches for similar cached responses. -/// -/// Wraps a Qdrant client. The actual embedding generation is the caller's -/// responsibility (e.g., via the backend's embedding endpoint). This struct -/// handles only the vector store operations. +/// Semantic cache that stores and searches response embeddings in Qdrant. pub struct SemanticCache { - #[allow(dead_code)] client: Qdrant, - #[allow(dead_code)] collection: String, - #[allow(dead_code)] threshold: f32, + /// Whether the collection has been verified/created. + collection_ready: AtomicBool, } impl SemanticCache { /// Create a new semantic cache connected to Qdrant. /// /// Returns `None` if `QDRANT_URL` is not set, enabling graceful - /// degradation: the proxy starts without semantic caching and logs a - /// warning at the call site. + /// degradation: the proxy starts without semantic caching. pub fn new() -> Option { let url = std::env::var("QDRANT_URL").ok()?; - let collection = std::env::var("QDRANT_COLLECTION") - .unwrap_or_else(|_| "anyllm_cache".to_string()); + let collection = + std::env::var("QDRANT_COLLECTION").unwrap_or_else(|_| "anyllm_cache".to_string()); let threshold: f32 = std::env::var("SEMANTIC_CACHE_THRESHOLD") .ok() .and_then(|v| v.parse().ok()) @@ -41,42 +43,215 @@ impl SemanticCache { client, collection, threshold, + collection_ready: AtomicBool::new(false), }) } + /// Ensure the Qdrant collection exists with the right vector dimensions. + /// Called lazily on first use to avoid blocking startup. + pub async fn ensure_collection(&self, vector_size: u64) -> Result<(), String> { + if self.collection_ready.load(Ordering::Acquire) { + return Ok(()); + } + + // Check if collection exists + let exists = self + .client + .collection_exists(&self.collection) + .await + .map_err(|e| format!("Qdrant collection_exists check failed: {e}"))?; + + if !exists { + // Concurrent callers may both attempt create_collection. + // Treat "already exists" as success to avoid TOCTOU race. + match self + .client + .create_collection( + CreateCollectionBuilder::new(&self.collection) + .vectors_config(VectorParamsBuilder::new(vector_size, Distance::Cosine)), + ) + .await + { + Ok(_) => { + tracing::info!( + collection = %self.collection, + vector_size = vector_size, + "created Qdrant collection for semantic cache" + ); + } + Err(e) => { + // If another caller already created it, that's fine. + let msg = e.to_string(); + if !msg.contains("already exists") { + return Err(format!("Qdrant create_collection failed: {e}")); + } + } + } + } + + self.collection_ready.store(true, Ordering::Release); + Ok(()) + } + /// Search for a semantically similar cached response. /// - /// Returns the cached response body if the top result's similarity + /// Returns the cached response body and model if the top result's similarity /// score meets or exceeds the configured threshold. - /// - /// Currently a placeholder: always returns `None`. A full implementation - /// would query the Qdrant collection with the embedding vector and - /// deserialize the payload into a `CacheEntry`. - pub async fn search(&self, _embedding: &[f32]) -> Option { - // Placeholder: actual implementation would: - // 1. Search qdrant collection with the embedding vector - // 2. If top result score >= self.threshold, deserialize payload - // 3. Otherwise return None - None + pub async fn search(&self, embedding: &[f32]) -> Option { + use qdrant_client::qdrant::with_payload_selector::SelectorOptions; + + let results = self + .client + .search_points( + SearchPointsBuilder::new(&self.collection, embedding.to_vec(), 1) + .with_payload(SelectorOptions::Enable(true)), + ) + .await + .ok()?; + + let point = results.result.first()?; + if point.score < self.threshold { + tracing::debug!( + score = point.score, + threshold = self.threshold, + "semantic cache miss (below threshold)" + ); + return None; + } + + // Extract string values from Qdrant payload (protobuf Value type). + let payload = &point.payload; + let response_body = extract_string_value(payload.get("response_body")?)?; + let model = extract_string_value(payload.get("model")?)?; + + tracing::debug!( + score = point.score, + model = %model, + "semantic cache hit" + ); + + Some(super::CacheEntry { + response_body: bytes::Bytes::from(response_body), + model, + created_at: std::time::Instant::now(), + ttl_secs: None, // Semantic cache does not use per-entry TTL + }) } - /// Store a response with its embedding vector. - /// - /// Currently a placeholder. A full implementation would upsert a point - /// into the Qdrant collection with the embedding as the vector and - /// the serialized `CacheEntry` + cache key as the payload. - pub async fn store( - &self, - _embedding: &[f32], - _entry: &super::CacheEntry, - _cache_key: &str, - ) { - // Placeholder: actual implementation would: - // 1. Upsert point into qdrant collection - // 2. Vector = embedding, payload = serialized CacheEntry + cache_key + /// Store a response with its embedding vector in Qdrant. + pub async fn store(&self, embedding: &[f32], entry: &super::CacheEntry, cache_key: &str) { + let payload = Payload::try_from(serde_json::json!({ + "response_body": String::from_utf8_lossy(&entry.response_body).to_string(), + "model": entry.model.clone(), + "cache_key": cache_key, + })); + let payload = match payload { + Ok(p) => p, + Err(e) => { + tracing::warn!(error = %e, "failed to build Qdrant payload"); + return; + } + }; + + let point = PointStruct::new( + uuid::Uuid::new_v4().to_string(), + embedding.to_vec(), + payload, + ); + + if let Err(e) = self + .client + .upsert_points(UpsertPointsBuilder::new(&self.collection, vec![point])) + .await + { + tracing::warn!(error = %e, "failed to store in semantic cache"); + } } } +/// Extract a string from a Qdrant protobuf Value. +fn extract_string_value(value: &qdrant_client::qdrant::Value) -> Option { + use qdrant_client::qdrant::value::Kind; + match &value.kind { + Some(Kind::StringValue(s)) => Some(s.clone()), + _ => None, + } +} + +/// Generate an embedding for the given text using the backend's embeddings endpoint. +/// +/// Calls `embeddings_passthrough` on the backend client with an OpenAI-format +/// embedding request. Returns `None` if the backend doesn't support embeddings +/// or if the request fails. +pub async fn embed_text( + backend: &crate::backend::BackendClient, + text: &str, + model: &str, +) -> Option> { + let body = serde_json::json!({ + "input": text, + "model": model, + }); + let bytes = serde_json::to_vec(&body).ok()?; + let (status, _, resp_body) = backend + .embeddings_passthrough(bytes::Bytes::from(bytes), "application/json") + .await + .ok()?; + + if !status.is_success() { + tracing::debug!( + status = %status, + "embedding request failed for semantic cache" + ); + return None; + } + + // OpenAI embeddings response: { "data": [{ "embedding": [...] }] } + // Deserialize into a minimal struct to avoid cloning the full JSON value. + #[derive(serde::Deserialize)] + struct EmbeddingData { + embedding: Vec, + } + #[derive(serde::Deserialize)] + struct EmbeddingResponse { + data: Vec, + } + let resp: EmbeddingResponse = serde_json::from_slice(&resp_body).ok()?; + resp.data.into_iter().next().map(|d| d.embedding) +} + +/// Extract the last user message text from an Anthropic MessageCreateRequest +/// for use as the semantic cache key. +pub fn extract_last_user_text( + request: &anyllm_translate::anthropic::MessageCreateRequest, +) -> Option { + use anyllm_translate::anthropic::{Content, ContentBlock, Role}; + + for msg in request.messages.iter().rev() { + if msg.role == Role::User { + match &msg.content { + Content::Text(text) => { + if !text.is_empty() { + return Some(text.clone()); + } + } + Content::Blocks(blocks) => { + let mut text_parts = Vec::new(); + for block in blocks { + if let ContentBlock::Text { text } = block { + text_parts.push(text.as_str()); + } + } + if !text_parts.is_empty() { + return Some(text_parts.join(" ")); + } + } + } + } + } + None +} + #[cfg(test)] mod tests { use super::*; @@ -84,9 +259,7 @@ mod tests { #[test] fn new_returns_none_without_env() { // Ensure QDRANT_URL is not set for this test. - // (It should not be set in CI or local dev by default.) if std::env::var("QDRANT_URL").is_ok() { - // Skip: can't unset env vars safely in parallel tests. return; } assert!( @@ -94,4 +267,59 @@ mod tests { "SemanticCache::new() should return None when QDRANT_URL is unset" ); } + + #[test] + fn extract_last_user_text_finds_text() { + let j = serde_json::json!({ + "model": "test", + "max_tokens": 100, + "messages": [ + {"role": "user", "content": "first message"}, + {"role": "assistant", "content": "reply"}, + {"role": "user", "content": "second message"} + ] + }); + let request: anyllm_translate::anthropic::MessageCreateRequest = + serde_json::from_value(j).unwrap(); + assert_eq!( + extract_last_user_text(&request), + Some("second message".to_string()) + ); + } + + #[test] + fn extract_last_user_text_empty_messages() { + let j = serde_json::json!({ + "model": "test", + "max_tokens": 100, + "messages": [] + }); + let request: anyllm_translate::anthropic::MessageCreateRequest = + serde_json::from_value(j).unwrap(); + assert_eq!(extract_last_user_text(&request), None); + } + + #[test] + fn parse_embedding_response() { + let resp = serde_json::json!({ + "data": [{ + "embedding": [0.1, 0.2, 0.3], + "index": 0, + "object": "embedding" + }], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 5, "total_tokens": 5} + }); + let embedding: Vec = serde_json::from_value( + resp.get("data") + .unwrap() + .get(0) + .unwrap() + .get("embedding") + .unwrap() + .clone(), + ) + .unwrap(); + assert_eq!(embedding, vec![0.1, 0.2, 0.3]); + } } diff --git a/crates/proxy/src/server/routes.rs b/crates/proxy/src/server/routes.rs index 6b1e1aa..277a62c 100644 --- a/crates/proxy/src/server/routes.rs +++ b/crates/proxy/src/server/routes.rs @@ -49,7 +49,12 @@ where /// Result of resolving a model name through the model router. pub(crate) enum ResolvedModel { /// Routed via model_list to a specific backend and actual model name. - Routed { backend_name: String, model: String }, + Routed { + backend_name: String, + model: String, + /// The deployment Arc for recording in-flight/latency stats. + deployment: Arc, + }, /// Model is known but all deployments are at their RPM limit. AllAtLimit, /// No model router, or model not in router. Used legacy ModelMapping. @@ -80,7 +85,8 @@ pub struct AppState { /// Optional response cache for non-streaming requests. pub cache: Option>, /// Model-level router for LiteLLM model_list configs. None for TOML/env configs. - pub model_router: Option>, + /// Wrapped in RwLock for dynamic model management via admin API. + pub model_router: Option>>, /// All backend states, for cross-backend model routing. None unless model_router is set. pub all_backends: Option>>, } @@ -101,11 +107,13 @@ impl AppState { /// Resolve a model name through the model router (if set) or fall back to ModelMapping. pub(crate) fn resolve_model(&self, model: &str) -> ResolvedModel { - if let Some(ref router) = self.model_router { + if let Some(ref router_lock) = self.model_router { + let router = router_lock.read().unwrap_or_else(|e| e.into_inner()); if let Some(routed) = router.route(model) { return ResolvedModel::Routed { backend_name: routed.backend_name.to_string(), model: routed.actual_model.to_string(), + deployment: routed.deployment.clone(), }; } if router.has_model(model) { @@ -115,18 +123,27 @@ impl AppState { ResolvedModel::Legacy(self.map_model(model)) } - /// Resolve model and return (mapped_model, effective AppState). + /// Resolve model and return (mapped_model, effective AppState, optional deployment). /// If the model routes to a different backend, the returned state is cloned from /// all_backends. Returns Err with a 429 response if all deployments are at limit. + /// The deployment Arc is returned so handlers can call record_start/record_finish. #[allow(clippy::result_large_err)] pub(crate) fn resolve_model_and_state( &self, model: &str, - ) -> Result<(String, AppState), Response> { + ) -> Result< + ( + String, + AppState, + Option>, + ), + Response, + > { match self.resolve_model(model) { ResolvedModel::Routed { backend_name, model: mapped, + deployment, } => { let effective = self .all_backends @@ -134,7 +151,7 @@ impl AppState { .and_then(|m| m.get(&backend_name)) .cloned() .unwrap_or_else(|| self.clone()); - Ok((mapped, effective)) + Ok((mapped, effective, Some(deployment))) } ResolvedModel::AllAtLimit => { let err = mapping::errors_map::create_anthropic_error( @@ -144,7 +161,7 @@ impl AppState { ); Err((StatusCode::TOO_MANY_REQUESTS, Json(err)).into_response()) } - ResolvedModel::Legacy(mapped) => Ok((mapped, self.clone())), + ResolvedModel::Legacy(mapped) => Ok((mapped, self.clone(), None)), } } @@ -179,7 +196,7 @@ pub fn app_multi(config: MultiConfig) -> Router { pub fn app_multi_with_shared( config: MultiConfig, shared: Option, - model_router: Option>, + model_router: Option>>, ) -> Router { let mut backend_metrics: HashMap = HashMap::new(); let mut router = Router::new(); @@ -298,13 +315,22 @@ pub fn app_multi_with_shared( .layer(axum::middleware::from_fn(super::middleware::validate_auth)); // Health is public (no auth required). - Router::new() + let mut final_router = Router::new() .route("/health", get(health)) .merge(metrics_route) .merge(router) .fallback(fallback_not_found) - .layer(axum::middleware::from_fn(super::middleware::add_request_id)) - .with_state(global_state) + .layer(axum::middleware::from_fn(super::middleware::add_request_id)); + + // Apply IP allowlist middleware before auth if IP_ALLOWLIST is configured. + if super::middleware::ip_allowlist_active() { + final_router = final_router.layer(axum::middleware::from_fn( + super::middleware::check_ip_allowlist, + )); + tracing::info!("IP allowlist middleware enabled"); + } + + final_router.with_state(global_state) } /// Return Anthropic-shaped 404 for any unmatched route (PRD US-004). @@ -415,35 +441,58 @@ pub(crate) struct ConcurrencyPermit( #[allow(dead_code)] pub(crate) Arc, ); -static MODELS_RESPONSE: std::sync::LazyLock = std::sync::LazyLock::new(|| { - serde_json::json!({ - "data": [ +/// Static Claude model entries, merged with model_list models at runtime. +static STATIC_CLAUDE_MODELS: std::sync::LazyLock> = std::sync::LazyLock::new( + || { + vec![ // Claude 4.x - {"id": "claude-opus-4-6", "display_name": "Claude Opus 4.6", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, - {"id": "claude-sonnet-4-6", "display_name": "Claude Sonnet 4.6", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, - {"id": "claude-opus-4-5", "display_name": "Claude Opus 4.5", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, - {"id": "claude-sonnet-4-5", "display_name": "Claude Sonnet 4.5", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, - {"id": "claude-haiku-4-5", "display_name": "Claude Haiku 4.5", "created_at": "2025-05-14T00:00:00Z", "type": "model"}, - {"id": "claude-haiku-4-5-20251001", "display_name": "Claude Haiku 4.5 (Oct 2025)","created_at": "2025-10-01T00:00:00Z", "type": "model"}, + serde_json::json!({"id": "claude-opus-4-6", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Opus 4.6"}), + serde_json::json!({"id": "claude-sonnet-4-6", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Sonnet 4.6"}), + serde_json::json!({"id": "claude-opus-4-5", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Opus 4.5"}), + serde_json::json!({"id": "claude-sonnet-4-5", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Sonnet 4.5"}), + serde_json::json!({"id": "claude-haiku-4-5", "object": "model", "created": 1715644800, "owned_by": "anthropic", "display_name": "Claude Haiku 4.5"}), + serde_json::json!({"id": "claude-haiku-4-5-20251001", "object": "model", "created": 1727740800, "owned_by": "anthropic", "display_name": "Claude Haiku 4.5 (Oct 2025)"}), // Claude 3.7 - {"id": "claude-3-7-sonnet-20250219", "display_name": "Claude 3.7 Sonnet", "created_at": "2025-02-19T00:00:00Z", "type": "model"}, + serde_json::json!({"id": "claude-3-7-sonnet-20250219", "object": "model", "created": 1708300800, "owned_by": "anthropic", "display_name": "Claude 3.7 Sonnet"}), // Claude 3.5 - {"id": "claude-3-5-sonnet-20241022", "display_name": "Claude 3.5 Sonnet (Oct 2024)","created_at": "2024-10-22T00:00:00Z", "type": "model"}, - {"id": "claude-3-5-sonnet-20240620", "display_name": "Claude 3.5 Sonnet (Jun 2024)","created_at": "2024-06-20T00:00:00Z", "type": "model"}, - {"id": "claude-3-5-haiku-20241022", "display_name": "Claude 3.5 Haiku", "created_at": "2024-10-22T00:00:00Z", "type": "model"}, + serde_json::json!({"id": "claude-3-5-sonnet-20241022", "object": "model", "created": 1729555200, "owned_by": "anthropic", "display_name": "Claude 3.5 Sonnet (Oct 2024)"}), + serde_json::json!({"id": "claude-3-5-sonnet-20240620", "object": "model", "created": 1718841600, "owned_by": "anthropic", "display_name": "Claude 3.5 Sonnet (Jun 2024)"}), + serde_json::json!({"id": "claude-3-5-haiku-20241022", "object": "model", "created": 1729555200, "owned_by": "anthropic", "display_name": "Claude 3.5 Haiku"}), // Claude 3 - {"id": "claude-3-opus-20240229", "display_name": "Claude 3 Opus", "created_at": "2024-02-29T00:00:00Z", "type": "model"}, - {"id": "claude-3-sonnet-20240229", "display_name": "Claude 3 Sonnet", "created_at": "2024-02-29T00:00:00Z", "type": "model"}, - {"id": "claude-3-haiku-20240307", "display_name": "Claude 3 Haiku", "created_at": "2024-03-07T00:00:00Z", "type": "model"}, - ], - "has_more": false, - "first_id": "claude-opus-4-6", - "last_id": "claude-3-haiku-20240307", - }) -}); + serde_json::json!({"id": "claude-3-opus-20240229", "object": "model", "created": 1709164800, "owned_by": "anthropic", "display_name": "Claude 3 Opus"}), + serde_json::json!({"id": "claude-3-sonnet-20240229", "object": "model", "created": 1709164800, "owned_by": "anthropic", "display_name": "Claude 3 Sonnet"}), + serde_json::json!({"id": "claude-3-haiku-20240307", "object": "model", "created": 1709769600, "owned_by": "anthropic", "display_name": "Claude 3 Haiku"}), + ] + }, +); -async fn models(State(_state): State) -> Json { - Json(MODELS_RESPONSE.clone()) +/// GET /v1/models -- returns static Claude models merged with model_list entries. +async fn models(State(state): State) -> Json { + let mut data: Vec = STATIC_CLAUDE_MODELS.clone(); + + // Merge models from the model router (LiteLLM model_list config). + if let Some(ref router_lock) = state.model_router { + let router = router_lock.read().unwrap_or_else(|e| e.into_inner()); + let static_ids: std::collections::HashSet = data + .iter() + .filter_map(|m| m["id"].as_str().map(|s| s.to_string())) + .collect(); + for model_name in router.known_models() { + if !static_ids.contains(model_name) { + data.push(serde_json::json!({ + "id": model_name, + "object": "model", + "created": 0, + "owned_by": "organization" + })); + } + } + } + + Json(serde_json::json!({ + "object": "list", + "data": data, + })) } async fn batches_legacy_stub() -> impl IntoResponse { @@ -507,6 +556,7 @@ pub(crate) async fn try_cache_response( response_body: resp_body, model, created_at: std::time::Instant::now(), + ttl_secs: cache_ttl, }, ttl, ) @@ -619,13 +669,22 @@ async fn messages( if state.log_bodies() { tracing::debug!(model = %body.model, "streaming request initiated"); } - let (mapped_model, effective) = match state.resolve_model_and_state(&body.model) { + let (mapped_model, effective, deployment) = match state.resolve_model_and_state(&body.model) + { Ok(v) => v, Err(resp) => return resp, }; + if let Some(ref d) = deployment { + d.record_start(); + } // Logging deferred until stream completes (inside messages_stream tasks). + let stream_start = std::time::Instant::now(); match messages_stream(effective, body, ctx, mapped_model, permit).await { Ok((rate_limits, sse)) => { + // For streaming, record_finish is approximate (headers sent, not stream end). + if let Some(ref d) = deployment { + d.record_finish(stream_start.elapsed().as_millis() as u64); + } let mut response = sse.into_response(); rate_limits.inject_anthropic_response_headers(response.headers_mut()); inject_degradation_header(response.headers_mut(), &warnings); @@ -636,6 +695,9 @@ async fn messages( return response; } Err(e) => { + if let Some(ref d) = deployment { + d.record_finish(stream_start.elapsed().as_millis() as u64); + } // Pre-stream backend error: return proper HTTP status instead of 200 OK return backend_error_to_response(e); } @@ -681,10 +743,14 @@ async fn messages( } // Resolve model routing (may switch to a different backend). - let (mapped_model, effective) = match state.resolve_model_and_state(&body.model) { + let (mapped_model, effective, deployment) = match state.resolve_model_and_state(&body.model) { Ok(v) => v, Err(resp) => return resp, }; + if let Some(ref d) = deployment { + d.record_start(); + } + let backend_start = std::time::Instant::now(); match &effective.backend { BackendClient::OpenAI(client) @@ -702,6 +768,9 @@ async fn messages( match client.chat_completion(&openai_req).await { Ok((openai_resp, _status, rate_limits)) => { + if let Some(ref d) = deployment { + d.record_finish(backend_start.elapsed().as_millis() as u64); + } state.metrics.record_success(); let anthropic_resp = mapping::message_map::openai_to_anthropic_response( &openai_resp, @@ -746,6 +815,9 @@ async fn messages( response } Err(e) => { + if let Some(ref d) = deployment { + d.record_finish(backend_start.elapsed().as_millis() as u64); + } state.metrics.record_error(); let status = e.status_code(); log_request( @@ -772,6 +844,9 @@ async fn messages( match client.responses(&responses_req).await { Ok((resp, _status, rate_limits)) => { + if let Some(ref d) = deployment { + d.record_finish(backend_start.elapsed().as_millis() as u64); + } state.metrics.record_success(); let anthropic_resp = mapping::responses_message_map::responses_to_anthropic_response( @@ -816,6 +891,9 @@ async fn messages( response } Err(e) => { + if let Some(ref d) = deployment { + d.record_finish(backend_start.elapsed().as_millis() as u64); + } state.metrics.record_error(); let status = e.status_code(); log_request( @@ -915,8 +993,21 @@ pub(crate) fn record_vk_tpm( } } -/// Log a completed request to the admin write buffer and broadcast to WebSocket clients. +/// Global webhook callback config, set once at startup. +static CALLBACKS: std::sync::OnceLock> = + std::sync::OnceLock::new(); + +/// Set the global webhook callback config (called once at startup). +pub fn set_callbacks(config: Arc) { + let _ = CALLBACKS.set(config); +} + +/// Log a completed request to the admin write buffer, broadcast to WebSocket clients, +/// and fire webhook callbacks if configured. pub(crate) fn log_request(shared: &Option, entry: RequestLogEntry) { + if let Some(cb) = CALLBACKS.get() { + cb.notify(&entry); + } if let Some(ref shared) = shared { let _ = shared .events_tx