litellm-rs 0.9.0

// Tests for the deployment-selection retry loops above.

use super::{execute_stream_with_selected_deployment, execute_with_selected_deployment};
use crate::core::providers::Provider;
use crate::core::providers::ProviderError;
use crate::core::providers::anthropic::{AnthropicConfig, AnthropicProvider};
use crate::core::providers::openai::OpenAIProvider;
use crate::core::router::RouterConfig;
use crate::core::router::{
    Deployment, DeploymentConfig, HealthStatus, UnifiedRouter, UnifiedRoutingStrategy,
};
use crate::core::types::model::ProviderCapability;
use crate::utils::error::gateway_error::GatewayError;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::Ordering;

async fn build_test_router() -> UnifiedRouter {
    let router = UnifiedRouter::default();
    let provider = Provider::OpenAI(
        OpenAIProvider::with_api_key("sk-test-key")
            .await
            .expect("test provider should build"),
    );

    router.add_deployment(Deployment::new(
        "deployment-1".to_string(),
        provider,
        "gpt-4o-mini".to_string(),
        "gpt-4".to_string(),
    ));

    router
}

async fn large_pending_settlement(marker: &u8) {
    let payload = [0_u8; 16 * 1024];
    std::future::pending::<()>().await;
    std::hint::black_box((payload, marker));
}

#[tokio::test]
async fn settlement_wrappers_box_large_borrowing_futures_before_polling() {
    let router = Arc::new(build_test_router().await);
    let deployment = router.get_deployment("deployment-1").unwrap();
    let (_, mut lease) = execute_stream_with_selected_deployment(
        router,
        "gpt-4",
        ProviderCapability::ChatCompletion,
        |_, _, _| async { Ok::<_, ProviderError>(()) },
    )
    .await
    .unwrap();
    let marker = 7_u8;
    assert!(std::mem::size_of_val(&large_pending_settlement(&marker)) >= 16 * 1024);

    let wait = lease.settle_terminal(42, None, large_pending_settlement(&marker));
    assert!(std::mem::size_of_val(&wait) < 1024);
    drop(wait);
    let wait = lease.settle_interrupted(42, None, large_pending_settlement(&marker));
    assert!(std::mem::size_of_val(&wait) < 1024);
    drop(wait);
    let wait = super::settle_stream_terminal(
        Some(&mut lease),
        42,
        None,
        large_pending_settlement(&marker),
    );
    assert!(std::mem::size_of_val(&wait) < 1024);
    drop(wait);
    let wait = super::settle_stream_terminal(None, 42, None, large_pending_settlement(&marker));
    assert!(std::mem::size_of_val(&wait) < 1024);
    drop(wait);
    // Eager allocation must not eagerly create terminal accounting intent.
    assert_eq!(deployment.state.total_requests.load(Ordering::Relaxed), 0);
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 1);
    drop(lease);
    assert_eq!(deployment.state.total_requests.load(Ordering::Relaxed), 0);
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
}

#[test]
fn unary_scope_does_not_inline_a_large_operation_future() {
    let router = UnifiedRouter::default();
    let request = execute_with_selected_deployment(
        &router,
        "gpt-4",
        ProviderCapability::ChatCompletion,
        |_, _, _| async {
            let payload = [0_u8; 128 * 1024];
            std::future::pending::<()>().await;
            std::hint::black_box(payload);
            Ok::<_, ProviderError>(((), 0))
        },
    );
    assert!(std::mem::size_of_val(&request) < 128 * 1024);
}

async fn build_mixed_capability_router() -> UnifiedRouter {
    let router = UnifiedRouter::new(RouterConfig {
        routing_strategy: UnifiedRoutingStrategy::PriorityBased,
        ..Default::default()
    });

    let chat_only_provider = Provider::Anthropic(
        AnthropicProvider::new(AnthropicConfig::new("sk-test-key"))
            .expect("test provider should build"),
    );
    let embedding_provider = Provider::OpenAI(
        OpenAIProvider::with_api_key("sk-test-key")
            .await
            .expect("test provider should build"),
    );

    router.add_deployment(
        Deployment::new(
            "chat-only".to_string(),
            chat_only_provider,
            "claude-3-haiku".to_string(),
            "shared-model".to_string(),
        )
        .with_config(DeploymentConfig {
            priority: 0,
            ..Default::default()
        }),
    );
    router.add_deployment(
        Deployment::new(
            "embedding-capable".to_string(),
            embedding_provider,
            "text-embedding-3-small".to_string(),
            "shared-model".to_string(),
        )
        .with_config(DeploymentConfig {
            priority: 10,
            ..Default::default()
        }),
    );

    router
}

async fn build_provider_budget_fallback_router() -> UnifiedRouter {
    let router = UnifiedRouter::new(RouterConfig {
        routing_strategy: UnifiedRoutingStrategy::PriorityBased,
        num_retries: 1,
        ..Default::default()
    });
    let primary = Provider::Anthropic(
        AnthropicProvider::new(AnthropicConfig::new("sk-test-key"))
            .expect("test provider should build"),
    );
    let fallback = Provider::OpenAI(
        OpenAIProvider::with_api_key("sk-test-key")
            .await
            .expect("test provider should build"),
    );

    router.add_deployment(
        Deployment::new(
            "primary-budget-exhausted".to_string(),
            primary,
            "claude-3-haiku".to_string(),
            "shared-model".to_string(),
        )
        .with_config(DeploymentConfig {
            priority: 0,
            ..Default::default()
        }),
    );
    router.add_deployment(
        Deployment::new(
            "fallback-provider".to_string(),
            fallback,
            "gpt-4o-mini".to_string(),
            "shared-model".to_string(),
        )
        .with_config(DeploymentConfig {
            priority: 10,
            ..Default::default()
        }),
    );

    router
}

