// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
// SPDX-License-Identifier: Apache-2.0

//! HTTP contract for static P/D buckets.
//!
//! Buckets narrow the candidate domain before policy selection. Prefill SLO
//! profiles may override rank, while decode uses `input_tokens + max_tokens`.

use axum::body::Body;
use axum::http::{Request, StatusCode};
use sgl_kv_indexer::{PrefixIndex, PrefixIndexError, PrefixMatch, PrefixOutcome};
use sgl_router::config::{
    AffinityConfig, BucketConfig, BucketSpec, BucketStage, CacheAwareConfig, CachePrefixProvider,
    Config, DiscoveryBackend, InflightLoadConfig, KvIndexerEndpointConfig, ModelConfig,
    ObservabilityConfig, PolicyKind, ProxyConfig, ServerConfig, SessionAffinityMode,
    SloBucketPolicy, StaticUrlsDiscoveryConfig,
};
use sgl_router::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
use sgl_router::policies::factory::build_registry_with_defaults;
use sgl_router::proxy::Proxy;
use sgl_router::server::app::build_router;
use sgl_router::server::app_context::AppContext;
use sgl_router::state::load_monitor::engine_reported_load::{LoadStat, NativeCacheRankLoad};
use sgl_router::tokenizer::TokenizerRegistry;
use sgl_router::workers::WorkerRegistry;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tower::ServiceExt;

fn bucket(id: &str, stage: BucketStage, rank: u32, worker_id: &str) -> BucketSpec {
    BucketSpec {
        id: id.into(),
        stage,
        rank,
        worker_ids: vec![worker_id.into()],
        min_extend_tokens: None,
        max_extend_tokens: None,
        min_sequence_tokens: None,
        max_sequence_tokens: None,
        max_context_tokens: Some(16_384),
        ttft_p95_at_capacity_ms: None,
        tps_p05_at_capacity: None,
        max_pending_prefill_tokens: None,
    }
}

fn build_app_context(
    specs: Vec<WorkerSpec>,
    bucket_config: BucketConfig,
    policy: PolicyKind,
    affinity: Option<AffinityConfig>,
) -> AppContext {
    let config = Config {
        server: ServerConfig {
            host: "0".into(),
            port: 0,
            ..Default::default()
        },
        observability: ObservabilityConfig::default(),
        model: ModelConfig {
            id: "tiny".into(),
            tokenizer_path: "tests/fixtures/tiny_tokenizer.json".into(),
            disable_input_ids_forwarding: false,
            policy,
            decode_policy: Default::default(),
            bucket_config: Some(bucket_config),
            circuit_breaker: None,
            cache_aware: None,
            sticky: None,
            affinity,
            fused: None,
            eligibility: None,
            sampling_overrides: Default::default(),
        },
        discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
            urls: vec!["http://placeholder:0".into()],
        }),
        proxy: ProxyConfig::default(),
        router_inflight_load: InflightLoadConfig::default(),
    };
    let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&config).unwrap());
    let registry = Arc::new(WorkerRegistry::default());
    for spec in specs {
        let _ = registry.add(spec);
    }
    let policies = Arc::new(build_registry_with_defaults(&config).unwrap());
    let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap());
    AppContext::new(config, tokenizers, proxy, registry, policies)
}

fn build_ctx(
    specs: Vec<WorkerSpec>,
    bucket_config: BucketConfig,
    policy: PolicyKind,
    affinity: Option<AffinityConfig>,
) -> Arc<AppContext> {
    Arc::new(build_app_context(specs, bucket_config, policy, affinity))
}

struct FakePrefixIndex {
    address: Option<String>,
    calls: AtomicUsize,
}

impl FakePrefixIndex {
    fn matched(address: String) -> Arc<Self> {
        Arc::new(Self {
            address: Some(address),
            calls: AtomicUsize::new(0),
        })
    }

    fn no_signal() -> Arc<Self> {
        Arc::new(Self {
            address: None,
            calls: AtomicUsize::new(0),
        })
    }
}

