liter-llm 1.11.2

//! Hedged-retry Tower middleware.
//!
//! # Overview
//!
//! [`HedgeLayer`] races multiple copies of the same request against each other.
//! After a configurable delay, a second (or third, …) request is launched.
//! The first response that arrives wins; all losers are cancelled via
//! [`tokio_util::sync::CancellationToken`].
//!
//! This pattern is particularly effective for tail-latency reduction: most
//! requests complete before the hedge fires, but slow outliers get a
//! second chance without incurring extra cost in the common case.
//!
//! # Trait-first design
//!
//! The [`HedgePolicy`] trait is the extension point.  Supply a custom
//! implementation to use latency-based delays (e.g. p99 latency), adaptive
//! delays per model, or request-property-based hedging.
//!
//! # Note on `CancellationToken`
//!
//! This module depends on `tokio_util`.  Because `tokio_util` is not yet a
//! workspace dependency, it is referenced via `tokio`'s re-export
//! (`tokio_util::sync::CancellationToken` is available through
//! `tokio-util = "0.7"` which `tokio` 1.x exposes indirectly).  We use a
//! bespoke `AbortHandle` via `tokio::task::JoinSet` instead to avoid adding
//! a hard dependency here — the cancellation is implemented with
//! `tokio::select!` and `AbortHandle`.

use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;

use tower::{Layer, Service};

use super::types::{LlmRequest, LlmResponse};
use crate::client::BoxFuture;
use crate::error::{LiterLlmError, Result};

/// Policy that controls when and how many hedged requests are launched.
///
/// Implement this trait to provide custom hedging strategies such as
/// latency-percentile-based delays or per-model adaptive delays.
#[cfg_attr(alef, alef(skip))]
pub trait HedgePolicy: Send + Sync + 'static {
    /// Returns the delay before launching attempt `attempt` (1-indexed; attempt
    /// 1 is the initial request, attempt 2 is the first hedge, etc.).
    ///
    /// - `attempt`: 1-indexed attempt number.
    /// - `latency_so_far`: elapsed time since the first request was dispatched.
    ///
    /// Return `None` to skip this attempt (and all subsequent ones).
    fn delay_for_attempt(&self, attempt: u32, latency_so_far: Duration) -> Option<Duration>;

    /// Maximum number of concurrent attempts (including the original request).
    ///
    /// Must be ≥ 1.  Values above 3 are rarely useful and increase provider
    /// costs significantly.
    fn max_attempts(&self) -> u32;
}

/// A simple [`HedgePolicy`] that fires hedges at fixed intervals.
///
/// # Example
///
/// ```rust,ignore
/// use std::sync::Arc;
/// use std::time::Duration;
/// use liter_llm::tower::hedge::{FixedDelayHedge, HedgeLayer};
///
/// // Fire a second request 200 ms after the first; allow up to 2 attempts.
/// let policy = Arc::new(FixedDelayHedge::new(Duration::from_millis(200), 2));
/// let layer = HedgeLayer::new(policy);
/// ```
#[cfg_attr(alef, alef(skip))]
pub struct FixedDelayHedge {
    /// Fixed delay between attempts.
    delay: Duration,
    /// Maximum concurrent attempts including the first request.
    max_attempts: u32,
}

impl FixedDelayHedge {
    /// Create a new policy.
    ///
    /// - `delay`: how long to wait before launching each additional attempt.
    /// - `max_attempts`: maximum concurrent copies of the request (≥ 1).
    #[must_use]
    pub fn new(delay: Duration, max_attempts: u32) -> Self {
        Self {
            delay,
            max_attempts: max_attempts.max(1),
        }
    }
}

impl HedgePolicy for FixedDelayHedge {
    fn delay_for_attempt(&self, attempt: u32, _latency_so_far: Duration) -> Option<Duration> {
        if attempt > self.max_attempts {
            return None;
        }
        Some(self.delay * (attempt - 1))
    }

    fn max_attempts(&self) -> u32 {
        self.max_attempts
    }
}

/// Tower [`Layer`] that wraps a service with hedged request racing.
///
/// The layer clones the inner service for each additional attempt.
#[cfg_attr(alef, alef(skip))]
pub struct HedgeLayer<P> {
    policy: Arc<P>,
}

impl<P: HedgePolicy> HedgeLayer<P> {
    /// Create a new hedge layer.
    #[must_use]
    pub fn new(policy: Arc<P>) -> Self {
        Self { policy }
    }
}