#[tokio::test]
async fn test_execute_with_selected_deployment_uses_actual_deployment_model() {
    let router = build_test_router().await;

    let model = execute_with_selected_deployment(
        &router,
        "gpt-4",
        ProviderCapability::ChatCompletion,
        |_provider, model, _deployment_id| async { Ok((model, 0)) },
    )
    .await
    .expect("execution should succeed");

    assert_eq!(model, "gpt-4o-mini");
}

#[tokio::test]
async fn test_execute_with_selected_deployment_uses_capability_selected_deployment() {
    let router = build_mixed_capability_router().await;

    let (provider, model) =
        execute_with_selected_deployment(
            &router,
            "shared-model",
            ProviderCapability::Embeddings,
            |provider, model, _deployment_id| async move {
                Ok(((provider.name().to_string(), model), 0))
            },
        )
        .await
        .expect("execution should use an embeddings-capable deployment");

    assert_eq!(provider, "openai");
    assert_eq!(model, "text-embedding-3-small");
}

#[tokio::test]
async fn test_execute_with_selected_deployment_rejects_unavailable_capability() {
    let router = build_mixed_capability_router().await;
    let deployment = router
        .get_deployment("embedding-capable")
        .expect("deployment should exist");
    deployment
        .state
        .health
        .store(HealthStatus::Unhealthy as u8, Ordering::Relaxed);
    drop(deployment);

    let err = execute_with_selected_deployment(
        &router,
        "shared-model",
        ProviderCapability::Embeddings,
        |_provider, _model, _deployment_id| async { Ok::<_, ProviderError>(("should not run", 0)) },
    )
    .await
    .expect_err("unavailable capability should fail before execution");

    assert!(matches!(
        err,
        GatewayError::Provider(ProviderError::ProviderUnavailable { .. })
    ));
}

#[tokio::test]
async fn test_execute_with_selected_deployment_rejects_unsupported_capability() {
    let router = build_mixed_capability_router().await;

    let err = execute_with_selected_deployment(
        &router,
        "shared-model",
        ProviderCapability::CodeExecution,
        |_provider, _model, _deployment_id| async { Ok::<_, ProviderError>(("should not run", 0)) },
    )
    .await
    .expect_err("unsupported capability should fail before execution");

    assert!(matches!(
        err,
        GatewayError::Provider(ProviderError::InvalidRequest { .. })
    ));
}

#[tokio::test]
async fn test_execute_with_selected_deployment_maps_provider_error() {
    let router = build_test_router().await;

    let err = execute_with_selected_deployment(
        &router,
        "gpt-4",
        ProviderCapability::ChatCompletion,
        |_provider, _model, _deployment_id| async {
            Err::<(String, u64), _>(ProviderError::timeout("test", "timed out"))
        },
    )
    .await
    .expect_err("provider error should be mapped");

    assert!(matches!(
        err,
        GatewayError::Provider(ProviderError::Timeout { .. })
    ));
}

#[tokio::test]
async fn native_media_mutations_fail_over_after_explicit_rate_limit_response() {
    for capability in [
        ProviderCapability::ImageGeneration,
        ProviderCapability::ImageEdit,
    ] {
        let router = UnifiedRouter::new(RouterConfig {
            routing_strategy: UnifiedRoutingStrategy::PriorityBased,
            num_retries: 1,
            ..Default::default()
        });
        for (deployment_id, priority) in [("image-primary", 0), ("image-secondary", 10)] {
            let provider = Provider::OpenAI(
                OpenAIProvider::with_api_key("sk-test-key")
                    .await
                    .expect("test provider should build"),
            );
            router.add_deployment(
                Deployment::new(
                    deployment_id.to_string(),
                    provider,
                    "gpt-image-1".to_string(),
                    "image-model".to_string(),
                )
                .with_config(DeploymentConfig {
                    priority,
                    ..Default::default()
                }),
            );
        }
        let attempts = Arc::new(Mutex::new(Vec::new()));

        let result =
            execute_with_selected_deployment(&router, "image-model", capability.clone(), {
                let attempts = attempts.clone();
                move |_provider, _model, deployment_id| {
                    let attempts = attempts.clone();
                    async move {
                        attempts
                            .lock()
                            .expect("attempt lock")
                            .push(deployment_id.clone());
                        if deployment_id == "image-primary" {
                            Err(ProviderError::api_error(
                                "native-media",
                                429,
                                "rate limited",
                            ))
                        } else {
                            Ok((deployment_id, 0))
                        }
                    }
                }
            })
            .await
            .expect("explicit rate-limit response should fail over");

        assert_eq!(result, "image-secondary", "{capability:?}");
        assert_eq!(
            attempts.lock().expect("attempt lock").as_slice(),
            ["image-primary", "image-secondary"],
            "{capability:?}"
        );
    }
}

#[tokio::test]
async fn test_execute_with_selected_deployment_excludes_provider_budget_failures() {
    let router = build_provider_budget_fallback_router().await;
    let attempts = Arc::new(Mutex::new(Vec::new()));

    let result = execute_with_selected_deployment(
        &router,
        "shared-model",
        ProviderCapability::ChatCompletion,
        {
            let attempts = attempts.clone();
            move |provider, model, _deployment_id| {
                let attempts = attempts.clone();
                async move {
                    let provider_name = provider.name().to_string();
                    attempts.lock().unwrap().push(provider_name.clone());
                    if provider_name == "anthropic" {
                        Err(ProviderError::quota_exceeded(
                            "budget",
                            "provider 'anthropic' budget exceeded",
                        ))
                    } else {
                        Ok(((provider_name, model), 0))
                    }
                }
            }
        },
    )
    .await
    .expect("fallback provider should be selected after provider budget exhaustion");

    assert_eq!(result.0, "openai");
    assert_eq!(attempts.lock().unwrap().as_slice(), ["anthropic", "openai"]);
    let primary = router
        .get_deployment("primary-budget-exhausted")
        .expect("primary deployment should exist");
    assert_eq!(primary.state.fail_requests.load(Ordering::Relaxed), 0);
}

