skippy-server 0.76.1

use super::support::*;
use super::*;

#[test]
fn proactive_eviction_attrs_are_bounded_and_request_free() {
    let attrs = proactive_eviction_attrs("error", Some("inactive_session"), 1024, 2, 768);

    assert_eq!(
        attrs.get("skippy.kv.decision"),
        Some(&json!("proactive_eviction"))
    );
    assert_eq!(
        attrs.get(attr_key::KV_PROACTIVE_EVICTION_STATUS),
        Some(&json!("error"))
    );
    assert_eq!(
        attrs.get(attr_key::KV_PROACTIVE_EVICTION_ERROR_KIND),
        Some(&json!("inactive_session"))
    );
    assert_eq!(
        attrs.get(attr_key::KV_PROACTIVE_EVICTION_TARGET_TOKENS),
        Some(&json!(1024))
    );
    assert_eq!(
        attrs.get(attr_key::KV_PROACTIVE_EVICTED_ENTRIES),
        Some(&json!(2))
    );
    assert_eq!(
        attrs.get(attr_key::KV_PROACTIVE_EVICTED_TOKENS),
        Some(&json!(768))
    );
    assert!(!attrs.contains_key(attr_key::REQUEST_ID));
    assert!(!attrs.contains_key(attr_key::SESSION_ID));
    assert!(!attrs.contains_key("openai.prompt_cache_key"));
    assert!(!attrs.contains_key("openai.prompt_cache_retention"));
}

#[test]
fn resident_capacity_rejection_is_side_effect_free_and_retryable() {
    let config = prefix_cache_test_config();
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let mut runtime = crate::runtime_state::RuntimeState::new_modelless_with_capacity_for_test(
        config.lane_count,
        8,
    );
    let before = runtime.session_stats();

    let first = kv
        .admit_resident_capacity(&mut runtime, "request", 9, 1, 1, None)
        .unwrap();
    let second = kv
        .admit_resident_capacity(&mut runtime, "request", 9, 1, 1, None)
        .unwrap();
    let recovered = kv
        .admit_resident_capacity(&mut runtime, "request", 4, 1, 1, None)
        .unwrap();

    assert!(!first.admitted);
    assert!(recovered.admitted);
    assert_eq!(second.active_tokens, first.active_tokens);
    assert_eq!(
        second.admission_deficit_tokens,
        first.admission_deficit_tokens
    );
    assert_eq!(runtime.session_stats(), before);

    let response = crate::frontend::local_generation::resident_capacity_admission_error(&first)
        .into_response();
    assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
    assert_eq!(response.headers().get("retry-after").unwrap(), "1");
}

#[test]
fn resident_capacity_unknown_fails_closed() {
    let config = prefix_cache_test_config();
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let mut runtime = crate::runtime_state::RuntimeState::new_modelless_for_test(1);

    let decision = kv
        .admit_resident_capacity(&mut runtime, "unknown", 1, 0, 0, None)
        .unwrap();

    assert!(!decision.capacity_known);
    assert!(!decision.admitted);
    assert_eq!(decision.admission_deficit_tokens, 1);
}

#[test]
fn resident_capacity_keeps_decode_watermark_separate_from_prompt_reservation() {
    let config = prefix_cache_test_config();
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let mut runtime = crate::runtime_state::RuntimeState::new_modelless_with_capacity_for_test(
        config.lane_count,
        31_613,
    );
    let _reservation = kv
        .reserve_resident_capacity("queued", 31_101)
        .unwrap()
        .expect("resident reservation");

    let decision = kv
        .admit_resident_capacity(&mut runtime, "active", 0, 512, 512, None)
        .unwrap();

    assert!(decision.admitted);
    assert_eq!(decision.active_tokens, 0);
    assert_eq!(decision.inflight_outstanding_tokens, 31_101);
    assert_eq!(decision.projected_free_tokens, 512);
    assert_eq!(decision.admission_deficit_tokens, 0);
}

