kcode_context_cache_policy/

lib.rs

//! Deterministic cache-stability and sparse context-marker policy.

#![forbid(unsafe_code)]

use std::collections::BTreeSet;

/// Minimum estimated context growth between size markers in one cache epoch.
pub const CONTEXT_SIZE_MARKER_INTERVAL_TOKENS: u64 = 10_000;

/// Cache behavior expected from the relationship between two submissions.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum CacheExpectation {
    /// No prior submitted cacheable projection exists.
    ColdStart,
    /// The prior cacheable projection is an exact prefix of the candidate.
    ExpectedWarm,
    /// A known rewrite invalidated the previous cacheable suffix.
    PlannedInvalidation { reason: String },
}

impl CacheExpectation {
    /// Stable structured-log label.
    pub const fn label(&self) -> &'static str {
        match self {
            Self::ColdStart => "cold_start",
            Self::ExpectedWarm => "expected_warm",
            Self::PlannedInvalidation { .. } => "planned_invalidation",
        }
    }

    /// Planned invalidation reason, when applicable.
    pub fn planned_reason(&self) -> Option<&str> {
        match self {
            Self::PlannedInvalidation { reason } => Some(reason),
            _ => None,
        }
    }

    /// Whether a provider cache hit is expected.
    pub const fn expects_cache_hit(&self) -> bool {
        matches!(self, Self::ExpectedWarm)
    }
}

/// Exact prior cacheable input and its non-text provider-material identity.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PreviousProjection<'a> {
    pub text: &'a str,
    pub material_fingerprint: &'a str,
}

/// Classifies one candidate projection without normalizing its bytes.
pub fn classify(
    previous: Option<PreviousProjection<'_>>,
    current_text: &str,
    current_material_fingerprint: &str,
    rewrite_reason: &str,
) -> CacheExpectation {
    let Some(previous) = previous else {
        return CacheExpectation::ColdStart;
    };
    if previous.material_fingerprint != current_material_fingerprint {
        return CacheExpectation::PlannedInvalidation {
            reason: "provider_material_changed".into(),
        };
    }
    if current_text.starts_with(previous.text) {
        return CacheExpectation::ExpectedWarm;
    }
    CacheExpectation::PlannedInvalidation {
        reason: nonblank_reason(rewrite_reason),
    }
}

fn nonblank_reason(reason: &str) -> String {
    let reason = reason.trim();
    if reason.is_empty() {
        "other_projection_rewrite".into()
    } else {
        reason.into()
    }
}

/// Marker state reconstructed for the current cache epoch.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct MarkerState {
    pub reported_stale_boxes: BTreeSet<u64>,
    pub last_context_size_tokens: Option<u64>,
}

/// Current inputs to sparse marker policy.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MarkerObservation {
    pub stale_boxes: Vec<u64>,
    pub current_context_tokens: u64,
    pub context_limit_tokens: u64,
    pub expectation: CacheExpectation,
}

/// One stale-box marker to append.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum StaleMarker {
    New(Vec<u64>),
    Consolidated(Vec<u64>),
}

impl StaleMarker {
    /// Exact model-visible marker text.
    pub fn render(&self) -> String {
        let (label, ids) = match self {
            Self::New(ids) => ("new stale boxes", ids),
            Self::Consolidated(ids) => ("stale boxes", ids),
        };
        let ids = ids
            .iter()
            .map(u64::to_string)
            .collect::<Vec<_>>()
            .join(", ");
        format!("[{label}: {ids}]")
    }
}

/// Complete marker transition selected for one provider preparation.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MarkerDecision {
    pub reset_epoch: bool,
    pub stale: Option<StaleMarker>,
    pub context_size_tokens: Option<u64>,
    pub next_state: MarkerState,
}

/// Selects all sparse markers due for one provider preparation.
pub fn decide_markers(state: &MarkerState, observation: MarkerObservation) -> MarkerDecision {
    let current_stale = observation.stale_boxes.into_iter().collect::<BTreeSet<_>>();
    let reset_epoch = !matches!(observation.expectation, CacheExpectation::ExpectedWarm);
    let mut next_state = if reset_epoch {
        MarkerState::default()
    } else {
        state.clone()
    };
    let stale = if reset_epoch {
        (!current_stale.is_empty())
            .then(|| StaleMarker::Consolidated(current_stale.iter().copied().collect()))
    } else {
        let new = current_stale
            .difference(&next_state.reported_stale_boxes)
            .copied()
            .collect::<Vec<_>>();
        (!new.is_empty()).then_some(StaleMarker::New(new))
    };
    next_state.reported_stale_boxes = current_stale;

    let context_size_tokens = size_marker_due(
        next_state.last_context_size_tokens,
        observation.current_context_tokens,
        observation.context_limit_tokens,
    );
    if let Some(tokens) = context_size_tokens {
        next_state.last_context_size_tokens = Some(tokens);
    }
    MarkerDecision {
        reset_epoch,
        stale,
        context_size_tokens,
        next_state,
    }
}

fn size_marker_due(last: Option<u64>, current: u64, limit: u64) -> Option<u64> {
    if limit == 0 || current.saturating_mul(10) <= limit.saturating_mul(3) {
        return None;
    }
    match last {
        None => Some(current),
        Some(last) if current.saturating_sub(last) >= CONTEXT_SIZE_MARKER_INTERVAL_TOKENS => {
            Some(current)
        }
        Some(_) => None,
    }
}

/// Final provider input-token dimensions used by cache diagnostics.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct InputTokens {
    pub total: u64,
    pub cached: u64,
}

