scryer-mcp 0.3.0

//! Response budgets that don't force a second call (Plan 19 Phase 2).
//!
//! The byte-oriented limiter in `telemetry.rs` is the last resort: it cuts wherever the budget
//! ends and leaves no way to ask for the rest. List-shaped tools instead drop whole items from
//! the end of the page until the response fits, then say how to fetch the remainder:
//! `complete: false`, `next_offset` and a `see_also` naming the call.

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

use crate::session::ResponseFormat;
use crate::telemetry::MAX_PAYLOAD_TOKENS;

/// The budget a tool gets when neither the call nor the session sets one, sized to what the
/// tool normally needs so the common answer arrives in one call.
/// `items` is how many things the call asks about (the lines of a batched scope lookup).
pub fn default_budget(tool: &str, items: usize) -> usize {
    match tool {
        "get_file_outline" => 2_000,
        "find_references" | "search_symbols" => 1_500,
        "get_enclosing_scope" => 600.max(100 * items.min(25)),
        "inspect_symbol" => MAX_PAYLOAD_TOKENS.max(500 * items.min(5)),
        _ => MAX_PAYLOAD_TOKENS,
    }
}

/// How to get the rest of a cut list, flattened into each list-shaped result.
#[derive(Debug, Clone, Default, Serialize, Deserialize, JsonSchema)]
pub struct Continuation {
    /// `false` when the list was cut to fit the budget. Absent otherwise.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub complete: Option<bool>,
    /// `offset` that returns the next item after the last one shown.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub next_offset: Option<usize>,
    /// The exact next call.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub see_also: Vec<String>,
    /// Index state for the text header (`index: 1,375 files, ...`); never serialised.
    #[serde(skip)]
    pub index_state: Option<String>,
}

/// `skip_serializing_if` for counts that are 0.
pub(crate) fn is_zero(n: &usize) -> bool {
    *n == 0
}

/// A result whose main field is a page of items that can be shortened from the end.
pub trait Paged: Serialize + Clone {
    const TOOL: &'static str;

    /// Items on this page.
    fn item_count(&self) -> usize;
    /// Offset of the first item of the page.
    fn page_offset(&self) -> usize;
    /// Keep only the first `kept` items and record that more follow.
    fn keep_first(&mut self, kept: usize);
    fn continuation_mut(&mut self) -> &mut Continuation;
    /// A header line, then one line per item.
    fn render_text(&self) -> String;

    /// The answer as it will be sent, which is what the budget is measured against.
    fn render(&self, format: ResponseFormat) -> String {
        match format {
            ResponseFormat::Json => serde_json::to_string(self).unwrap_or_default(),
            ResponseFormat::Text => self.render_text(),
        }
    }
}

/// Shorten `out` until it fits `budget` tokens, keeping at least one item. Does nothing when
/// `no_truncate` is set, the budget is 0 or the result already fits.
pub fn fit_list_to_budget<P: Paged>(
    out: &mut P,
    budget: Option<usize>,
    no_truncate: bool,
    format: ResponseFormat,
    count_tokens: impl Fn(&str) -> usize,
) {
    let budget = budget.unwrap_or(MAX_PAYLOAD_TOKENS);
    let size = |r: &P| count_tokens(&r.render(format));
    if no_truncate || budget == 0 || size(out) <= budget {
        return;
    }
    let full = out.clone();
    let len = full.item_count();
    if len < 2 {
        return;
    }
    let offset = full.page_offset();
    let base = full.clone().continuation_mut().clone();

    let apply = |out: &mut P, kept: usize| {
        *out = full.clone();
        out.keep_first(kept);
        let next = offset + kept;
        let hint = format!("{}(same arguments, offset: {next}) for the rest", P::TOOL);
        let c = out.continuation_mut();
        *c = base.clone();
        c.complete = Some(false);
        c.next_offset = Some(next);
        c.see_also.push(hint);
    };

    // Largest prefix that fits; at least one item so every call makes progress.
    let (mut lo, mut hi) = (1, len - 1);
    while lo < hi {
        let mid = (lo + hi).div_ceil(2);
        apply(out, mid);
        if size(out) <= budget {
            lo = mid;
        } else {
            hi = mid - 1;
        }
    }
    apply(out, lo);
}

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

    #[derive(Clone, Serialize)]
    struct Page {
        items: Vec<String>,
        offset: usize,
        returned: usize,
        has_more: bool,
        #[serde(flatten)]
        next: Continuation,
    }

    impl Paged for Page {
        const TOOL: &'static str = "demo";
        fn item_count(&self) -> usize {
            self.items.len()
        }
        fn page_offset(&self) -> usize {
            self.offset
        }
        fn keep_first(&mut self, kept: usize) {
            self.items.truncate(kept);
            self.returned = kept;
            self.has_more = true;
        }
        fn continuation_mut(&mut self) -> &mut Continuation {
            &mut self.next
        }
        fn render_text(&self) -> String {
            self.items.join("\n")
        }
    }

    // One token per character is enough to exercise the search.
    fn chars(s: &str) -> usize {
        s.len()
    }

    fn page(n: usize) -> Page {
        Page {
            items: (0..n).map(|i| format!("item{i:03}")).collect(),
            offset: 10,
            returned: n,
            has_more: false,
            next: Continuation::default(),
        }
    }

    #[test]
    fn a_page_that_fits_is_left_alone() {
        let mut p = page(3);
        fit_list_to_budget(&mut p, Some(10_000), false, ResponseFormat::Json, chars);
        assert_eq!(p.items.len(), 3);
        assert!(p.next.complete.is_none() && p.next.see_also.is_empty());
    }

    #[test]
    fn an_oversized_page_keeps_the_largest_prefix_and_names_the_next_call() {
        let mut p = page(200);
        fit_list_to_budget(&mut p, Some(600), false, ResponseFormat::Json, chars);
        let json = serde_json::to_string(&p).unwrap();
        assert!(json.len() <= 600, "{}", json.len());
        let kept = p.items.len();
        assert!((1..200).contains(&kept));
        assert_eq!(p.items[kept - 1], format!("item{:03}", kept - 1));
        assert_eq!(p.returned, kept);
        assert!(p.has_more);
        assert_eq!(p.next.complete, Some(false));
        assert_eq!(p.next.next_offset, Some(10 + kept));
        assert!(p.next.see_also[0].contains(&format!("offset: {}", 10 + kept)));
    }

    #[test]
    fn no_truncate_and_a_zero_budget_disable_fitting() {
        let mut p = page(200);
        fit_list_to_budget(&mut p, Some(100), true, ResponseFormat::Json, chars);
        fit_list_to_budget(&mut p, Some(0), false, ResponseFormat::Json, chars);
        assert_eq!(p.items.len(), 200);
    }
}