#[test]
fn resident_capacity_admission_evicts_for_the_aggregate_active_four_wave() {
    let config = StageConfig {
        ctx_size: 131_072,
        lane_count: 4,
        kv_cache: Some(StageKvCacheConfig {
            max_entries: 32,
            min_tokens: 64,
            ..prefix_cache_test_config()
                .kv_cache
                .expect("test cache config")
        }),
        ..prefix_cache_test_config()
    };
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let mut runtime = crate::runtime_state::RuntimeState::new_modelless_with_capacity_for_test(
        config.lane_count,
        config.ctx_size,
    );

    // The failed qualification's warm-up left 98,028 resident cells. Model
    // that physical occupancy with sixteen independent resident paths.
    for index in 0..16 {
        let token_count = if index < 12 { 6_127 } else { 6_126 };
        let mut base = prefix_cache_test_base();
        base.chat_template_id = Some(format!("warm-family-{index}"));
        let tokens = (0..token_count).collect::<Vec<_>>();
        let identity = kv.prefill_identity(&config, &base, 0, &tokens);
        seed_resident_prefix(&kv, &identity);
    }
    assert_eq!(kv.radix.lock().unwrap().stats().resident_tokens, 98_028);

    // Four unrelated requests restore only the shared chat-template prefix.
    // Their outstanding suffix plus bounded decode demand is 47,419 tokens.
    let wave = [
        ("measured-1", 147_u64, 9_209_u64),
        ("measured-2", 81, 10_551),
        ("measured-3", 146, 12_439),
        ("measured-4", 146, 15_092),
    ];
    let mut reservations = Vec::new();
    for (session_id, restored_tokens, suffix_tokens) in wave {
        runtime.track_session_tokens_for_test(session_id, restored_tokens);
        reservations.push(
            kv.reserve_resident_capacity(
                session_id,
                restored_tokens
                    .saturating_add(suffix_tokens)
                    .saturating_add(32),
            )
            .unwrap()
            .expect("resident reservation"),
        );
    }

    let decision = kv
        .admit_resident_capacity(&mut runtime, "measured-1", 9_241, 2_048, 2_080, None)
        .unwrap();

    assert_eq!(decision.inflight_reservations, 4);
    assert_eq!(decision.inflight_outstanding_tokens, 47_419);
    assert_eq!(decision.request_tokens, 47_419);
    assert!(decision.admitted);
    assert!(decision.evicted_entries >= 3);
    assert!(decision.physical_evicted_tokens >= 16_455);
    assert!(decision.projected_free_tokens >= 2_048);

    drop(reservations);
    let after_release = kv
        .admit_resident_capacity(&mut runtime, "measured-1", 0, 0, 0, None)
        .unwrap();
    assert_eq!(after_release.inflight_reservations, 0);
    assert_eq!(after_release.inflight_outstanding_tokens, 0);
}

#[test]
fn openai_cache_stats_default_to_disabled() {
    let stats = GenerationCacheStats::default();

    assert_eq!(stats.status, "disabled");
    assert_eq!(stats.cached_prompt_tokens, 0);
    assert_eq!(stats.matched_prefix_tokens, 0);
    assert_eq!(stats.suffix_prefill_tokens, 0);
    assert_eq!(stats.hit_kind, None);
}

#[test]
fn cache_identity_reuses_repeated_prompts_without_client_cache_key() {
    let config = prefix_cache_test_config();
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let first_request = prefix_cache_base_with_request("request-a", "session-a");
    let second_request = prefix_cache_base_with_request("request-b", "session-b");
    let tokens = (0..1024).collect::<Vec<_>>();

    let recorded = kv.prefill_identity(&config, &first_request, 0, &tokens);
    let looked_up = kv.prefill_identity(&config, &second_request, 0, &tokens);

    assert_eq!(recorded.page_id, looked_up.page_id);
    assert_eq!(
        recorded.identity.prefix_hash,
        looked_up.identity.prefix_hash
    );
    assert_ne!(recorded.identity.session_id, looked_up.identity.session_id);

    seed_resident_prefix(&kv, &recorded);
    let hit = kv
        .probe_resident_prefix(&looked_up)
        .expect("repeated prompt should hit without a client cache key");
    assert_eq!(hit.page_id, recorded.page_id);
    assert_eq!(hit.token_count, tokens.len());
}

#[test]
fn cache_identity_namespaces_explicit_prompt_cache_keys() {
    let config = prefix_cache_test_config();
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let tokens = (0..1024).collect::<Vec<_>>();
    let mut default_namespace = prefix_cache_test_base();
    default_namespace.chat_template_id = None;
    let mut explicit_namespace = prefix_cache_test_base();
    explicit_namespace.chat_template_id = Some("openai:prompt_cache_key:tenant-a".to_string());

    let default_identity = kv.prefill_identity(&config, &default_namespace, 0, &tokens);
    let explicit_identity = kv.prefill_identity(&config, &explicit_namespace, 0, &tokens);

    assert_ne!(default_identity.page_id, explicit_identity.page_id);
    assert_ne!(
        default_identity.identity.prefix_hash,
        explicit_identity.identity.prefix_hash
    );
}

#[test]
fn disabled_cache_config_has_no_stage_integration() {
    let config = StageConfig {
        kv_cache: Some(StageKvCacheConfig {
            mode: StageKvCacheMode::Disabled,
            ..prefix_cache_test_config()
                .kv_cache
                .expect("test cache config")
        }),
        ..prefix_cache_test_config()
    };

    let kv =
        KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense).unwrap();

    assert!(kv.is_none());
}

#[test]
fn cold_resident_prefix_lookup_misses_before_recording() {
    let config = prefix_cache_test_config();
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let identity = kv.prefill_identity(
        &config,
        &prefix_cache_test_base(),
        0,
        &(0..1024).collect::<Vec<_>>(),
    );

    assert!(kv.probe_resident_prefix(&identity).is_none());
}