impl InputTokens {
    pub const fn uncached(self) -> u64 {
        self.total.saturating_sub(self.cached)
    }

    pub fn cached_ratio(self) -> Option<f64> {
        (self.total > 0).then(|| self.cached as f64 / self.total as f64)
    }
}

/// Observed cache result after provider usage becomes available.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CacheOutcome {
    Hit,
    UnplannedMiss,
    Excluded,
    UsageUnavailable,
}

impl CacheOutcome {
    pub const fn label(self) -> &'static str {
        match self {
            Self::Hit => "hit",
            Self::UnplannedMiss => "unplanned_miss",
            Self::Excluded => "excluded",
            Self::UsageUnavailable => "usage_unavailable",
        }
    }
}

/// Classifies final provider cache usage without inventing unavailable data.
pub fn observe_cache(expectation: &CacheExpectation, tokens: Option<InputTokens>) -> CacheOutcome {
    let Some(tokens) = tokens.filter(|tokens| tokens.total > 0) else {
        return CacheOutcome::UsageUnavailable;
    };
    if tokens.cached > 0 {
        CacheOutcome::Hit
    } else if expectation.expects_cache_hit() {
        CacheOutcome::UnplannedMiss
    } else {
        CacheOutcome::Excluded
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn exact_prefix_and_material_determine_cache_expectation() {
        let previous = PreviousProjection {
            text: "stable",
            material_fingerprint: "material-a",
        };
        assert_eq!(
            classify(None, "stable", "material-a", "rewrite"),
            CacheExpectation::ColdStart
        );
        assert_eq!(
            classify(Some(previous), "stable plus", "material-a", "rewrite"),
            CacheExpectation::ExpectedWarm
        );
        assert_eq!(
            classify(Some(previous), "changed", "material-a", "dehydrated"),
            CacheExpectation::PlannedInvalidation {
                reason: "dehydrated".into()
            }
        );
        assert_eq!(
            classify(Some(previous), "stable", "material-b", ""),
            CacheExpectation::PlannedInvalidation {
                reason: "provider_material_changed".into()
            }
        );
    }

    #[test]
    fn warm_epochs_emit_only_new_stale_boxes() {
        let state = MarkerState {
            reported_stale_boxes: BTreeSet::from([2, 4]),
            last_context_size_tokens: None,
        };
        let decision = decide_markers(
            &state,
            MarkerObservation {
                stale_boxes: vec![2, 4, 7],
                current_context_tokens: 20,
                context_limit_tokens: 100,
                expectation: CacheExpectation::ExpectedWarm,
            },
        );
        assert!(!decision.reset_epoch);
        assert_eq!(decision.stale, Some(StaleMarker::New(vec![7])));
        assert_eq!(
            decision.stale.as_ref().unwrap().render(),
            "[new stale boxes: 7]"
        );
        assert_eq!(decision.context_size_tokens, None);
    }

    #[test]
    fn invalidation_consolidates_stale_state_and_resets_size_baseline() {
        let state = MarkerState {
            reported_stale_boxes: BTreeSet::from([1]),
            last_context_size_tokens: Some(75_000),
        };
        let decision = decide_markers(
            &state,
            MarkerObservation {
                stale_boxes: vec![1, 9],
                current_context_tokens: 40_000,
                context_limit_tokens: 100_000,
                expectation: CacheExpectation::PlannedInvalidation {
                    reason: "summarized".into(),
                },
            },
        );
        assert!(decision.reset_epoch);
        assert_eq!(decision.stale, Some(StaleMarker::Consolidated(vec![1, 9])));
        assert_eq!(decision.context_size_tokens, Some(40_000));
        assert_eq!(decision.next_state.last_context_size_tokens, Some(40_000));
    }

    #[test]
    fn size_markers_are_strictly_over_thirty_percent_and_ten_thousand_apart() {
        for (current, last, expected) in [
            (30_000, None, None),
            (30_001, None, Some(30_001)),
            (40_000, Some(30_001), None),
            (40_001, Some(30_001), Some(40_001)),
        ] {
            assert_eq!(size_marker_due(last, current, 100_000), expected);
        }
        assert_eq!(size_marker_due(None, 100, 0), None);
    }

    #[test]
    fn cache_outcomes_separate_health_from_expected_misses() {
        let warm = CacheExpectation::ExpectedWarm;
        let cold = CacheExpectation::ColdStart;
        assert_eq!(
            observe_cache(
                &warm,
                Some(InputTokens {
                    total: 100,
                    cached: 80
                })
            ),
            CacheOutcome::Hit
        );
        assert_eq!(
            observe_cache(
                &warm,
                Some(InputTokens {
                    total: 100,
                    cached: 0
                })
            ),
            CacheOutcome::UnplannedMiss
        );
        assert_eq!(
            observe_cache(
                &cold,
                Some(InputTokens {
                    total: 100,
                    cached: 0
                })
            ),
            CacheOutcome::Excluded
        );
        assert_eq!(observe_cache(&warm, None), CacheOutcome::UsageUnavailable);
    }

    #[test]
    fn token_dimensions_saturate_and_report_ratios() {
        let tokens = InputTokens {
            total: 100,
            cached: 125,
        };
        assert_eq!(tokens.uncached(), 0);
        assert_eq!(tokens.cached_ratio(), Some(1.25));
        assert_eq!(
            InputTokens {
                total: 0,
                cached: 0
            }
            .cached_ratio(),
            None
        );
    }
}