Cache Breakpoint Planner — Rust source

//! Cache Breakpoint Planner — find the blocks a set of prompts share and
//! place cache breakpoints where they pay.
//!
//! Language: Rust (edition 2021, standard library only)
//! Port of src/lib/cacheBreakpointPlanner.ts (the canonical TypeScript
//! implementation). javascript.js in this set carries the same port.
//! Tool page: https://dev.cosmolabs.org/tools/cache-breakpoint-planner



/// Cached reads bill at ~0.1x — the saving on the cached share is ~90%.
pub const CACHE_READ_DISCOUNT: f64 = 0.1;

/// One prompt session: an id plus its ordered blocks.
pub struct PromptSession<'a> {
    pub id: &'a str,
    pub blocks: &'a [&'a str],
}

/// Place the cache breakpoint AFTER this block index (0-based).
#[derive(Debug, Clone, PartialEq)]
pub struct Breakpoint {
    pub after_block: i64,
    pub label: &'static str,
    pub reason: String,
    pub cached_tokens: i64,
}

#[derive(Debug, Clone, PartialEq)]
pub struct PerSessionRow {
    pub id: String,
    pub total_tokens: i64,
    pub unique_tokens: i64,
    pub cached_ratio: f64,
}

/// The full plan: shared blocks, breakpoints, per-session rows, savings.
#[derive(Debug, Clone, Default)]
pub struct BreakpointPlan {
    pub prefix_blocks: Vec<String>,
    pub prefix_tokens: i64,
    pub suffix_blocks: Vec<String>,
    pub suffix_tokens: i64,
    pub breakpoints: Vec<Breakpoint>,
    pub per_session: Vec<PerSessionRow>,
    /// Estimated cost saving across the sessions vs no caching (0-1).
    pub estimated_savings: f64,
    pub warnings: Vec<String>,
}

/// The `type: 'prose'` path of the tokenEstimator, inlined: every non-empty
/// line costs max(1, round(len / 4)) tokens (UTF-16 code units ≈ bytes here
/// for ASCII display snippets); empty text is 0.
fn tok(text: &str) -> i64 {
    if text.is_empty() {
        return 0;
    }
    let mut tokens: i64 = 0;
    for line in text.split('\n') {
        if !line.is_empty() {
            tokens += std::cmp::max(1, (line.len() as f64 / 4.0).round() as i64);
        }
    }
    tokens
}

/// Plan cache breakpoints for a set of prompt sessions: find the common
/// leading/trailing blocks across every session and place breakpoints where
/// the cache pays. `sessions` with a non-array `blocks` are ignored at the
/// caller's boundary (TS filters with Array.isArray); pass only valid ones.
pub fn plan_breakpoints(sessions: &[PromptSession<'_>]) -> BreakpointPlan {
    let mut warnings: Vec<String> = Vec::new();

    if sessions.is_empty() {
        return BreakpointPlan {
            warnings: vec!["No sessions given — paste at least two prompts to compare.".into()],
            ..Default::default()
        };
    }
    if sessions.len() == 1 {
        warnings.push("Only one session — a prefix needs at least two prompts to detect.".into());
    }

    // Common leading blocks by position.
    let shortest = sessions.iter().map(|s| s.blocks.len()).min().unwrap_or(0);
    let mut prefix_end = 0usize;
    while prefix_end < shortest
        && sessions
            .iter()
            .all(|s| s.blocks[prefix_end] == sessions[0].blocks[prefix_end])
    {
        prefix_end += 1;
    }

    // Common trailing blocks, matched from each session's own tail, never
    // overlapping the prefix.
    let mut suffix_len = 0usize;
    while suffix_len < shortest - prefix_end
        && sessions.iter().all(|s| {
            let own = s.blocks.len() - 1 - suffix_len;
            let first = sessions[0].blocks.len() - 1 - suffix_len;
            s.blocks[own] == sessions[0].blocks[first]
        })
    {
        suffix_len += 1;
    }

    let prefix_blocks: Vec<String> = sessions[0].blocks[..prefix_end]
        .iter()
        .map(|s| s.to_string())
        .collect();
    let suffix_blocks: Vec<String> = if suffix_len > 0 {
        let n = sessions[0].blocks.len();
        sessions[0].blocks[n - suffix_len..]
            .iter()
            .map(|s| s.to_string())
            .collect()
    } else {
        Vec::new()
    };
    let prefix_tokens = tok(&prefix_blocks.join("\n"));
    let suffix_tokens = tok(&suffix_blocks.join("\n"));

    let mut breakpoints: Vec<Breakpoint> = Vec::new();
    if !prefix_blocks.is_empty() {
        breakpoints.push(Breakpoint {
            after_block: prefix_end as i64 - 1,
            label: "after the shared prefix",
            reason: format!(
                "{} block(s) identical across every session — cache once, hit on every request.",
                prefix_blocks.len()
            ),
            cached_tokens: prefix_tokens,
        });
    }
    if suffix_len > 0 {
        breakpoints.push(Breakpoint {
            after_block: -1, // terminal: the shared tail sits at the end of each request
            label: "shared tail",
            reason: format!(
                "{} trailing block(s) also identical — extend the cache segment or accept the re-read.",
                suffix_len
            ),
            cached_tokens: suffix_tokens,
        });
    }
    if breakpoints.is_empty() {
        warnings.push(
            "No shared leading or trailing blocks — nothing to cache across these sessions.".into(),
        );
    }

    let per_session: Vec<PerSessionRow> = sessions
        .iter()
        .map(|s| {
            let total_tokens = tok(&s.blocks.join("\n"));
            let unique_tokens = std::cmp::max(total_tokens - prefix_tokens - suffix_tokens, 0);
            let cached_ratio = if total_tokens > 0 {
                ((prefix_tokens + suffix_tokens) as f64 / total_tokens as f64).min(1.0)
            } else {
                0.0
            };
            PerSessionRow {
                id: s.id.to_string(),
                total_tokens,
                unique_tokens,
                cached_ratio,
            }
        })
        .collect();

    let avg_total = per_session.iter().map(|p| p.total_tokens).sum::<i64>() as f64
        / per_session.len() as f64;
    let cached_share = if avg_total > 0.0 {
        ((prefix_tokens + suffix_tokens) as f64 / avg_total).min(1.0)
    } else {
        0.0
    };
    let estimated_savings = cached_share * (1.0 - CACHE_READ_DISCOUNT);

    BreakpointPlan {
        prefix_blocks,
        prefix_tokens,
        suffix_blocks,
        suffix_tokens,
        breakpoints,
        per_session,
        estimated_savings,
        warnings,
    }
}