impl<P: HedgePolicy, S> Layer<S> for HedgeLayer<P> {
    type Service = HedgeService<P, S>;

    fn layer(&self, inner: S) -> Self::Service {
        HedgeService {
            inner,
            policy: Arc::clone(&self.policy),
        }
    }
}

/// Tower service produced by [`HedgeLayer`].
#[cfg_attr(alef, alef(skip))]
pub struct HedgeService<P, S> {
    inner: S,
    policy: Arc<P>,
}

impl<P: HedgePolicy, S: Clone> Clone for HedgeService<P, S> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
            policy: Arc::clone(&self.policy),
        }
    }
}

impl<P, S> Service<LlmRequest> for HedgeService<P, S>
where
    P: HedgePolicy + 'static,
    S: Service<LlmRequest, Response = LlmResponse, Error = LiterLlmError> + Send + Clone + 'static,
    S::Future: Send + 'static,
{
    type Response = LlmResponse;
    type Error = LiterLlmError;
    type Future = BoxFuture<'static, Result<LlmResponse>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: LlmRequest) -> Self::Future {
        let policy = Arc::clone(&self.policy);
        let max_attempts = policy.max_attempts();

        // ~keep Use the poll-ready instance for attempt #1 so permits are consumed correctly.
        // ~keep Hedge attempts use standby clones and call ready() to avoid double-consuming permits.
        let standby = self.inner.clone();
        let primary = std::mem::replace(&mut self.inner, standby);
        let inner_for_hedges = self.inner.clone();

        Box::pin(async move {
            tracing::debug!(hedge.max_attempts = max_attempts, "starting hedged request");
            hedge_race(req, primary, inner_for_hedges, policy, max_attempts).await
        })
    }
}