#[tokio::test]
async fn test_execute_with_selected_deployment_excludes_model_budget_failures() {
    let router = build_provider_budget_fallback_router().await;
    let attempts = Arc::new(Mutex::new(Vec::new()));

    let result = execute_with_selected_deployment(
        &router,
        "shared-model",
        ProviderCapability::ChatCompletion,
        {
            let attempts = attempts.clone();
            move |provider, model, _deployment_id| {
                let attempts = attempts.clone();
                async move {
                    attempts.lock().unwrap().push(model.clone());
                    if model == "claude-3-haiku" {
                        Err(ProviderError::quota_exceeded(
                            "budget",
                            "model 'claude-3-haiku' budget exceeded",
                        ))
                    } else {
                        Ok(((provider.name().to_string(), model), 0))
                    }
                }
            }
        },
    )
    .await
    .expect("fallback model should be selected after model budget exhaustion");

    assert_eq!(result.0, "openai");
    assert_eq!(result.1, "gpt-4o-mini");
    assert_eq!(
        attempts.lock().unwrap().as_slice(),
        ["claude-3-haiku", "gpt-4o-mini"]
    );
    let primary = router
        .get_deployment("primary-budget-exhausted")
        .expect("primary deployment should exist");
    assert_eq!(primary.state.fail_requests.load(Ordering::Relaxed), 0);
}

#[tokio::test]
async fn test_execute_stream_holds_deployment_active_until_success() {
    let router = Arc::new(build_test_router().await);

    let (_stream, lease) = execute_stream_with_selected_deployment(
        router.clone(),
        "gpt-4",
        ProviderCapability::ChatCompletionStream,
        |_provider, model, _selected_deployment_id| async move { Ok(model) },
    )
    .await
    .expect("stream creation should succeed");

    let deployment = router
        .get_deployment("deployment-1")
        .expect("deployment should exist");
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 1);
    assert_eq!(deployment.state.success_requests.load(Ordering::Relaxed), 0);
    drop(deployment);

    lease.finish_success(42).await;

    let deployment = router
        .get_deployment("deployment-1")
        .expect("deployment should exist");
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
    assert_eq!(deployment.state.success_requests.load(Ordering::Relaxed), 1);
    assert_eq!(deployment.state.tpm_current.load(Ordering::Relaxed), 42);
}

#[tokio::test]
async fn test_stream_lease_drop_releases_without_recording_outcome() {
    let router = Arc::new(build_test_router().await);

    let (_stream, lease) = execute_stream_with_selected_deployment(
        router.clone(),
        "gpt-4",
        ProviderCapability::ChatCompletionStream,
        |_provider, model, _selected_deployment_id| async move { Ok(model) },
    )
    .await
    .expect("stream creation should succeed");

    drop(lease);

    let deployment = router
        .get_deployment("deployment-1")
        .expect("deployment should exist");
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
    assert_eq!(deployment.state.total_requests.load(Ordering::Relaxed), 0);
}

#[tokio::test]
async fn test_execute_stream_records_stream_failure() {
    let router = Arc::new(build_test_router().await);

    let (_stream, lease) = execute_stream_with_selected_deployment(
        router.clone(),
        "gpt-4",
        ProviderCapability::ChatCompletionStream,
        |_provider, model, _selected_deployment_id| async move { Ok(model) },
    )
    .await
    .expect("stream creation should succeed");

    let error = ProviderError::rate_limit("test", Some(1));
    lease.finish_failure(&error).await;

    let deployment = router
        .get_deployment("deployment-1")
        .expect("deployment should exist");
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
    assert_eq!(deployment.state.fail_requests.load(Ordering::Relaxed), 1);
}

#[tokio::test]
async fn test_execute_stream_excludes_provider_budget_failures() {
    let router = Arc::new(build_provider_budget_fallback_router().await);
    let attempts = Arc::new(Mutex::new(Vec::new()));

    let ((provider_name, model), lease) = execute_stream_with_selected_deployment(
        router.clone(),
        "shared-model",
        ProviderCapability::ChatCompletionStream,
        {
            let attempts = attempts.clone();
            move |provider, model, _selected_deployment_id| {
                let attempts = attempts.clone();
                async move {
                    let provider_name = provider.name().to_string();
                    attempts.lock().unwrap().push(provider_name.clone());
                    if provider_name == "anthropic" {
                        Err(ProviderError::quota_exceeded(
                            "budget",
                            "provider 'anthropic' budget exceeded",
                        ))
                    } else {
                        Ok((provider_name, model))
                    }
                }
            }
        },
    )
    .await
    .expect("fallback stream provider should be selected after provider budget exhaustion");

    assert_eq!(provider_name, "openai");
    assert_eq!(model, "gpt-4o-mini");
    assert_eq!(attempts.lock().unwrap().as_slice(), ["anthropic", "openai"]);
    let primary = router
        .get_deployment("primary-budget-exhausted")
        .expect("primary deployment should exist");
    assert_eq!(primary.state.fail_requests.load(Ordering::Relaxed), 0);
    lease.finish_success(0).await;
}