#[test]
fn resident_prefix_cache_hits_radix_common_prefix_without_a_record_ladder() {
    let config = prefix_cache_test_config();
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let base = prefix_cache_test_base();
    let recorded_tokens = (0..2214).collect::<Vec<_>>();
    let mut lookup_tokens = recorded_tokens.clone();
    lookup_tokens[2176] = 100_000;
    lookup_tokens.extend(100_001..100_017);
    let record_plan = crate::frontend::prefix_cache::stage0_full_prefill_record_identities(
        &kv,
        &config,
        &base,
        &recorded_tokens,
    );
    let lookup_plan = kv.lookup_identities(&config, &base, 0, &lookup_tokens);
    assert_eq!(record_plan.len(), 1);
    assert_eq!(lookup_plan.len(), 1);
    let recorded = &record_plan[0];
    let lookup = &lookup_plan[0];

    seed_resident_prefix(&kv, recorded);
    let hit = kv
        .probe_resident_prefix(lookup)
        .expect("different-tail prompt should hit the radix common prefix");

    assert_eq!(hit.page_id, recorded.page_id);
    assert_eq!(hit.token_count, 2176);
}

#[test]
fn resident_prefix_cache_rejects_common_prefix_below_configured_minimum() {
    let config = StageConfig {
        kv_cache: Some(StageKvCacheConfig {
            min_tokens: 64,
            ..prefix_cache_test_config()
                .kv_cache
                .expect("test cache config")
        }),
        ..prefix_cache_test_config()
    };
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let donor = prefix_cache_base_with_request("donor-request", "donor-session");
    let receiver = prefix_cache_base_with_request("receiver-request", "receiver-session");
    let recorded_tokens = (0..20_751).collect::<Vec<_>>();
    let mut lookup_tokens = recorded_tokens[..27].to_vec();
    lookup_tokens.extend(100_000..128_065);
    let recorded = kv.prefill_identity(&config, &donor, 0, &recorded_tokens);
    let lookup = kv.prefill_identity(&config, &receiver, 0, &lookup_tokens);

    seed_resident_prefix(&kv, &recorded);

    assert!(kv.probe_resident_prefix(&lookup).is_none());
    assert_eq!(
        kv.peek_cache_affinity(&config, &[lookup]),
        skippy_scheduler::CacheAffinity::default()
    );
}

#[test]
fn stage0_full_prefill_uses_one_radix_path_per_request() {
    let config = prefix_cache_test_config();
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let base = prefix_cache_test_base();
    let recorded_tokens = (0..2214).collect::<Vec<_>>();
    let mut lookup_tokens = recorded_tokens.clone();
    lookup_tokens.extend(100_000..100_017);

    let record_plan = crate::frontend::prefix_cache::stage0_full_prefill_record_identities(
        &kv,
        &config,
        &base,
        &recorded_tokens,
    );
    let lookup_plan = kv.lookup_identities(&config, &base, 0, &lookup_tokens);

    let record_counts = record_plan
        .iter()
        .map(|identity| identity.identity.token_count)
        .collect::<Vec<_>>();
    let lookup_counts = lookup_plan
        .iter()
        .map(|identity| identity.identity.token_count)
        .collect::<Vec<_>>();

    assert_eq!(record_counts, vec![2214]);
    assert_eq!(lookup_counts, vec![2231]);
    assert_eq!(record_plan[0].namespace, lookup_plan[0].namespace);
    assert_ne!(record_plan[0].page_id, lookup_plan[0].page_id);
}

#[test]
fn stage0_chunked_prefill_uses_one_radix_path_per_request() {
    let config = prefix_cache_test_config();
    let kv = KvStageIntegration::from_config(&config, skippy_runtime::ModelStateKind::Dense)
        .unwrap()
        .expect("resident prefix cache enabled");
    let base = prefix_cache_test_base();
    let recorded_tokens = (0..2214).collect::<Vec<_>>();
    let mut lookup_tokens = recorded_tokens.clone();
    lookup_tokens.extend(100_000..100_017);

    let record_plan = crate::frontend::prefix_cache::stage0_prefill_record_identities(
        &kv,
        &config,
        &base,
        0,
        &recorded_tokens,
    );
    let lookup_plan = kv.lookup_identities(&config, &base, 0, &lookup_tokens);

    let record_counts = record_plan
        .iter()
        .map(|identity| identity.identity.token_count)
        .collect::<Vec<_>>();
    let lookup_counts = lookup_plan
        .iter()
        .map(|identity| identity.identity.token_count)
        .collect::<Vec<_>>();

    assert_eq!(record_counts, vec![2214]);
    assert_eq!(lookup_counts, vec![2231]);
    assert_eq!(record_plan[0].namespace, lookup_plan[0].namespace);
    assert_ne!(record_plan[0].page_id, lookup_plan[0].page_id);
}