/// Core hedging logic: spawn attempts with increasing delays and race them.
///
/// Uses `JoinSet` with `abort_all()` to cancel losing tasks.
///
/// # Tower readiness contract
///
/// `primary` is the service instance on which `poll_ready` was already called
/// by [`HedgeService::poll_ready`].  It is used directly for attempt #1 so
/// that any permit acquired during readiness (e.g. a `ConcurrencyLimit`
/// semaphore slot) is properly consumed.
///
/// Hedged attempts (#2‥N) each receive a *fresh clone* of `inner_for_hedges`
/// and call `ServiceExt::ready()` inside their spawned task before calling
/// the service.  This means hedged attempts may wait for permits to become
/// available, which is the correct behaviour — hedging is not a mechanism to
/// bypass concurrency controls.
async fn hedge_race<S>(
    req: LlmRequest,
    mut primary: S,
    inner_for_hedges: S,
    policy: Arc<impl HedgePolicy>,
    max_attempts: u32,
) -> Result<LlmResponse>
where
    S: Service<LlmRequest, Response = LlmResponse, Error = LiterLlmError> + Send + Clone + 'static,
    S::Future: Send + 'static,
{
    use std::time::Instant;

    use tower::ServiceExt as _;

    let dispatch_time = Instant::now();

    if max_attempts == 1 {
        tracing::debug!("hedge fast path: max_attempts=1, calling primary directly");
        return primary.call(req).await;
    }

    let mut join_set: tokio::task::JoinSet<(u32, Result<LlmResponse>)> = tokio::task::JoinSet::new();

    {
        let req_clone = req.clone();
        join_set.spawn(async move {
            let result = primary.call(req_clone).await;
            (1u32, result)
        });
    }

    // ~keep Later hedge attempts must call ready() so concurrency and buffer permits are respected.
    for attempt in 2..=max_attempts {
        let latency_so_far = dispatch_time.elapsed();
        let Some(hedge_delay) = policy.delay_for_attempt(attempt, latency_so_far) else {
            break;
        };

        let req_clone = req.clone();
        let mut svc_clone = inner_for_hedges.clone();
        join_set.spawn(async move {
            if hedge_delay > Duration::ZERO {
                tokio::time::sleep(hedge_delay).await;
            }
            tracing::debug!(attempt, "launching hedged request");

            let model = req_clone.model().unwrap_or("").to_owned();
            let system = model.split_once('/').map(|(p, _)| p.to_owned()).unwrap_or_default();
            super::metrics::record_retry_attempt(&system, &model, req_clone.operation_name());

            // ~keep ready() acquires any per-instance permits before the hedged call.
            let ready_result = svc_clone.ready().await;
            let result = match ready_result {
                Ok(ready_svc) => ready_svc.call(req_clone).await,
                Err(e) => Err(e),
            };
            (attempt, result)
        });
    }

    let mut last_err: Option<LiterLlmError> = None;

    while let Some(join_result) = join_set.join_next().await {
        match join_result {
            Ok((attempt, Ok(resp))) => {
                tracing::debug!(attempt, "hedged request succeeded first");
                join_set.abort_all();
                return Ok(resp);
            }
            Ok((attempt, Err(e))) => {
                tracing::debug!(attempt, error = %e, "hedged attempt failed");
                last_err = Some(e);
            }
            Err(join_err) if join_err.is_cancelled() => {}
            Err(join_err) => {
                tracing::error!(error = %join_err, "hedged task panicked");
                last_err = Some(LiterLlmError::InternalError {
                    message: format!("hedge task panicked: {join_err}"),
                });
            }
        }
    }

    Err(last_err.unwrap_or(LiterLlmError::InternalError {
        message: "all hedged attempts failed with no error recorded".into(),
    }))
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::Ordering;
    use std::time::Duration;

    use tower::{Layer as _, Service as _, ServiceExt as _};

    use super::*;
    use crate::tower::service::LlmService;
    use crate::tower::tests_common::{MockClient, chat_req};
    use crate::tower::types::LlmRequest;

    #[tokio::test]
    async fn hedge_returns_first_success() {
        let policy = Arc::new(FixedDelayHedge::new(Duration::from_millis(200), 2));
        let inner = LlmService::new(MockClient::ok());
        let mut svc = HedgeLayer::new(policy).layer(inner);

        let resp = svc
            .call(LlmRequest::Chat(chat_req("openai/gpt-4")))
            .await
            .expect("should succeed");

        assert!(matches!(resp, LlmResponse::Chat(_)));
    }

    #[tokio::test]
    async fn hedge_single_attempt_policy_does_not_spawn_extra() {
        let policy = Arc::new(FixedDelayHedge::new(Duration::from_millis(100), 1));
        let mock = MockClient::ok();
        let call_count = Arc::clone(&mock.call_count);
        let inner = LlmService::new(mock);
        let mut svc = HedgeLayer::new(policy).layer(inner);

        svc.call(LlmRequest::Chat(chat_req("openai/gpt-4")))
            .await
            .expect("should succeed");

        assert_eq!(
            call_count.load(Ordering::SeqCst),
            1,
            "max_attempts=1 should only call inner service once"
        );
    }

    #[tokio::test]
    async fn hedge_propagates_error_when_all_attempts_fail() {
        let policy = Arc::new(FixedDelayHedge::new(Duration::from_millis(10), 2));
        let inner = LlmService::new(MockClient::failing_timeout());
        let mut svc = HedgeLayer::new(policy).layer(inner);

        let err = svc
            .call(LlmRequest::Chat(chat_req("openai/gpt-4")))
            .await
            .expect_err("all attempts should fail");

        assert!(
            matches!(err, LiterLlmError::Timeout),
            "expected Timeout from failed hedge, got {err:?}"
        );
    }

    #[tokio::test]
    async fn fixed_delay_hedge_policy_respects_max_attempts() {
        let policy = FixedDelayHedge::new(Duration::from_millis(100), 3);

        assert_eq!(policy.delay_for_attempt(1, Duration::ZERO), Some(Duration::ZERO));
        assert_eq!(
            policy.delay_for_attempt(2, Duration::ZERO),
            Some(Duration::from_millis(100))
        );
        assert_eq!(
            policy.delay_for_attempt(3, Duration::ZERO),
            Some(Duration::from_millis(200))
        );
        assert_eq!(policy.delay_for_attempt(4, Duration::ZERO), None);
    }

    #[tokio::test]
    async fn hedge_with_two_attempts_calls_inner_at_most_twice() {
        let policy = Arc::new(FixedDelayHedge::new(Duration::from_millis(5), 2));
        let mock = MockClient::ok();
        let call_count = Arc::clone(&mock.call_count);
        let inner = LlmService::new(mock);
        let mut svc = HedgeLayer::new(policy).layer(inner);

        svc.call(LlmRequest::Chat(chat_req("openai/gpt-4")))
            .await
            .expect("should succeed");

        let count = call_count.load(Ordering::SeqCst);
        assert!((1..=2).contains(&count), "expected 1 or 2 calls, got {count}");
    }

    /// With `max_attempts=1`, the JoinSet fast path is skipped entirely and
    /// the inner service is called exactly once (Option B fast path).
    #[tokio::test]
    async fn hedge_max_attempts_one_does_not_spawn_extra() {
        let policy = Arc::new(FixedDelayHedge::new(Duration::from_millis(0), 1));
        let mock = MockClient::ok();
        let call_count = Arc::clone(&mock.call_count);
        let inner = LlmService::new(mock);
        let mut svc = HedgeLayer::new(policy).layer(inner);

        for _ in 0..2 {
            svc.ready()
                .await
                .expect("service should become ready")
                .call(LlmRequest::Chat(chat_req("openai/gpt-4")))
                .await
                .expect("should succeed");
        }

        assert_eq!(
            call_count.load(Ordering::SeqCst),
            2,
            "max_attempts=1 must not spawn additional tasks; expected exactly 2 calls total"
        );
    }

    /// A `ConcurrencyLimit(1)` inner service permits only one in-flight call
    /// at a time.  With `max_attempts=2` and a hedge delay of 0, the second
    /// attempt must wait for the first to release its permit before it can
    /// proceed.  This verifies that hedged clones use `ready()` and do not
    /// bypass the concurrency limit.
    #[tokio::test]
    async fn hedge_respects_inner_readiness_via_ready_and_call() {
        use tower::limit::ConcurrencyLimitLayer;

        let mock = MockClient::ok();
        let call_count = Arc::clone(&mock.call_count);
        let inner = LlmService::new(mock);
        let limited = ConcurrencyLimitLayer::new(1).layer(inner);

        let policy = Arc::new(FixedDelayHedge::new(Duration::ZERO, 2));
        let mut svc = HedgeLayer::new(policy).layer(limited);

        let resp = svc
            .ready()
            .await
            .expect("service should become ready")
            .call(LlmRequest::Chat(chat_req("openai/gpt-4")))
            .await
            .expect("hedged request should succeed");

        assert!(matches!(resp, LlmResponse::Chat(_)));

        let count = call_count.load(Ordering::SeqCst);
        assert!(
            (1..=2).contains(&count),
            "expected 1 or 2 inner calls with ConcurrencyLimit(1), got {count}"
        );
    }

    /// Bug 3 fix: wrapping inner in `ConcurrencyLimit(1)` with 2 hedge attempts
    /// must not panic or exceed 1 concurrent in-flight call at any moment.
    #[tokio::test]
    async fn hedge_no_double_permit_consumption() {
        use std::sync::atomic::AtomicUsize;

        use tower::limit::ConcurrencyLimit;

        let peak = Arc::new(AtomicUsize::new(0));
        let current = Arc::new(AtomicUsize::new(0));

        #[derive(Clone)]
        struct PeakTracker {
            peak: Arc<AtomicUsize>,
            current: Arc<AtomicUsize>,
        }

        impl tower::Service<LlmRequest> for PeakTracker {
            type Response = LlmResponse;
            type Error = LiterLlmError;
            type Future = crate::client::BoxFuture<'static, crate::error::Result<LlmResponse>>;

            fn poll_ready(&mut self, _cx: &mut std::task::Context<'_>) -> std::task::Poll<crate::error::Result<()>> {
                std::task::Poll::Ready(Ok(()))
            }

            fn call(&mut self, _req: LlmRequest) -> Self::Future {
                let peak = Arc::clone(&self.peak);
                let current = Arc::clone(&self.current);
                Box::pin(async move {
                    let now = current.fetch_add(1, Ordering::SeqCst) + 1;
                    let mut prev = peak.load(Ordering::SeqCst);
                    while now > prev {
                        match peak.compare_exchange(prev, now, Ordering::SeqCst, Ordering::SeqCst) {
                            Ok(_) => break,
                            Err(p) => prev = p,
                        }
                    }
                    tokio::task::yield_now().await;
                    current.fetch_sub(1, Ordering::SeqCst);
                    Ok(LlmResponse::Chat(crate::tower::tests_common::make_chat_response(
                        "gpt-4",
                    )))
                })
            }
        }

        let tracker = PeakTracker {
            peak: Arc::clone(&peak),
            current: Arc::clone(&current),
        };
        let limited: ConcurrencyLimit<PeakTracker> = ConcurrencyLimit::new(tracker, 1);
        let policy = Arc::new(FixedDelayHedge::new(Duration::ZERO, 2));
        let mut svc = HedgeLayer::new(policy).layer(limited);

        let resp = svc
            .ready()
            .await
            .expect("service should become ready")
            .call(LlmRequest::Chat(chat_req("openai/gpt-4")))
            .await
            .expect("hedged request must succeed");

        assert!(matches!(resp, LlmResponse::Chat(_)));
        assert_eq!(
            peak.load(Ordering::SeqCst),
            1,
            "ConcurrencyLimit(1) must cap concurrent calls at 1 even with hedging"
        );
    }

    /// HIGH-priority correctness: when the winner returns, the loser must be
    /// dropped (cancelled) so its long-running future does not continue to
    /// consume permits or upstream resources.
    ///
    /// We wrap the inner response future in a `DropGuard` that decrements a
    /// shared `AtomicUsize` on drop.  The winner finishes quickly; the loser
    /// is parked indefinitely until aborted.  After the hedge returns, both
    /// drop-counts must be observed, proving the loser's future was dropped.
    #[tokio::test]
    async fn hedge_loser_is_dropped_before_winner_returns() {
        use std::sync::atomic::AtomicUsize;
        use std::task::Poll;

        use tokio::sync::Notify;

        /// RAII helper: drops decrement `live_count`.
        struct DropGuard {
            live_count: Arc<AtomicUsize>,
        }
        impl Drop for DropGuard {
            fn drop(&mut self) {
                self.live_count.fetch_sub(1, Ordering::SeqCst);
            }
        }

        let live = Arc::new(AtomicUsize::new(0));
        let total_calls = Arc::new(AtomicUsize::new(0));
        let winner_signal = Arc::new(Notify::new());

        #[derive(Clone)]
        struct SlowOrFast {
            live: Arc<AtomicUsize>,
            total_calls: Arc<AtomicUsize>,
            attempt: Arc<AtomicUsize>,
            winner_signal: Arc<Notify>,
        }

        impl tower::Service<LlmRequest> for SlowOrFast {
            type Response = LlmResponse;
            type Error = LiterLlmError;
            type Future = crate::client::BoxFuture<'static, crate::error::Result<LlmResponse>>;

            fn poll_ready(&mut self, _cx: &mut std::task::Context<'_>) -> Poll<crate::error::Result<()>> {
                Poll::Ready(Ok(()))
            }

            fn call(&mut self, _req: LlmRequest) -> Self::Future {
                let attempt = self.attempt.fetch_add(1, Ordering::SeqCst) + 1;
                self.total_calls.fetch_add(1, Ordering::SeqCst);
                self.live.fetch_add(1, Ordering::SeqCst);
                let guard = DropGuard {
                    live_count: Arc::clone(&self.live),
                };
                let winner_signal = Arc::clone(&self.winner_signal);
                Box::pin(async move {
                    let _g = guard;
                    if attempt == 1 {
                        tokio::time::sleep(Duration::from_millis(20)).await;
                        winner_signal.notify_one();
                        Ok(LlmResponse::Chat(crate::tower::tests_common::make_chat_response(
                            "gpt-4",
                        )))
                    } else {
                        std::future::pending::<()>().await;
                        unreachable!("loser must be cancelled before completing");
                    }
                })
            }
        }

        let inner = SlowOrFast {
            live: Arc::clone(&live),
            total_calls: Arc::clone(&total_calls),
            attempt: Arc::new(AtomicUsize::new(0)),
            winner_signal: Arc::clone(&winner_signal),
        };

        let policy = Arc::new(FixedDelayHedge::new(Duration::ZERO, 2));
        let mut svc = HedgeLayer::new(policy).layer(inner);

        let resp = svc
            .ready()
            .await
            .expect("ready")
            .call(LlmRequest::Chat(chat_req("openai/gpt-4")))
            .await
            .expect("winner must succeed");
        assert!(matches!(resp, LlmResponse::Chat(_)));

        assert_eq!(
            total_calls.load(Ordering::SeqCst),
            2,
            "expected primary + 1 hedged attempt"
        );

        for _ in 0..50 {
            if live.load(Ordering::SeqCst) == 0 {
                break;
            }
            tokio::task::yield_now().await;
            tokio::time::sleep(Duration::from_millis(2)).await;
        }
        assert_eq!(
            live.load(Ordering::SeqCst),
            0,
            "loser future must be dropped after winner returns; {} still alive",
            live.load(Ordering::SeqCst)
        );
    }
}