#[tokio::test]
async fn test_execute_stream_startup_abort_releases_parallel_slot() {
    let router = Arc::new(UnifiedRouter::default());
    let Ok(openai_provider) = OpenAIProvider::with_api_key("sk-test-key").await else {
        panic!("test provider should build");
    };
    let provider = Provider::OpenAI(openai_provider);
    router.add_deployment(
        Deployment::new(
            "deployment-1".to_string(),
            provider,
            "gpt-4o-mini".to_string(),
            "gpt-4".to_string(),
        )
        .with_config(DeploymentConfig {
            max_parallel_requests: Some(1),
            ..Default::default()
        }),
    );

    let operation_entered = Arc::new(tokio::sync::Notify::new());
    let handle = {
        let router = router.clone();
        let operation_entered = operation_entered.clone();
        tokio::spawn(async move {
            execute_stream_with_selected_deployment(
                router,
                "gpt-4",
                ProviderCapability::ChatCompletionStream,
                move |_provider, _model, _selected_deployment_id| {
                    let operation_entered = operation_entered.clone();
                    async move {
                        operation_entered.notify_one();
                        std::future::pending::<Result<String, ProviderError>>().await
                    }
                },
            )
            .await
        })
    };

    operation_entered.notified().await;

    let Some(deployment) = router.get_deployment("deployment-1") else {
        panic!("deployment should exist");
    };
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 1);
    drop(deployment);

    handle.abort();
    match handle.await {
        Err(cancelled) => assert!(cancelled.is_cancelled()),
        Ok(_) => panic!("task should be cancelled"),
    }

    let Some(deployment) = router.get_deployment("deployment-1") else {
        panic!("deployment should exist");
    };
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
}

#[tokio::test]
async fn stream_finalization_transient_failure_gated_by_thresholds() {
    let router = Arc::new(UnifiedRouter::new(RouterConfig {
        allowed_fails: 3,
        cooldown_time_secs: 60,
        min_requests: 1,
        ..Default::default()
    }));
    let provider = Provider::OpenAI(
        OpenAIProvider::with_api_key("sk-test-key")
            .await
            .expect("test provider should build"),
    );
    router.add_deployment(Deployment::new(
        "deployment-1".to_string(),
        provider,
        "gpt-4o-mini".to_string(),
        "gpt-4".to_string(),
    ));

    let (_stream, lease) = execute_stream_with_selected_deployment(
        router.clone(),
        "gpt-4",
        ProviderCapability::ChatCompletionStream,
        |_provider, model, _selected_deployment_id| async move { Ok(model) },
    )
    .await
    .expect("stream creation should succeed");

    // A first mid-stream transient failure only counts toward the breaker
    // thresholds; it must not cool the deployment down on its own.
    lease
        .finish_failure(&ProviderError::timeout("openai", "mid-stream timeout"))
        .await;

    let deployment = router
        .get_deployment("deployment-1")
        .expect("deployment should exist");
    assert_eq!(deployment.state.fail_requests.load(Ordering::Relaxed), 1);
    assert!(!deployment.is_in_cooldown());
}

#[tokio::test]
async fn stream_finalization_failures_trip_cooldown_after_allowed_fails() {
    let router = Arc::new(UnifiedRouter::new(RouterConfig {
        allowed_fails: 3,
        cooldown_time_secs: 60,
        min_requests: 1,
        ..Default::default()
    }));
    let provider = Provider::OpenAI(
        OpenAIProvider::with_api_key("sk-test-key")
            .await
            .expect("test provider should build"),
    );
    router.add_deployment(Deployment::new(
        "deployment-1".to_string(),
        provider,
        "gpt-4o-mini".to_string(),
        "gpt-4".to_string(),
    ));

    for _ in 0..3 {
        let (_stream, lease) = execute_stream_with_selected_deployment(
            router.clone(),
            "gpt-4",
            ProviderCapability::ChatCompletionStream,
            |_provider, model, _selected_deployment_id| async move { Ok(model) },
        )
        .await
        .expect("stream creation should succeed");
        lease
            .finish_failure(&ProviderError::timeout("openai", "mid-stream timeout"))
            .await;
    }

    let deployment = router
        .get_deployment("deployment-1")
        .expect("deployment should exist");
    assert!(deployment.is_in_cooldown());
}

#[tokio::test]
async fn stream_finalization_fail_fast_errors_trip_immediate_cooldown() {
    let errors = [
        ProviderError::rate_limit("openai", Some(30)),
        ProviderError::api_error("openai", 429, "rate limited"),
        ProviderError::authentication("openai", "revoked key"),
        ProviderError::model_not_found("openai", "missing-model"),
    ];

    for error in errors {
        let router = Arc::new(UnifiedRouter::new(RouterConfig {
            allowed_fails: 3,
            cooldown_time_secs: 60,
            min_requests: 1,
            ..Default::default()
        }));
        let provider = Provider::OpenAI(
            OpenAIProvider::with_api_key("sk-test-key")
                .await
                .expect("test provider should build"),
        );
        router.add_deployment(Deployment::new(
            "deployment-1".to_string(),
            provider,
            "gpt-4o-mini".to_string(),
            "gpt-4".to_string(),
        ));

        let (_stream, lease) = execute_stream_with_selected_deployment(
            router.clone(),
            "gpt-4",
            ProviderCapability::ChatCompletionStream,
            |_provider, model, _selected_deployment_id| async move { Ok(model) },
        )
        .await
        .expect("stream creation should succeed");
        lease.finish_failure(&error).await;

        let deployment = router
            .get_deployment("deployment-1")
            .expect("deployment should exist");
        assert!(
            deployment.is_in_cooldown(),
            "{error:?} should bypass transient failure thresholds"
        );
    }
}