#[tonic::async_trait]
impl PrefixIndex for FakePrefixIndex {
    async fn match_prefix(&self, hashes: Vec<i64>) -> Result<PrefixOutcome, PrefixIndexError> {
        self.calls.fetch_add(1, Ordering::Relaxed);
        let Some(address) = &self.address else {
            return Ok(PrefixOutcome::Empty);
        };
        let matched_prefix_blocks =
            u32::try_from(hashes.len().saturating_sub(1)).unwrap_or(u32::MAX);
        Ok(PrefixOutcome::Matched {
            matches: vec![PrefixMatch {
                address: address.clone(),
                matched_prefix_blocks,
                worker_id: "fake-index-worker".into(),
            }],
            best_prefix_blocks: matched_prefix_blocks,
        })
    }
}

struct TwoPrefixIndex {
    best_address: String,
    lower_ranked_address: String,
}

impl TwoPrefixIndex {
    fn new(best_address: String, lower_ranked_address: String) -> Arc<Self> {
        Arc::new(Self {
            best_address,
            lower_ranked_address,
        })
    }
}

#[tonic::async_trait]
impl PrefixIndex for TwoPrefixIndex {
    async fn match_prefix(&self, hashes: Vec<i64>) -> Result<PrefixOutcome, PrefixIndexError> {
        let best_prefix_blocks = u32::try_from(hashes.len().saturating_sub(1)).unwrap_or(u32::MAX);
        let lower_ranked_prefix_blocks = (best_prefix_blocks / 2).max(1);
        Ok(PrefixOutcome::Matched {
            matches: vec![
                PrefixMatch {
                    address: self.best_address.clone(),
                    matched_prefix_blocks: best_prefix_blocks,
                    worker_id: "best-index-worker".into(),
                },
                PrefixMatch {
                    address: self.lower_ranked_address.clone(),
                    matched_prefix_blocks: lower_ranked_prefix_blocks,
                    worker_id: "lower-index-worker".into(),
                },
            ],
            best_prefix_blocks,
        })
    }
}

fn build_cache_ctx(
    specs: Vec<WorkerSpec>,
    bucket_config: BucketConfig,
    prefix_index: Arc<dyn PrefixIndex>,
) -> Arc<AppContext> {
    build_cache_ctx_with_affinity(
        specs,
        bucket_config,
        prefix_index,
        AffinityConfig::default(),
    )
}

fn build_cache_ctx_with_affinity(
    specs: Vec<WorkerSpec>,
    bucket_config: BucketConfig,
    prefix_index: Arc<dyn PrefixIndex>,
    affinity: AffinityConfig,
) -> Arc<AppContext> {
    let mut context =
        build_app_context(specs, bucket_config, PolicyKind::CacheAware, Some(affinity));
    context.config.model.cache_aware = Some(CacheAwareConfig {
        prefix_provider: CachePrefixProvider::Indexer,
        kv_indexer_endpoint: Some(KvIndexerEndpointConfig {
            url: "http://fake-indexer".into(),
            query_timeout_ms: 100,
            query_max_inflight: 32,
        }),
    });
    context.prefix_index = Some(prefix_index);
    context.block_size_oracle.try_set(1).unwrap();
    Arc::new(context)
}

fn worker_spec(id: &str, url: String, mode: WorkerMode) -> WorkerSpec {
    WorkerSpec {
        id: WorkerId(id.into()),
        url,
        mode,
        model_ids: vec![ModelId("tiny".into())],
        bootstrap_port: (mode == WorkerMode::Prefill).then_some(8997),
    }
}

fn set_native_load(
    ctx: &AppContext,
    worker_url: &str,
    num_total_tokens: u64,
    max_total_num_tokens: u64,
) {
    set_native_load_with_waiting(ctx, worker_url, num_total_tokens, max_total_num_tokens, 0);
}

fn set_native_load_with_waiting(
    ctx: &AppContext,
    worker_url: &str,
    num_total_tokens: u64,
    max_total_num_tokens: u64,
    num_waiting_uncached_tokens: u64,
) {
    ctx.engine_reported_load.set(
        worker_url,
        0,
        LoadStat {
            num_running_reqs: 0,
            num_waiting_reqs: 0,
            num_tokens: num_total_tokens,
            max_total_num_tokens,
            native_cache: Some(NativeCacheRankLoad {
                num_waiting_uncached_tokens,
                num_total_tokens,
                max_running_requests: 64,
                total_prefill_uncached_tokens: 1,
                total_prefill_busy_us: 1,
            }),
        },
        Instant::now(),
    );
}

