use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use dashmap::DashMap;
use tower::{Layer, Service};
use crate::client::BoxFuture;
use crate::error::LiterLlmError;
use crate::error::Result as LiterResult;
use crate::tower::cache::CachedResponse;
use crate::tower::types::{LlmRequest, LlmRequestKind, LlmResponse};
const IDEM_HASH_SEED_0: u64 = 0x6964_656d_706f_7465;
const IDEM_HASH_SEED_1: u64 = 0x6e63_795f_6861_7368;
const IDEM_HASH_SEED_2: u64 = 0x5f73_6565_6430_5f76;
const IDEM_HASH_SEED_3: u64 = 0x315f_6c6c_6d00_0000;
fn idem_random_state() -> &'static ahash::RandomState {
use std::sync::OnceLock;
static STATE: OnceLock<ahash::RandomState> = OnceLock::new();
STATE.get_or_init(|| {
ahash::RandomState::generate_with(IDEM_HASH_SEED_0, IDEM_HASH_SEED_1, IDEM_HASH_SEED_2, IDEM_HASH_SEED_3)
})
}
fn compute_body_hash(request: &LlmRequest) -> Option<String> {
let json = serde_json::to_string(&request.kind).ok()?;
let h = idem_random_state().hash_one(&json);
Some(format!("{h:016x}:{}", &json[..json.len().min(64)]))
}
#[derive(Clone)]
pub struct IdempotencyEntry {
pub body_hash: String,
pub response: Option<CachedResponse>,
pub inserted_at: Instant,
pub ttl: Duration,
}
impl std::fmt::Debug for IdempotencyEntry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("IdempotencyEntry")
.field("body_hash", &self.body_hash)
.field("has_response", &self.response.is_some())
.field("inserted_at", &self.inserted_at)
.field("ttl", &self.ttl)
.finish()
}
}
impl IdempotencyEntry {
fn is_expired(&self) -> bool {
self.inserted_at.elapsed() > self.ttl
}
}
#[derive(Debug, thiserror::Error)]
pub enum IdempotencyStoreError {
#[error("idempotency store backend error: {0}")]
Backend(String),
}
pub trait IdempotencyStore: Send + Sync + 'static {
fn get<'a>(
&'a self,
key: &'a str,
) -> Pin<Box<dyn Future<Output = Result<Option<IdempotencyEntry>, IdempotencyStoreError>> + Send + 'a>>;
fn try_insert<'a>(
&'a self,
key: &'a str,
body_hash: &'a str,
ttl: Duration,
) -> Pin<Box<dyn Future<Output = Result<bool, IdempotencyStoreError>> + Send + 'a>>;
fn store_response<'a>(
&'a self,
key: &'a str,
response: CachedResponse,
) -> Pin<Box<dyn Future<Output = Result<(), IdempotencyStoreError>> + Send + 'a>>;
fn remove<'a>(
&'a self,
key: &'a str,
) -> Pin<Box<dyn Future<Output = Result<(), IdempotencyStoreError>> + Send + 'a>>;
}
#[derive(Default)]
pub struct InMemoryIdempotencyStore {
map: DashMap<String, IdempotencyEntry>,
}
impl InMemoryIdempotencyStore {
#[must_use]
pub fn new() -> Self {
Self::default()
}
}
impl IdempotencyStore for InMemoryIdempotencyStore {
fn get<'a>(
&'a self,
key: &'a str,
) -> Pin<Box<dyn Future<Output = Result<Option<IdempotencyEntry>, IdempotencyStoreError>> + Send + 'a>> {
let result = self
.map
.get(key)
.and_then(|entry| if entry.is_expired() { None } else { Some(entry.clone()) });
Box::pin(std::future::ready(Ok(result)))
}
fn try_insert<'a>(
&'a self,
key: &'a str,
body_hash: &'a str,
ttl: Duration,
) -> Pin<Box<dyn Future<Output = Result<bool, IdempotencyStoreError>> + Send + 'a>> {
use dashmap::mapref::entry::Entry;
let inserted = match self.map.entry(key.to_owned()) {
Entry::Vacant(slot) => {
slot.insert(IdempotencyEntry {
body_hash: body_hash.to_owned(),
response: None,
inserted_at: Instant::now(),
ttl,
});
true
}
Entry::Occupied(entry) => {
if entry.get().is_expired() {
entry.replace_entry(IdempotencyEntry {
body_hash: body_hash.to_owned(),
response: None,
inserted_at: Instant::now(),
ttl,
});
true
} else {
false
}
}
};
Box::pin(std::future::ready(Ok(inserted)))
}
fn store_response<'a>(
&'a self,
key: &'a str,
response: CachedResponse,
) -> Pin<Box<dyn Future<Output = Result<(), IdempotencyStoreError>> + Send + 'a>> {
if let Some(mut entry) = self.map.get_mut(key) {
entry.response = Some(response);
}
Box::pin(std::future::ready(Ok(())))
}
fn remove<'a>(
&'a self,
key: &'a str,
) -> Pin<Box<dyn Future<Output = Result<(), IdempotencyStoreError>> + Send + 'a>> {
self.map.remove(key);
Box::pin(std::future::ready(Ok(())))
}
}
#[cfg_attr(alef, alef(skip))]
pub struct IdempotencyLayer<S: IdempotencyStore> {
store: Arc<S>,
ttl: Duration,
}
impl<S: IdempotencyStore> IdempotencyLayer<S> {
#[must_use]
pub fn new(store: S) -> Self {
Self::with_ttl(store, Duration::from_secs(24 * 60 * 60))
}
#[must_use]
pub fn with_ttl(store: S, ttl: Duration) -> Self {
Self {
store: Arc::new(store),
ttl,
}
}
}
impl<S: IdempotencyStore> Clone for IdempotencyLayer<S> {
fn clone(&self) -> Self {
Self {
store: Arc::clone(&self.store),
ttl: self.ttl,
}
}
}
impl<I, S: IdempotencyStore> Layer<I> for IdempotencyLayer<S> {
type Service = IdempotencyService<I, S>;
fn layer(&self, inner: I) -> Self::Service {
IdempotencyService {
inner,
store: Arc::clone(&self.store),
ttl: self.ttl,
}
}
}
#[cfg_attr(alef, alef(skip))]
pub struct IdempotencyService<I, S: IdempotencyStore> {
inner: I,
store: Arc<S>,
ttl: Duration,
}
impl<I: Clone, S: IdempotencyStore> Clone for IdempotencyService<I, S> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
store: Arc::clone(&self.store),
ttl: self.ttl,
}
}
}
impl<I, S> Service<LlmRequest> for IdempotencyService<I, S>
where
I: Service<LlmRequest, Response = LlmResponse, Error = LiterLlmError> + Clone + Send + 'static,
I::Future: Send + 'static,
S: IdempotencyStore,
{
type Response = LlmResponse;
type Error = LiterLlmError;
type Future = BoxFuture<'static, LiterResult<LlmResponse>>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<LiterResult<()>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, request: LlmRequest) -> Self::Future {
let standby = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, standby);
let store = Arc::clone(&self.store);
let ttl = self.ttl;
Box::pin(async move {
let Some(ref raw_key) = request.idempotency_key.clone() else {
return inner.call(request).await;
};
let tenant_prefix = request.tenant_id.as_ref().map(|t| t.as_ref()).unwrap_or("_");
let key = format!("{tenant_prefix}:{raw_key}");
let body_hash = match compute_body_hash(&request) {
Some(h) => h,
None => {
return inner.call(request).await;
}
};
if let Some(entry) = store.get(&key).await.map_err(store_err)? {
if entry.body_hash != body_hash {
return Err(LiterLlmError::IdempotencyConflict { key: raw_key.clone() });
}
if let Some(cached) = entry.response {
return cached.into_llm_response();
}
return Err(LiterLlmError::IdempotencyInFlight { key: raw_key.clone() });
}
let inserted = store.try_insert(&key, &body_hash, ttl).await.map_err(store_err)?;
if !inserted && let Some(entry) = store.get(&key).await.map_err(store_err)? {
if entry.body_hash != body_hash {
return Err(LiterLlmError::IdempotencyConflict { key: raw_key.clone() });
}
if let Some(cached) = entry.response {
return cached.into_llm_response();
}
return Err(LiterLlmError::IdempotencyInFlight { key: raw_key.clone() });
}
let result = inner.call(request).await;
match &result {
Ok(resp) => {
let cached = match resp {
LlmResponse::Chat(r) => Some(CachedResponse::Chat(r.clone())),
LlmResponse::Embed(r) => Some(CachedResponse::Embed(r.clone())),
_ => None,
};
if let Some(cached_resp) = cached {
let _ = store.store_response(&key, cached_resp).await;
} else {
let _ = store.remove(&key).await;
}
}
Err(_) => {
let _ = store.remove(&key).await;
}
}
result
})
}
}
#[inline]
fn store_err(e: IdempotencyStoreError) -> LiterLlmError {
LiterLlmError::InternalError {
message: format!("idempotency store: {e}"),
}
}
#[must_use]
#[allow(dead_code)]
pub(crate) fn is_cacheable_kind(kind: &LlmRequestKind) -> bool {
matches!(kind, LlmRequestKind::Chat(_) | LlmRequestKind::Embed(_))
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tower::{Layer as _, Service as _};
use super::*;
use crate::error::LiterLlmError;
use crate::tower::service::LlmService;
use crate::tower::tests_common::{MockClient, chat_req};
use crate::tower::types::{LlmRequest, LlmResponse};
fn make_layer() -> IdempotencyLayer<InMemoryIdempotencyStore> {
IdempotencyLayer::new(InMemoryIdempotencyStore::new())
}
fn req_with_key(model: &str, key: &str) -> LlmRequest {
LlmRequest::Chat(chat_req(model)).with_idempotency_key(key)
}
#[tokio::test]
async fn store_get_returns_none_on_miss() {
let store = InMemoryIdempotencyStore::new();
let result = store.get("missing-key").await.unwrap();
assert!(result.is_none());
}
#[tokio::test]
async fn store_try_insert_wins_first_caller() {
let store = InMemoryIdempotencyStore::new();
let inserted = store.try_insert("k1", "hash1", Duration::from_secs(60)).await.unwrap();
assert!(inserted, "first caller must win insertion");
let second = store.try_insert("k1", "hash1", Duration::from_secs(60)).await.unwrap();
assert!(!second, "second caller must lose insertion race");
}
#[tokio::test]
async fn store_try_insert_wins_after_expiry() {
let store = InMemoryIdempotencyStore::new();
store.try_insert("k2", "hash", Duration::from_nanos(1)).await.unwrap();
tokio::time::sleep(Duration::from_millis(2)).await;
let inserted = store.try_insert("k2", "hash", Duration::from_secs(60)).await.unwrap();
assert!(inserted, "insertion after TTL expiry must succeed");
}
#[tokio::test]
async fn store_get_returns_none_for_expired_entry() {
let store = InMemoryIdempotencyStore::new();
store.try_insert("k3", "hash", Duration::from_nanos(1)).await.unwrap();
tokio::time::sleep(Duration::from_millis(2)).await;
let result = store.get("k3").await.unwrap();
assert!(result.is_none(), "expired entry must not be returned");
}
#[tokio::test]
async fn store_store_response_populates_entry() {
let store = InMemoryIdempotencyStore::new();
store.try_insert("k4", "hash", Duration::from_secs(60)).await.unwrap();
let resp = CachedResponse::Chat(crate::tower::tests_common::make_chat_response("gpt-4"));
store.store_response("k4", resp).await.unwrap();
let entry = store.get("k4").await.unwrap().expect("entry must exist");
assert!(entry.response.is_some(), "response must be populated");
}
#[tokio::test]
async fn store_remove_deletes_entry() {
let store = InMemoryIdempotencyStore::new();
store.try_insert("k5", "hash", Duration::from_secs(60)).await.unwrap();
store.remove("k5").await.unwrap();
let result = store.get("k5").await.unwrap();
assert!(result.is_none(), "removed entry must not be present");
}
#[tokio::test]
async fn first_request_hits_inner() {
let layer = make_layer();
let client = MockClient::ok();
let call_count = Arc::clone(&client.call_count);
let mut svc = layer.layer(LlmService::new(client));
let result = svc.call(req_with_key("gpt-4", "key-001")).await;
assert!(result.is_ok(), "first request must succeed");
assert_eq!(call_count.load(Ordering::SeqCst), 1, "inner must be called once");
}
#[tokio::test]
async fn repeat_same_key_same_body_returns_cached() {
let layer = make_layer();
let client = MockClient::ok();
let call_count = Arc::clone(&client.call_count);
let mut svc = layer.layer(LlmService::new(client));
svc.call(req_with_key("gpt-4", "key-002"))
.await
.expect("first call must succeed");
assert_eq!(call_count.load(Ordering::SeqCst), 1);
let result = svc.call(req_with_key("gpt-4", "key-002")).await;
assert!(result.is_ok(), "second call must succeed");
assert_eq!(
call_count.load(Ordering::SeqCst),
1,
"inner must NOT be called on second request with same key+body"
);
}
#[tokio::test]
async fn repeat_same_key_different_body_returns_conflict() {
let layer = make_layer();
let client = MockClient::ok();
let mut svc = layer.layer(LlmService::new(client));
svc.call(req_with_key("gpt-4", "key-003"))
.await
.expect("first call must succeed");
let result = svc.call(req_with_key("gpt-3.5-turbo", "key-003")).await;
assert!(
matches!(result, Err(LiterLlmError::IdempotencyConflict { .. })),
"different body for same key must return IdempotencyConflict; got {result:?}"
);
}
#[tokio::test]
async fn no_key_passes_through() {
let layer = make_layer();
let client = MockClient::ok();
let call_count = Arc::clone(&client.call_count);
let mut svc = layer.layer(LlmService::new(client));
let result = svc.call(LlmRequest::Chat(chat_req("gpt-4"))).await;
assert!(result.is_ok(), "request without key must succeed");
assert_eq!(
call_count.load(Ordering::SeqCst),
1,
"inner must be called for keyless request"
);
}
#[tokio::test]
async fn inner_error_does_not_cache() {
let layer = make_layer();
let client = MockClient::failing_rate_limited();
let call_count = Arc::clone(&client.call_count);
let mut svc = layer.layer(LlmService::new(client));
let first = svc.call(req_with_key("gpt-4", "key-err")).await;
assert!(first.is_err(), "first call must fail");
assert_eq!(call_count.load(Ordering::SeqCst), 1);
let second = svc.call(req_with_key("gpt-4", "key-err")).await;
assert!(second.is_err(), "second call must also fail (same inner error)");
assert_eq!(
call_count.load(Ordering::SeqCst),
2,
"inner must be called again after first failed call"
);
}
#[tokio::test]
#[ignore = "moka time-mocking not available; TTL expiry tested via InMemoryIdempotencyStore unit tests"]
async fn ttl_expiry_allows_new_invocation() {}
#[tokio::test]
async fn different_keys_are_independent() {
let layer = make_layer();
let client = MockClient::ok();
let call_count = Arc::clone(&client.call_count);
let mut svc = layer.layer(LlmService::new(client));
svc.call(req_with_key("gpt-4", "key-A"))
.await
.expect("call A must succeed");
svc.call(req_with_key("gpt-4", "key-B"))
.await
.expect("call B must succeed");
assert_eq!(
call_count.load(Ordering::SeqCst),
2,
"different keys must both hit inner"
);
svc.call(req_with_key("gpt-4", "key-A"))
.await
.expect("repeat A must succeed");
svc.call(req_with_key("gpt-4", "key-B"))
.await
.expect("repeat B must succeed");
assert_eq!(
call_count.load(Ordering::SeqCst),
2,
"repeated calls with same key+body must not hit inner"
);
}
#[tokio::test]
async fn returned_response_matches_original() {
let layer = make_layer();
let client = MockClient::ok();
let mut svc = layer.layer(LlmService::new(client));
let first = svc
.call(req_with_key("gpt-4", "key-content"))
.await
.expect("first call");
let first_model = match &first {
LlmResponse::Chat(r) => r.model.clone(),
_ => panic!("expected Chat response"),
};
let second = svc
.call(req_with_key("gpt-4", "key-content"))
.await
.expect("second call");
let second_model = match &second {
LlmResponse::Chat(r) => r.model.clone(),
_ => panic!("expected Chat response"),
};
assert_eq!(first_model, second_model, "cached response must match original");
}
#[test]
fn idempotency_body_hash_deterministic_across_instances() {
let req = LlmRequest::Chat(chat_req("gpt-4"));
let hashes: Vec<_> = (0..10).map(|_| compute_body_hash(&req)).collect();
let first = hashes[0].as_ref().expect("hash must be Some");
for (i, h) in hashes.iter().enumerate() {
assert_eq!(
h.as_ref().expect("hash must be Some"),
first,
"hash #{i} differs from hash #0 — ahash seed is not fixed"
);
}
}
#[tokio::test]
async fn idempotency_tenant_scoped_keys_dont_collide() {
use crate::tower::types::LlmResponse;
let store = Arc::new(InMemoryIdempotencyStore::new());
let layer_a = IdempotencyLayer::new(InMemoryIdempotencyStore::new());
let layer_b = IdempotencyLayer::new(InMemoryIdempotencyStore::new());
let _ = (store, layer_a, layer_b);
let shared_store = Arc::new(InMemoryIdempotencyStore::new());
let make_layer_shared = || IdempotencyLayer {
store: Arc::clone(&shared_store),
ttl: Duration::from_secs(60),
};
let client_a = MockClient::ok();
let call_count_a = Arc::clone(&client_a.call_count);
let mut svc_a = make_layer_shared().layer(LlmService::new(client_a));
let client_b = MockClient::ok();
let call_count_b = Arc::clone(&client_b.call_count);
let mut svc_b = make_layer_shared().layer(LlmService::new(client_b));
let req_a = LlmRequest::Chat(chat_req("gpt-4"))
.with_idempotency_key("shared-key")
.with_tenant_id("tenant-A");
let req_b = LlmRequest::Chat(chat_req("gpt-4"))
.with_idempotency_key("shared-key")
.with_tenant_id("tenant-B");
let resp_a = svc_a.call(req_a.clone()).await.expect("tenant A first call");
assert!(matches!(resp_a, LlmResponse::Chat(_)));
assert_eq!(call_count_a.load(Ordering::SeqCst), 1, "inner called for tenant A");
let resp_b = svc_b.call(req_b.clone()).await.expect("tenant B first call");
assert!(matches!(resp_b, LlmResponse::Chat(_)));
assert_eq!(
call_count_b.load(Ordering::SeqCst),
1,
"inner called for tenant B (no cross-tenant hit)"
);
svc_a.call(req_a).await.expect("tenant A repeat");
assert_eq!(
call_count_a.load(Ordering::SeqCst),
1,
"inner NOT called on tenant A repeat"
);
svc_b.call(req_b).await.expect("tenant B repeat");
assert_eq!(
call_count_b.load(Ordering::SeqCst),
1,
"inner NOT called on tenant B repeat"
);
}
}