#[tokio::test]
async fn test_execute_stream_returns_last_error_when_budget_excludes_every_deployment() {
    let router = Arc::new(build_provider_budget_fallback_router().await);
    let attempts = Arc::new(Mutex::new(Vec::new()));

    let result = execute_stream_with_selected_deployment(
        router,
        "shared-model",
        ProviderCapability::ChatCompletionStream,
        {
            let attempts = attempts.clone();
            move |_provider, _model, deployment_id| {
                let attempts = attempts.clone();
                async move {
                    attempts.lock().unwrap().push(deployment_id.clone());
                    Err::<String, _>(ProviderError::quota_exceeded(
                        "budget",
                        format!("provider budget exhausted for {deployment_id}"),
                    ))
                }
            }
        },
    )
    .await;
    let error = match result {
        Err(error) => error,
        Ok((_stream, lease)) => {
            drop(lease);
            panic!("hard-excluded stream candidates should preserve the last budget error");
        }
    };

    match error {
        GatewayError::Provider(ProviderError::QuotaExceeded { message, .. }) => {
            assert_eq!(message, "provider budget exhausted for fallback-provider");
        }
        other => panic!("expected the last budget error, got {other:?}"),
    }
    assert_eq!(
        attempts.lock().unwrap().as_slice(),
        ["primary-budget-exhausted", "fallback-provider"]
    );
}

#[cfg(feature = "websockets")]
#[tokio::test]
async fn cancelled_realtime_admission_restores_shared_rpm() {
    use crate::config::models::storage::RedisConfig;
    use crate::storage::redis::RedisPool;
    let Ok(url) = std::env::var("REDIS_URL") else {
        assert!(std::env::var("CI").is_err(), "REDIS_URL is required in CI");
        return;
    };
    let pool = Arc::new(
        RedisPool::new(&RedisConfig {
            url,
            enabled: true,
            allow_degraded: false,
            ..RedisConfig::default()
        })
        .await
        .unwrap(),
    );
    let router = UnifiedRouter::default().with_admission_redis(pool.clone());
    let id = uuid::Uuid::new_v4().to_string();
    let provider = Provider::OpenAI(
        OpenAIProvider::with_api_key("sk-admission-test")
            .await
            .unwrap(),
    );
    router.add_deployment(
        Deployment::new(id.clone(), provider, "gpt-4o-mini".into(), "shared".into()).with_config(
            DeploymentConfig {
                rpm_limit: Some(1),
                max_parallel_requests: Some(1),
                ..Default::default()
            },
        ),
    );
    let (_, mut lease) = execute_stream_with_selected_deployment(
        Arc::new(router),
        "shared",
        ProviderCapability::ChatCompletion,
        |_, _, _| async { Ok::<_, ProviderError>(()) },
    )
    .await
    .unwrap();
    // As on an idle Realtime socket, release the handshake before each generation.
    lease.cancel_response().await;
    for _ in 0..3 {
        lease.begin_response(1).await.unwrap();
        // A rejected budget never forwards the generation and must cancel RPM.
        lease.cancel_response().await;
    }
    pool.delete(&RedisPool::admission_key(&id)).await.unwrap();
}

async fn cancellation_test_router() -> Option<(
    Arc<UnifiedRouter>,
    Arc<crate::storage::redis::RedisPool>,
    String,
)> {
    cancellation_test_router_with_local_limit(None).await
}

async fn cancellation_test_router_with_local_limit(
    local_tpm: Option<u64>,
) -> Option<(
    Arc<UnifiedRouter>,
    Arc<crate::storage::redis::RedisPool>,
    String,
)> {
    use crate::config::models::storage::RedisConfig;
    let Ok(url) = std::env::var("REDIS_URL") else {
        assert!(std::env::var("CI").is_err(), "REDIS_URL is required in CI");
        return None;
    };
    let pool = Arc::new(
        crate::storage::redis::RedisPool::new(&RedisConfig {
            url,
            enabled: true,
            allow_degraded: false,
            ..Default::default()
        })
        .await
        .unwrap(),
    );
    let router = UnifiedRouter::default().with_circuit_redis(pool.clone());
    let router = if local_tpm.is_some() {
        router
    } else {
        router.with_admission_redis(pool.clone())
    };
    let id = uuid::Uuid::new_v4().to_string();
    router.add_deployment(
        Deployment::new(
            id.clone(),
            Provider::OpenAI(OpenAIProvider::with_api_key("sk-test-key").await.unwrap()),
            "gpt-4o-mini".into(),
            "gpt-4".into(),
        )
        .with_config(DeploymentConfig {
            max_parallel_requests: local_tpm.is_none().then_some(1),
            tpm_limit: local_tpm,
            ..Default::default()
        }),
    );
    Some((Arc::new(router), pool, id))
}

fn assert_one_completed_request(deployment: &Deployment) {
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
    assert_eq!(deployment.state.success_requests.load(Ordering::Relaxed), 1);
    assert_eq!(deployment.state.total_requests.load(Ordering::Relaxed), 1);
    assert_eq!(deployment.state.rpm_current.load(Ordering::Relaxed), 1);
    assert_eq!(deployment.state.tpm_current.load(Ordering::Relaxed), 42);
}