fn chat_request(ttft_slo_ms: Option<u64>, max_tokens: Option<u64>) -> Request<Body> {
    chat_request_with_content("bucket routing", ttft_slo_ms, max_tokens, None)
}

fn chat_request_with_content(
    content: &str,
    ttft_slo_ms: Option<u64>,
    max_tokens: Option<u64>,
    session_id: Option<&str>,
) -> Request<Body> {
    let mut builder = Request::builder()
        .method("POST")
        .uri("/v1/chat/completions")
        .header("content-type", "application/json");
    if let Some(ttft_slo_ms) = ttft_slo_ms {
        builder = builder.header("x-sgl-ttft-slo-ms", ttft_slo_ms.to_string());
    }
    if let Some(session_id) = session_id {
        builder = builder.header("x-session-id", session_id);
    }
    builder
        .body(Body::from(
            serde_json::to_vec(&serde_json::json!({
                "model": "tiny",
                "messages": [{"role": "user", "content": content}],
                "max_tokens": max_tokens,
            }))
            .unwrap(),
        ))
        .unwrap()
}

async fn wait_for_prefill(mock: &crate::common::mock_worker::MockWorker) {
    tokio::time::timeout(Duration::from_secs(2), async {
        loop {
            if mock.captured.lock().unwrap().last_body.is_some() {
                return;
            }
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await
    .expect("selected prefill worker must receive the detached request");
}

async fn wait_for_prefill_body_containing(
    mock: &crate::common::mock_worker::MockWorker,
    expected: &str,
) -> Vec<u8> {
    tokio::time::timeout(Duration::from_secs(2), async {
        loop {
            let captured = mock.captured.lock().unwrap().last_body.clone();
            if let Some(body) = captured {
                if String::from_utf8_lossy(&body).contains(expected) {
                    return body.to_vec();
                }
            }
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
    })
    .await
    .expect("selected prefill worker must receive the expected request body")
}

#[tokio::test]
async fn prefill_slo_first_uses_eligible_ttft_bucket_before_lower_rank_bucket() {
    let cheap = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let fast = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut cheap_bucket = bucket("p-cheap", BucketStage::Prefill, 10, "p-cheap");
    cheap_bucket.ttft_p95_at_capacity_ms = Some(400);
    let mut fast_bucket = bucket("p-fast", BucketStage::Prefill, 20, "p-fast");
    fast_bucket.ttft_p95_at_capacity_ms = Some(100);
    let bucket_config = BucketConfig {
        buckets: vec![
            cheap_bucket,
            fast_bucket,
            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::SloFirst,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let ctx = build_ctx(
        vec![
            worker_spec("p-cheap", cheap.url.clone(), WorkerMode::Prefill),
            worker_spec("p-fast", fast.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        PolicyKind::PowerOfTwo,
        None,
    );

    let response = build_router(ctx)
        .oneshot(chat_request(Some(200), Some(16)))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    wait_for_prefill(&fast).await;
    assert!(
        cheap.captured.lock().unwrap().last_body.is_none(),
        "lower-rank but TTFT-ineligible P Bucket must not be dispatched first"
    );
}

#[tokio::test]
async fn prefill_tries_later_compatible_bucket_before_capacity_fallback() {
    let full = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let available = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let bucket_config = BucketConfig {
        buckets: vec![
            bucket("p-full", BucketStage::Prefill, 10, "p-full"),
            bucket("p-available", BucketStage::Prefill, 20, "p-available"),
            bucket("d", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let ctx = build_ctx(
        vec![
            worker_spec("p-full", full.url.clone(), WorkerMode::Prefill),
            worker_spec("p-available", available.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        PolicyKind::PowerOfTwo,
        None,
    );
    set_native_load(&ctx, &full.url, 100, 100);
    set_native_load(&ctx, &available.url, 0, 10_000);

    let response = build_router(ctx)
        .oneshot(chat_request(None, Some(16)))
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    wait_for_prefill(&available).await;
    assert!(
        full.captured.lock().unwrap().last_body.is_none(),
        "capacity fallback must wait until all compatible prefill buckets are exhausted"
    );
}

#[tokio::test]
async fn decode_bucket_uses_input_plus_requested_output_budget() {
    let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let short_decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let long_decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut short_bucket = bucket("d-short", BucketStage::Decode, 20, "d-short");
    short_bucket.max_sequence_tokens = Some(1_024);
    let mut long_bucket = bucket("d-long", BucketStage::Decode, 30, "d-long");
    long_bucket.min_sequence_tokens = Some(1_025);
    let bucket_config = BucketConfig {
        buckets: vec![
            bucket("p", BucketStage::Prefill, 10, "p"),
            short_bucket,
            long_bucket,
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let ctx = build_ctx(
        vec![
            worker_spec("p", prefill.url.clone(), WorkerMode::Prefill),
            worker_spec("d-short", short_decode.url.clone(), WorkerMode::Decode),
            worker_spec("d-long", long_decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        PolicyKind::PowerOfTwo,
        None,
    );

    let response = build_router(ctx)
        .oneshot(chat_request(None, Some(2_000)))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    assert_eq!(
        response
            .headers()
            .get("x-sgl-decode-url")
            .and_then(|value| value.to_str().ok()),
        Some(long_decode.url.as_str()),
        "peak sequence length must exclude the short Decode Bucket"
    );
    assert!(
        long_decode.captured.lock().unwrap().last_body.is_some(),
        "the selected long Decode worker is awaited before the response"
    );
    assert!(
        short_decode.captured.lock().unwrap().last_body.is_none(),
        "the incompatible short Decode Bucket must not receive the request"
    );
}

#[tokio::test]
async fn decode_tries_later_compatible_bucket_before_capacity_fallback() {
    let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let full = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let available = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let bucket_config = BucketConfig {
        buckets: vec![
            bucket("p", BucketStage::Prefill, 10, "p"),
            bucket("d-full", BucketStage::Decode, 20, "d-full"),
            bucket("d-available", BucketStage::Decode, 30, "d-available"),
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let ctx = build_ctx(
        vec![
            worker_spec("p", prefill.url.clone(), WorkerMode::Prefill),
            worker_spec("d-full", full.url.clone(), WorkerMode::Decode),
            worker_spec("d-available", available.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        PolicyKind::PowerOfTwo,
        None,
    );
    set_native_load(&ctx, &full.url, 100, 100);
    set_native_load(&ctx, &available.url, 0, 10_000);

    let response = build_router(ctx)
        .oneshot(chat_request(None, Some(16)))
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    assert_eq!(
        response
            .headers()
            .get("x-sgl-decode-url")
            .and_then(|value| value.to_str().ok()),
        Some(available.url.as_str()),
        "capacity fallback must wait until all compatible decode buckets are exhausted"
    );
}

#[tokio::test]
async fn prefill_only_bucket_configuration_keeps_global_decode_routing() {
    let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let bucket_config = BucketConfig {
        buckets: vec![bucket("p", BucketStage::Prefill, 10, "p")],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let ctx = build_ctx(
        vec![
            worker_spec("p", prefill.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        PolicyKind::PowerOfTwo,
        None,
    );

    let response = build_router(ctx)
        .oneshot(chat_request(None, Some(16)))
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    assert_eq!(
        response
            .headers()
            .get("x-sgl-decode-url")
            .and_then(|value| value.to_str().ok()),
        Some(decode.url.as_str()),
        "a Prefill-only Bucket rollout must retain the Step 1 global Decode domain"
    );
}

#[tokio::test]
async fn global_rebind_session_affinity_can_keep_a_cross_length_bucket_primary() {
    let short = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let long = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-short");
    short_bucket.max_extend_tokens = Some(256);
    short_bucket.max_context_tokens = Some(16_384);
    short_bucket.ttft_p95_at_capacity_ms = Some(80);
    let mut long_bucket = bucket("p-long", BucketStage::Prefill, 20, "p-long");
    long_bucket.min_extend_tokens = Some(257);
    long_bucket.max_context_tokens = Some(16_384);
    long_bucket.ttft_p95_at_capacity_ms = Some(300);
    let bucket_config = BucketConfig {
        buckets: vec![
            short_bucket,
            long_bucket,
            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::SloFirst,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let ctx = build_ctx(
        vec![
            worker_spec("p-short", short.url.clone(), WorkerMode::Prefill),
            worker_spec("p-long", long.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        PolicyKind::SessionAware,
        Some(AffinityConfig {
            session_affinity_mode: SessionAffinityMode::GlobalRebind,
            ..Default::default()
        }),
    );
    let app = build_router(ctx);

    let first = app
        .clone()
        .oneshot(chat_request_with_content(
            "short",
            Some(120),
            Some(8),
            Some("s-1"),
        ))
        .await
        .unwrap();
    assert_eq!(first.status(), StatusCode::OK);
    wait_for_prefill(&short).await;

    let long_content = "length ".repeat(128);
    let second = app
        .oneshot(chat_request_with_content(
            &long_content,
            Some(120),
            Some(8),
            Some("s-1"),
        ))
        .await
        .unwrap();
    assert_eq!(second.status(), StatusCode::OK);
    let short_body = wait_for_prefill_body_containing(&short, &long_content).await;
    assert!(
        String::from_utf8_lossy(&short_body).contains(&long_content),
        "the second, long request must retain the existing cross-Bucket session primary"
    );
    assert!(
        long.captured.lock().unwrap().last_body.is_none(),
        "target length Bucket is skipped only because the primary's own Hard TTFT profile is eligible"
    );
}

#[tokio::test]
async fn global_preserve_establishes_then_reuses_a_new_assignment() {
    let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let bucket_config = BucketConfig {
        buckets: vec![
            bucket("p", BucketStage::Prefill, 10, "p"),
            bucket("d", BucketStage::Decode, 20, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let ctx = build_ctx(
        vec![
            worker_spec("p", prefill.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        PolicyKind::SessionAware,
        Some(AffinityConfig {
            session_affinity_mode: SessionAffinityMode::GlobalPreserve,
            ..Default::default()
        }),
    );
    let app = build_router(Arc::clone(&ctx));

    for content in ["first global request", "second global request"] {
        let response = app
            .clone()
            .oneshot(chat_request_with_content(
                content,
                None,
                Some(8),
                Some("global-session"),
            ))
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::OK);
    }

    let metrics = ctx.metrics.render();
    assert!(
        metrics.contains(
            r#"sgl_router_policy_decisions_total{policy="session_aware",reason="assigned"} 1"#
        ),
        "the first global-preserve request must establish an assignment: {metrics}"
    );
    assert!(
        metrics.contains(
            r#"sgl_router_policy_decisions_total{policy="session_aware",reason="session_primary"} 1"#
        ),
        "the second global-preserve request must reuse the assignment: {metrics}"
    );
}

#[tokio::test]
async fn cache_winner_uses_target_uncached_work_before_prompt_length_bucket() {
    let short = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let long = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-short");
    short_bucket.max_extend_tokens = Some(8);
    let mut long_bucket = bucket("p-long", BucketStage::Prefill, 20, "p-long");
    long_bucket.min_extend_tokens = Some(9);
    let bucket_config = BucketConfig {
        buckets: vec![
            short_bucket,
            long_bucket,
            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let index = FakePrefixIndex::matched(short.url.clone());
    let prefix_index: Arc<dyn PrefixIndex> = index.clone();
    let ctx = build_cache_ctx(
        vec![
            worker_spec("p-short", short.url.clone(), WorkerMode::Prefill),
            worker_spec("p-long", long.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        prefix_index,
    );

    let content = "cached-prefix ".repeat(128);
    let response = build_router(ctx)
        .oneshot(chat_request_with_content(&content, None, Some(8), None))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    wait_for_prefill(&short).await;
    assert!(
        long.captured.lock().unwrap().last_body.is_none(),
        "a cache winner with small target-specific uncached work must not be replaced by the full-length Bucket"
    );
    assert_eq!(
        index.calls.load(Ordering::Relaxed),
        1,
        "the async Indexer query must run once at ingress, not once per Bucket"
    );
}

#[tokio::test]
async fn unbucketed_global_cache_winner_remains_eligible() {
    let best = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let lower_ranked = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut lower_ranked_bucket = bucket("p-lower", BucketStage::Prefill, 20, "p-lower");
    lower_ranked_bucket.min_extend_tokens = Some(32);
    let bucket_config = BucketConfig {
        buckets: vec![
            lower_ranked_bucket,
            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let index: Arc<dyn PrefixIndex> =
        TwoPrefixIndex::new(best.url.clone(), lower_ranked.url.clone());
    let ctx = build_cache_ctx_with_affinity(
        vec![
            worker_spec("p-best", best.url.clone(), WorkerMode::Prefill),
            worker_spec("p-lower", lower_ranked.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        index,
        AffinityConfig {
            cache_candidate_min_workers: 1,
            cache_candidate_ratio: 0.0,
            cache_candidate_max_workers: 1,
            ..AffinityConfig::default()
        },
    );

    let content = "cached bucket candidate ".repeat(256);
    let response = build_router(Arc::clone(&ctx))
        .oneshot(chat_request_with_content(&content, None, Some(8), None))
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    wait_for_prefill(&best).await;
    assert!(
        lower_ranked.captured.lock().unwrap().last_body.is_none(),
        "an admitted unbucketed cache holder must remain globally eligible"
    );
    assert!(
        ctx.metrics.render().contains(
            r#"sgl_router_policy_decisions_total{policy="cache_aware",reason="cache_candidate"} 1"#
        ),
        "the global cache winner must remain a cache candidate"
    );
}

#[tokio::test]
async fn cache_candidate_over_context_limit_falls_back_to_compatible_bucket() {
    let cached = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let fallback = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-cached");
    short_bucket.max_context_tokens = Some(32);
    let bucket_config = BucketConfig {
        buckets: vec![
            short_bucket,
            bucket("p-long", BucketStage::Prefill, 20, "p-fallback"),
            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let prefix_index: Arc<dyn PrefixIndex> = FakePrefixIndex::matched(cached.url.clone());
    let ctx = build_cache_ctx(
        vec![
            worker_spec("p-cached", cached.url.clone(), WorkerMode::Prefill),
            worker_spec("p-fallback", fallback.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        prefix_index,
    );

    let content = "context limit ".repeat(128);
    let response = build_router(ctx)
        .oneshot(chat_request_with_content(&content, None, Some(8), None))
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    wait_for_prefill(&fallback).await;
    assert!(
        cached.captured.lock().unwrap().last_body.is_none(),
        "a cache holder that cannot serve the full context must not receive the request"
    );
}

#[tokio::test]
async fn cache_candidate_outside_slo_first_tier_falls_back_to_eligible_bucket() {
    let cached = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let fallback = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut slow_bucket = bucket("p-slow", BucketStage::Prefill, 10, "p-cached");
    slow_bucket.ttft_p95_at_capacity_ms = Some(400);
    let mut fast_bucket = bucket("p-fast", BucketStage::Prefill, 20, "p-fallback");
    fast_bucket.ttft_p95_at_capacity_ms = Some(100);
    let bucket_config = BucketConfig {
        buckets: vec![
            slow_bucket,
            fast_bucket,
            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::SloFirst,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let prefix_index: Arc<dyn PrefixIndex> = FakePrefixIndex::matched(cached.url.clone());
    let ctx = build_cache_ctx(
        vec![
            worker_spec("p-cached", cached.url.clone(), WorkerMode::Prefill),
            worker_spec("p-fallback", fallback.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        prefix_index,
    );

    let content = "ttft tier ".repeat(128);
    let response = build_router(ctx)
        .oneshot(chat_request_with_content(
            &content,
            Some(200),
            Some(8),
            None,
        ))
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    wait_for_prefill(&fallback).await;
    assert!(
        cached.captured.lock().unwrap().last_body.is_none(),
        "SloFirst must exclude a cache holder outside the request's TTFT tier"
    );
}

#[tokio::test]
async fn cache_candidate_respects_its_bucket_pending_prefill_budget() {
    let cached = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let fallback = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut limited_bucket = bucket("p-limited", BucketStage::Prefill, 10, "p-cached");
    limited_bucket.max_pending_prefill_tokens = Some(64);
    let bucket_config = BucketConfig {
        buckets: vec![
            limited_bucket,
            bucket("p-fallback", BucketStage::Prefill, 20, "p-fallback"),
            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let prefix_index: Arc<dyn PrefixIndex> = FakePrefixIndex::matched(cached.url.clone());
    let ctx = build_cache_ctx(
        vec![
            worker_spec("p-cached", cached.url.clone(), WorkerMode::Prefill),
            worker_spec("p-fallback", fallback.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        prefix_index,
    );
    set_native_load_with_waiting(&ctx, &cached.url, 0, 100_000, 64);

    let content = "pending budget ".repeat(128);
    let response = build_router(Arc::clone(&ctx))
        .oneshot(chat_request_with_content(&content, None, Some(8), None))
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    wait_for_prefill(&fallback).await;
    assert!(
        cached.captured.lock().unwrap().last_body.is_none(),
        "a cache holder beyond its Bucket pending-prefill budget must not receive the request"
    );
    assert!(ctx
        .metrics
        .render()
        .contains("sgl_router_cache_admission_rejected_total 1"));
}

#[tokio::test]
async fn rejected_global_cache_candidate_falls_back_to_prompt_length_bucket() {
    let cached = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let fallback = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-cached");
    short_bucket.max_extend_tokens = Some(8);
    let mut long_bucket = bucket("p-long", BucketStage::Prefill, 20, "p-fallback");
    long_bucket.min_extend_tokens = Some(9);
    let bucket_config = BucketConfig {
        buckets: vec![
            short_bucket,
            long_bucket,
            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let prefix_index: Arc<dyn PrefixIndex> = FakePrefixIndex::matched(cached.url.clone());
    let ctx = build_cache_ctx(
        vec![
            worker_spec("p-cached", cached.url.clone(), WorkerMode::Prefill),
            worker_spec("p-fallback", fallback.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        prefix_index,
    );
    set_native_load(&ctx, &cached.url, 1, 1);

    let content = "admission fallback ".repeat(128);
    let response = build_router(Arc::clone(&ctx))
        .oneshot(chat_request_with_content(&content, None, Some(8), None))
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    wait_for_prefill(&fallback).await;
    assert!(
        cached.captured.lock().unwrap().last_body.is_none(),
        "an admission-rejected cache holder must not receive the request"
    );
    let metrics = ctx.metrics.render();
    assert!(metrics.contains("sgl_router_cache_admission_rejected_total 1"));
    assert!(
        !metrics.contains(
            r#"sgl_router_policy_decisions_total{policy="cache_aware",reason="cache_candidate"} 1"#
        ),
        "Bucket fallback must not be reported as a cache-candidate decision"
    );
}

#[tokio::test]
async fn cache_no_signal_restarts_normal_prompt_length_bucket_fallback() {
    let short = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let long = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let decode = crate::common::mock_worker::MockWorker::start(vec![]).await;
    let mut short_bucket = bucket("p-short", BucketStage::Prefill, 10, "p-short");
    short_bucket.max_extend_tokens = Some(8);
    let mut long_bucket = bucket("p-long", BucketStage::Prefill, 20, "p-long");
    long_bucket.min_extend_tokens = Some(9);
    let bucket_config = BucketConfig {
        buckets: vec![
            short_bucket,
            long_bucket,
            bucket("d-catch-all", BucketStage::Decode, 30, "d"),
        ],
        ttft_slo_policy: SloBucketPolicy::Disabled,
        tps_slo_policy: SloBucketPolicy::Disabled,
    };
    let index = FakePrefixIndex::no_signal();
    let prefix_index: Arc<dyn PrefixIndex> = index.clone();
    let ctx = build_cache_ctx(
        vec![
            worker_spec("p-short", short.url.clone(), WorkerMode::Prefill),
            worker_spec("p-long", long.url.clone(), WorkerMode::Prefill),
            worker_spec("d", decode.url.clone(), WorkerMode::Decode),
        ],
        bucket_config,
        prefix_index,
    );

    let content = "uncached-prompt ".repeat(128);
    let response = build_router(ctx)
        .oneshot(chat_request_with_content(&content, None, Some(8), None))
        .await
        .unwrap();
    assert_eq!(response.status(), StatusCode::OK);
    wait_for_prefill(&long).await;
    assert!(
        short.captured.lock().unwrap().last_body.is_none(),
        "without a cache winner the request must restart the normal full-input Bucket path"
    );
    assert_eq!(index.calls.load(Ordering::Relaxed), 1);
}