#[tokio::test(flavor = "current_thread")]
async fn local_http_observation_replaces_estimate_before_cancelled_circuit_wait() {
    const CHILD: &str = "LITELLM_HTTP_LOCAL_OBSERVATION_CHILD";
    if std::env::var_os(CHILD).is_none() {
        // The deliberate 64-slot pause belongs to this isolated test process.
        let output = tokio::task::spawn_blocking(|| {
            std::process::Command::new(std::env::current_exe().unwrap())
                .args([
                    "--exact",
                    "server::routes::ai::execution::tests::local_http_observation_replaces_estimate_before_cancelled_circuit_wait",
                    "--nocapture",
                ])
                .env(CHILD, "1")
                .output()
                .unwrap()
        })
        .await
        .unwrap();
        let stdout = String::from_utf8_lossy(&output.stdout);
        assert!(
            output.status.success() && stdout.contains("1 passed; 0 failed"),
            "isolated HTTP observation must pass once:\n{stdout}\n{}",
            String::from_utf8_lossy(&output.stderr)
        );
        return;
    }

    for mode in [
        "unary",
        "success",
        "failure",
        #[cfg(feature = "websockets")]
        "interrupted",
    ] {
        let Some((router, pool, id)) = cancellation_test_router_with_local_limit(Some(10)).await
        else {
            return;
        };
        let deployment = router.get_deployment(&id).unwrap();
        let paused = Arc::new(Mutex::new(None));
        let error = ProviderError::timeout("openai", "completed stream failure");
        let mut stream = if mode == "unary" {
            None
        } else {
            Some(
                execute_stream_with_selected_deployment(
                    router.clone(),
                    "gpt-4",
                    ProviderCapability::ChatCompletionStream,
                    |_, _, _| async { Ok(()) },
                )
                .await
                .unwrap()
                .1,
            )
        };
        if mode != "unary" {
            *paused.lock().unwrap() = Some(crate::core::router::circuit::pause_circuit_io().await);
        }
        let pause_in_provider = paused.clone();
        let mut completion = Box::pin(async {
            if mode == "unary" {
                execute_with_selected_deployment(
                    &router,
                    "gpt-4",
                    ProviderCapability::ChatCompletion,
                    move |_, _, _| {
                        let paused = pause_in_provider.clone();
                        async move {
                            *paused.lock().unwrap() =
                                Some(crate::core::router::circuit::pause_circuit_io().await);
                            Ok(((), 6))
                        }
                    },
                )
                .await
                .unwrap();
            } else {
                let lease = stream.as_mut().unwrap();
                match mode {
                    "success" => lease.complete_response(6, None).await,
                    "failure" => lease.complete_response(6, Some(&error)).await,
                    #[cfg(feature = "websockets")]
                    "interrupted" => lease.finish_interrupted(6, Some(&error)).await,
                    _ => unreachable!(),
                }
            }
        });
        let observed = async {
            tokio::time::timeout(std::time::Duration::from_secs(3), async {
                while deployment.state.tpm_current.load(Ordering::Relaxed) != 6 {
                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
                }
            })
            .await
            .expect("known tokens must be observed before the paused circuit wait");
        };
        tokio::select! {
            () = completion.as_mut() => panic!("shared circuit I/O must remain paused"),
            () = observed => {}
        }
        assert!(futures::poll!(completion.as_mut()).is_pending());
        let expected_rpm = u64::from(mode != "failure");
        let expected_success = u64::from(matches!(mode, "unary" | "success"));
        assert_eq!(
            deployment.state.rpm_current.load(Ordering::Relaxed),
            expected_rpm
        );
        assert_eq!(deployment.state.total_requests.load(Ordering::Relaxed), 1);
        assert_eq!(
            deployment.state.success_requests.load(Ordering::Relaxed),
            expected_success
        );
        assert_eq!(
            deployment.state.fail_requests.load(Ordering::Relaxed),
            1 - expected_success
        );
        // Probe the production local-admission boundary directly: selecting
        // through the router would itself wait on the deliberately paused
        // circuit bridge. Six observed tokens plus four must fit a limit of
        // ten; the unreplaced one-token HTTP estimate would reject it.
        let next = deployment
            .state
            .reserve_local_tokens(4, 10)
            .expect("known usage must replace the owned estimate before yielding");
        drop(next);
        assert_eq!(shared_failure_counts(&pool, &id).await, (0, 0));
        drop(completion);
        if let Some(lease) = &mut stream {
            match mode {
                "success" => lease.complete_response(6, None).await,
                "failure" => lease.complete_response(6, Some(&error)).await,
                #[cfg(feature = "websockets")]
                "interrupted" => lease.finish_interrupted(6, Some(&error)).await,
                _ => unreachable!(),
            }
        }
        drop(stream);
        drop(paused.lock().unwrap().take());
        assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
        assert_eq!(deployment.state.tpm_current.load(Ordering::Relaxed), 6);
        assert_eq!(
            deployment.state.rpm_current.load(Ordering::Relaxed),
            expected_rpm
        );
        assert_eq!(deployment.state.total_requests.load(Ordering::Relaxed), 1);
        assert_eq!(
            deployment.state.success_requests.load(Ordering::Relaxed),
            expected_success
        );
        assert_eq!(
            deployment.state.fail_requests.load(Ordering::Relaxed),
            1 - expected_success
        );
        let next = router
            .select_deployment_lease_with_tokens("gpt-4", 4)
            .unwrap();
        drop(next);
        // Cancellation before a circuit slot is acquired publishes no ACK.
        assert_eq!(shared_failure_counts(&pool, &id).await, (0, 0));
        pool.delete(&crate::storage::redis::RedisPool::circuit_key(&id))
            .await
            .unwrap();
    }
}

#[tokio::test(flavor = "current_thread")]
async fn local_http_terminal_cancellation_replaces_estimate_exactly_once() {
    for mode in ["success", "failure", "interrupted_failure", "interrupted"] {
        let Some((router, _, id)) = cancellation_test_router_with_local_limit(Some(10)).await
        else {
            return;
        };
        let (_, mut lease) = execute_stream_with_selected_deployment(
            router.clone(),
            "gpt-4",
            ProviderCapability::ChatCompletionStream,
            |_, _, _| async { Ok(()) },
        )
        .await
        .unwrap();
        let deployment = router.get_deployment(&id).unwrap();
        let error = ProviderError::timeout("openai", "terminal cancellation");
        let mut settlement = Box::pin(async {
            match mode {
                "success" => {
                    lease
                        .settle_terminal(6, None, std::future::pending::<()>())
                        .await
                }
                "failure" => {
                    lease
                        .settle_terminal(6, Some(&error), std::future::pending::<()>())
                        .await
                }
                "interrupted_failure" => {
                    lease
                        .settle_interrupted(6, Some(&error), std::future::pending::<()>())
                        .await
                }
                "interrupted" => {
                    lease
                        .settle_interrupted(6, None, std::future::pending::<()>())
                        .await
                }
                _ => unreachable!(),
            }
        });
        assert!(futures::poll!(settlement.as_mut()).is_pending());
        drop(settlement);
        lease.complete_response(6, None).await;
        drop(lease);
        assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
        assert_eq!(deployment.state.tpm_current.load(Ordering::Relaxed), 6);
        assert_eq!(
            deployment.state.rpm_current.load(Ordering::Relaxed),
            u64::from(mode != "failure")
        );
        assert_eq!(
            deployment.state.total_requests.load(Ordering::Relaxed),
            u64::from(mode != "interrupted")
        );
        assert_eq!(
            deployment.state.success_requests.load(Ordering::Relaxed),
            u64::from(mode == "success")
        );
        assert_eq!(
            deployment.state.fail_requests.load(Ordering::Relaxed),
            u64::from(matches!(mode, "failure" | "interrupted_failure"))
        );
        let next = router
            .select_deployment_lease_with_tokens("gpt-4", 4)
            .unwrap();
        drop(next);
    }
}

#[cfg(feature = "websockets")]
#[tokio::test(flavor = "current_thread")]
async fn local_http_neutral_completion_replaces_estimate_exactly_once() {
    let Some((router, _, id)) = cancellation_test_router_with_local_limit(Some(10)).await else {
        return;
    };
    let (_, mut lease) = execute_stream_with_selected_deployment(
        router.clone(),
        "gpt-4",
        ProviderCapability::ChatCompletionStream,
        |_, _, _| async { Ok(()) },
    )
    .await
    .unwrap();
    lease.finish_neutral(6).await;
    lease.finish_neutral(6).await;
    drop(lease);
    let deployment = router.get_deployment(&id).unwrap();
    assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
    assert_eq!(deployment.state.tpm_current.load(Ordering::Relaxed), 6);
    assert_eq!(deployment.state.rpm_current.load(Ordering::Relaxed), 0);
    assert_eq!(deployment.state.total_requests.load(Ordering::Relaxed), 0);
    let next = router
        .select_deployment_lease_with_tokens("gpt-4", 4)
        .unwrap();
    drop(next);
}

#[tokio::test(flavor = "current_thread")]
async fn cancelled_stream_completion_records_local_success_exactly_once() {
    let Some((router, pool, id)) = cancellation_test_router().await else {
        return;
    };
    let (_, mut lease) = execute_stream_with_selected_deployment(
        router.clone(),
        "gpt-4",
        ProviderCapability::ChatCompletionStream,
        |_, _, _| async { Ok(()) },
    )
    .await
    .unwrap();
    let deployment = router.get_deployment(&id).unwrap();
    let slots = crate::core::router::admission::pause_admission_io().await;
    let mut completion = Box::pin(lease.complete_response(42, None));
    tokio::select! {
        () = completion.as_mut() => panic!("admission settlement must remain paused"),
        () = wait_for_shared_success(&pool, &id) => {}
    }
    assert!(futures::poll!(completion.as_mut()).is_pending());
    assert_one_completed_request(&deployment);
    drop(completion);
    // Retry after cancellation, then drop the lease: neither may count twice
    // nor replace the queued actual-token settlement with cancellation.
    lease.complete_response(42, None).await;
    assert_one_completed_request(&deployment);
    drop(lease);
    drop(slots);
    assert_one_completed_request(&deployment);
    wait_for_actual_settlement(&pool, &id).await;
    assert_shared_success_once(&pool, &id).await;
}

#[tokio::test(flavor = "current_thread")]
async fn cancelled_gateway_unary_settlement_retains_local_success() {
    let Some((router, pool, id)) = cancellation_test_router().await else {
        return;
    };
    let (started, ready) = tokio::sync::oneshot::channel();
    let started = Arc::new(Mutex::new(Some(started)));
    let worker = router.clone();
    let task = tokio::spawn(async move {
        execute_with_selected_deployment(
            &worker,
            "gpt-4",
            ProviderCapability::ChatCompletion,
            move |_, _, _| {
                let started = started.clone();
                async move {
                    let slots = crate::core::router::admission::pause_admission_io().await;
                    let _ = started.lock().unwrap().take().unwrap().send(());
                    Ok((slots, 42))
                }
            },
        )
        .await
    });
    tokio::time::timeout(std::time::Duration::from_secs(10), ready)
        .await
        .unwrap()
        .unwrap();
    wait_for_shared_success(&pool, &id).await;
    task.abort();
    assert!(task.await.unwrap_err().is_cancelled());
    assert_one_completed_request(&router.get_deployment(&id).unwrap());
    wait_for_actual_settlement(&pool, &id).await;
    assert_shared_success_once(&pool, &id).await;
}

#[tokio::test(flavor = "current_thread")]
async fn cancelled_stream_failure_publishes_shared_failure_before_admission_cleanup() {
    for tokens in [0, 42] {
        let Some((router, pool, id)) = cancellation_test_router().await else {
            return;
        };
        let (_, mut lease) = execute_stream_with_selected_deployment(
            router.clone(),
            "gpt-4",
            ProviderCapability::ChatCompletionStream,
            |_, _, _| async { Ok(()) },
        )
        .await
        .unwrap();
        let deployment = router.get_deployment(&id).unwrap();
        let key = crate::storage::redis::RedisPool::admission_key(&id);
        let mut conn = pool.open_live_connection().await.unwrap();
        let slots = crate::core::router::admission::pause_admission_io().await;
        let error = ProviderError::timeout("openai", "mid-stream timeout");
        let mut completion = Box::pin(lease.complete_response(tokens, Some(&error)));
        let shared_failure = async {
            tokio::time::timeout(std::time::Duration::from_secs(3), async {
                loop {
                    if shared_failure_counts(&pool, &id).await == (1, 1) {
                        break;
                    }
                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
                }
            })
            .await
            .expect("shared circuit failure must precede paused admission cleanup");
        };
        tokio::select! {
            () = completion.as_mut() => panic!("admission cleanup must remain paused"),
            () = shared_failure => {}
        }
        assert!(futures::poll!(completion.as_mut()).is_pending());
        let parallel: i64 = redis::cmd("HGET")
            .arg(&key)
            .arg("p")
            .query_async(&mut conn)
            .await
            .unwrap();
        assert_eq!(parallel, 1, "admission liability remains held");
        drop(completion);
        // Retrying a finalized failure and dropping its lease must not replace
        // the prepared settlement/cancellation or publish the failure twice.
        lease.complete_response(tokens, Some(&error)).await;
        drop(lease);
        assert_eq!(deployment.state.active_requests.load(Ordering::Relaxed), 0);
        assert_eq!(deployment.state.success_requests.load(Ordering::Relaxed), 0);
        assert_eq!(deployment.state.fail_requests.load(Ordering::Relaxed), 1);
        assert_eq!(deployment.state.total_requests.load(Ordering::Relaxed), 1);
        assert_eq!(deployment.state.rpm_current.load(Ordering::Relaxed), 0);
        assert_eq!(deployment.state.tpm_current.load(Ordering::Relaxed), tokens);
        drop(slots);
        tokio::time::timeout(std::time::Duration::from_secs(3), async {
            loop {
                let state: (i64, i64, i64) = redis::cmd("HMGET")
                    .arg(&key)
                    .arg(&["p", "r", "t"])
                    .query_async(&mut conn)
                    .await
                    .unwrap();
                let fields: Vec<String> = redis::cmd("HKEYS")
                    .arg(&key)
                    .query_async(&mut conn)
                    .await
                    .unwrap();
                if state == (0, 0, tokens as i64)
                    && !fields.iter().any(|field| field.starts_with("l:"))
                {
                    break;
                }
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
            }
        })
        .await
        .expect("prepared admission outcome must survive cancellation");
        assert_eq!(shared_failure_counts(&pool, &id).await, (1, 1));
        pool.delete(&key).await.unwrap();
        pool.delete(&crate::storage::redis::RedisPool::circuit_key(&id))
            .await
            .unwrap();
    }
}

async fn shared_failure_counts(pool: &crate::storage::redis::RedisPool, id: &str) -> (i64, i64) {
    let key = crate::storage::redis::RedisPool::circuit_key(id);
    let mut conn = pool.open_live_connection().await.unwrap();
    let counts: (Option<i64>, Option<i64>) = redis::cmd("HMGET")
        .arg(&key)
        .arg(&["tot", "fail"])
        .query_async(&mut conn)
        .await
        .unwrap();
    (counts.0.unwrap_or(0), counts.1.unwrap_or(0))
}

async fn shared_success_counts(pool: &crate::storage::redis::RedisPool, id: &str) -> (i64, i64) {
    let key = crate::storage::redis::RedisPool::circuit_key(id);
    let mut conn = pool.open_live_connection().await.unwrap();
    let counts: (Option<i64>, Option<i64>) = redis::cmd("HMGET")
        .arg(&key)
        .arg(&["tot", "r"])
        .query_async(&mut conn)
        .await
        .unwrap();
    (counts.0.unwrap_or(0), counts.1.unwrap_or(0))
}

async fn wait_for_shared_success(pool: &crate::storage::redis::RedisPool, id: &str) {
    tokio::time::timeout(std::time::Duration::from_secs(3), async {
        loop {
            if shared_success_counts(pool, id).await == (1, 1) {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }
    })
    .await
    .expect("shared circuit success must precede the paused admission settlement");
}

async fn assert_shared_success_once(pool: &crate::storage::redis::RedisPool, id: &str) {
    assert_eq!(shared_success_counts(pool, id).await, (1, 1));
    pool.delete(&crate::storage::redis::RedisPool::circuit_key(id))
        .await
        .unwrap();
}

async fn wait_for_actual_settlement(pool: &crate::storage::redis::RedisPool, id: &str) {
    let key = crate::storage::redis::RedisPool::admission_key(id);
    let mut conn = pool.open_live_connection().await.unwrap();
    tokio::time::timeout(std::time::Duration::from_secs(3), async {
        loop {
            let state: (i64, i64) = redis::cmd("HMGET")
                .arg(&key)
                .arg(&["p", "t"])
                .query_async(&mut conn)
                .await
                .unwrap();
            if state == (0, 42) {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }
    })
    .await
    .unwrap();
    pool.delete(&key).await.unwrap();
}