ai_memory/

sizes.rs

// Copyright 2026 AlphaOne LLC
// SPDX-License-Identifier: Apache-2.0

//! v0.6.4-005 — Static schema-size table.
//!
//! Computes the per-tool BPE token cost of every MCP tool registered by
//! `crate::mcp::tool_definitions`, using the `tiktoken-rs` `cl100k_base`
//! tokenizer (the same BPE Claude/GPT use for context-window accounting,
//! and the same one v0.6.3.1 P6/R1 already wires for `budget_tokens`).
//!
//! The table is computed lazily on first access and cached behind a
//! `OnceLock`. The cost of the first call is one full pass over every
//! tool schema (~7 ms on Apple M2) followed by cache hits forever after.
//!
//! ## Why lazy and not literally compile-time
//!
//! The "build-time" framing in the v0.6.4 issue spec referred to the
//! desire that operators be able to query the table without running
//! the full MCP `register_tools()` dance — the runtime cache satisfies
//! that constraint. A real build-time approach would need either a
//! proc-macro or a `build.rs` that re-parsed the JSON-emitting Rust
//! source, both of which trade simplicity for marginal warm-cache
//! performance that nobody is paying for here. The lazy approach also
//! keeps the BPE table out of `cargo bench` cold paths — every place
//! that *doesn't* run `doctor --tokens` pays exactly nothing.
//!
//! ## CI gate
//!
//! `tool_sizes_under_ci_gate()` returns the largest single tool cost.
//! The unit test `no_tool_exceeds_1500_tokens` enforces the v0.6.4-005
//! acceptance gate that no individual tool definition exceeds 1500
//! tokens. The number is high enough to permit growth on the more
//! schema-heavy KG/governance tools and low enough that doubling a
//! tool's schema by accident lands in CI red.

use std::sync::OnceLock;

use serde_json::Value;
use tiktoken_rs::CoreBPE;

/// Single-tool cost report. The `total` is what counts against the
/// per-request prefix; the `name_tokens` and `schema_tokens` split is
/// useful for the doctor's diagnostic output.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolSize {
    pub name: String,
    pub schema_tokens: usize,
    pub name_tokens: usize,
    pub total_tokens: usize,
}

/// pm-v3.1 PR8 (issue #1174) — shared one-shot cache for the verbose
/// table. Replaces the prior pair of function-local `static TABLE`
/// declarations (one in `tool_sizes`, one in `trimmed_tool_sizes`)
/// with a single module-level slot per shape. The cache is computed
/// lazily on first use and reused forever after.
static VERBOSE_TABLE: OnceLock<Vec<ToolSize>> = OnceLock::new();

/// pm-v3.1 PR8 (issue #1174) — shared cache for the trimmed (wire-
/// form) table, matching `VERBOSE_TABLE` above.
static TRIMMED_TABLE: OnceLock<Vec<ToolSize>> = OnceLock::new();

/// Runtime-computed table of every tool's tokenized schema cost
/// **at the verbose ceiling** — every optional param, every default,
/// every per-property description. This is the upper bound a host
/// can ever pay (only reachable via
/// `memory_capabilities { verbose=true, family=…, include_schema=true }`
/// since v0.7 C4).
///
/// Returns a static slice on every call after the first invocation
/// (which performs the one-time BPE pass).
pub fn tool_sizes() -> &'static [ToolSize] {
    VERBOSE_TABLE
        .get_or_init(|| compute_table(false))
        .as_slice()
}

/// v0.7 C4 + #859 — runtime-computed table of every tool's tokenized
/// schema cost **as actually shipped on `tools/list`**. Per-property
/// `description` prose is stripped, the top-level tool `description`
/// is compacted to the first sentence, the `docs` field is dropped,
/// but every property entry survives so MCP clients can discover the
/// call surface (per [`crate::mcp::trim_optional_params`] +
/// [`crate::mcp::strip_docs_from_tools`] + `wire_compact_descriptions`).
/// This is what an MCP host pays per request on the default code path.
///
/// **Wire-form invariant.** This table is computed by feeding the
/// output of [`crate::mcp::tool_definitions_for_profile`] (full
/// profile) into the cl100k_base tokenizer; the budget gate at
/// `tests/c2_tool_docs_field.rs::c2_tools_list_token_budget_is_under_post_859_ceiling`
/// pins the sum at ≤ 5000 cl100k tokens (post-#859 floor; was 3500
/// pre-#859 when the trim hid optional property keys entirely).
pub fn trimmed_tool_sizes() -> &'static [ToolSize] {
    TRIMMED_TABLE.get_or_init(|| compute_table(true)).as_slice()
}

/// Highest-cost tool in the verbose table. Used by the CI gate.
pub fn tool_sizes_under_ci_gate() -> usize {
    tool_sizes()
        .iter()
        .map(|t| t.total_tokens)
        .max()
        .unwrap_or(0)
}

/// Sum of every tool's `total_tokens` (verbose schema) — the
/// worst-case prefix cost on a `verbose=true` opt-in harness with
/// `--profile full`. The actually-paid cost on the default code path
/// is reported by [`trimmed_full_profile_total_tokens`].
pub fn full_profile_total_tokens() -> usize {
    tool_sizes().iter().map(|t| t.total_tokens).sum()
}

/// v0.7 C4 — sum of every tool's `total_tokens` after the C4 trim
/// (optionals hidden). This is the bare `tools/list` payload cost
/// under `--profile full`.
pub fn trimmed_full_profile_total_tokens() -> usize {
    trimmed_tool_sizes().iter().map(|t| t.total_tokens).sum()
}

/// Lookup a single tool by name in the verbose table. `O(n)` but
/// `n ≤ 57` (v0.7.0 L1-5 added 5 skill tools).
pub fn tool_size(name: &str) -> Option<&'static ToolSize> {
    tool_sizes().iter().find(|t| t.name == name)
}

fn compute_table(trimmed: bool) -> Vec<ToolSize> {
    let bpe = bpe();
    // #859 — to keep the budget model in lockstep with the actually-
    // shipped wire payload, delegate to `tool_definitions_for_profile`
    // for the trimmed case (which now performs the full wire shape:
    // properties preserved, per-property prose stripped, top-level
    // description compacted). For the verbose case we measure the raw
    // `tool_definitions()` table as it would appear on the
    // `memory_capabilities { verbose=true }` opt-in path.
    let defs = if trimmed {
        crate::mcp::tool_definitions_for_profile(&crate::profile::Profile::full())
    } else {
        crate::mcp::tool_definitions()
    };
    let tools = defs
        .get("tools")
        .and_then(Value::as_array)
        .cloned()
        .unwrap_or_default();

    tools
        .into_iter()
        .filter_map(|tool| size_one_tool(&bpe, &tool))
        .collect()
}

fn size_one_tool(bpe: &CoreBPE, tool: &Value) -> Option<ToolSize> {
    let name = tool.get("name").and_then(Value::as_str)?.to_string();
    // The cost the host pays is the serialized JSON of the entire tool
    // object — name + description + inputSchema. We use the canonical
    // serde_json serialization (no pretty-printing) because that is
    // what every MCP host transmits over stdio.
    let schema_json = serde_json::to_string(tool).ok()?;
    let schema_tokens = bpe.encode_with_special_tokens(&schema_json).len();
    let name_tokens = bpe.encode_with_special_tokens(&name).len();
    Some(ToolSize {
        name,
        schema_tokens,
        name_tokens,
        total_tokens: schema_tokens,
    })
}

fn bpe() -> CoreBPE {
    // We construct a fresh BPE on each compute_table call (only ever
    // called once) because `cl100k_base` returns an owned `CoreBPE`
    // and stashing it forever in a static would leak ~1.7 MB for a
    // table that only gets walked at startup. Cheap to throw away.
    tiktoken_rs::cl100k_base().expect("cl100k_base BPE table embedded in tiktoken-rs")
}

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

    /// CI gate per v0.6.4-005 acceptance criteria. If any tool's schema
    /// crosses 1500 tokens, the whole build fails. The number is roughly
    /// 2.5× today's largest tool (memory_store at ~620 tokens) so we have
    /// runway, but not so much runway that a 3× regression slips through.
    #[test]
    fn no_tool_exceeds_1500_tokens() {
        let max = tool_sizes_under_ci_gate();
        assert!(
            max <= 1500,
            "v0.6.4-005 CI gate: largest tool schema is {max} tokens (limit: 1500). \
             Inspect `cargo run -- doctor --tokens --raw-table` to find the offender."
        );
    }

    /// Sanity: the table must be populated. Catches accidental empty
    /// `tool_definitions()` regressions that would silently hide other
    /// failures.
    #[test]
    fn table_entry_count_matches_full_profile() {
        // v0.7.0 refactor PR-2 (#793) — tool-count SSOT. Anchor the
        // assertion on `Profile::full().expected_tool_count()` (derived
        // from the per-Family `tool_names` slices) rather than a
        // hardcoded literal, so adding a new MCP tool touches ONE site
        // (the family slice) instead of N assertions across the
        // codebase.
        let n = tool_sizes().len();
        let expected = crate::profile::Profile::full().expected_tool_count();
        assert_eq!(
            n, expected,
            "tool_sizes() must hold exactly {expected} tools (the full-profile \
             SSOT count); got {n}. If these diverge, a tool was added to one \
             surface but not the other."
        );
    }

    /// Every tool should have non-zero name + schema costs. Zero would
    /// mean either an empty schema or a tokenizer wiring break.
    #[test]
    fn every_tool_has_nonzero_cost() {
        for t in tool_sizes() {
            assert!(t.schema_tokens > 0, "tool {} schema_tokens = 0", t.name);
            assert!(t.name_tokens > 0, "tool {} name_tokens = 0", t.name);
        }
    }

    /// Full-profile total cost — measured against `cl100k_base` (the
    /// tokenizer Claude / GPT actually use for input accounting).
    ///
    /// **Truthfulness note (v0.6.4-005, 2026-05-04):** the v0.6.4 RFC
    /// claimed ~25,800 tokens for the full surface, derived from "~600
    /// tokens/tool × 43" measured against MiniLM. MiniLM is a sentence-
    /// embedding vocabulary (~30K tokens) that systematically over-counts
    /// JSON by ~4× vs. `cl100k_base` (100K-token chat-completion BPE).
    /// The actual measured cost in `cl100k_base` is ~6,000 tokens for
    /// the full surface — still material, still worth the v0.6.4 ship,
    /// but the public claims need a 4× downward correction (tracked in
    /// v0.6.4-014 + v0.6.4-015 docs work).
    ///
    /// **v0.7 C2 update (2026-05-06):** the canonical
    /// `tool_definitions()` now carries an additional per-tool `docs`
    /// field (long-form description + examples) that the bare
    /// `tools/list` payload strips before transmission. The numbers
    /// in this table reflect the **source of truth** (verbose +
    /// short), not the wire payload. The bare-wire C5 budget is
    /// pinned separately at ≤ 3500 tokens by
    /// `tests/c2_tool_docs_field.rs::c2_tools_list_token_budget_is_under_post_859_ceiling`.
    /// The savings *percentage* from `core` is unchanged; the
    /// always-on payload is now ~85% smaller than the source.
    #[test]
    fn full_profile_total_in_honest_measured_range() {
        let total = full_profile_total_tokens();
        // **v0.7.0 #829 update.** Prior bound was 5K..=16K to soak the
        // multi-paragraph `docs` prose that every tool carried. After
        // the #829 trim every `docs` field is a single condensed
        // sentence with issue refs + tier annotations preserved.
        //
        // **v0.7.0 #987 update.** D1.6 collapsed `tool_definitions()`
        // to iterate over per-tool `McpTool` impls; the schemars-derived
        // `inputSchema` carries additional metadata the legacy
        // hand-coded macro didn't: `additionalProperties: false`,
        // `default: null` on optional fields, `$schema` reference,
        // `title`, request-struct-level `description`. Measured total
        // settles at ~15K. Hard ceiling at 17K to leave 2K headroom for
        // the next field-addition without re-bumping. Floor stays at 5K
        // to catch a wiring break that drops the catalog entirely.
        assert!(
            (5_000..=17_000).contains(&total),
            "full-profile total {total} tokens is outside the measured \
             cl100k_base range (5K-17K, post-#987 D1.6). If the schema \
             grew intentionally, update `tests/token_budget_guard.rs::\
             VERBOSE_FULL_PROFILE_CEILING_TOKENS` AND this bound together."
        );
    }

    /// Lookup by name should resolve a known tool.
    #[test]
    fn tool_size_resolves_memory_store() {
        let t = tool_size("memory_store").expect("memory_store should exist");
        assert!(t.total_tokens > 0);
        assert!(t.total_tokens < 1500);
    }

    /// Lookup of a nonexistent tool should return None, not panic.
    #[test]
    fn tool_size_returns_none_for_unknown() {
        assert!(tool_size("memory_does_not_exist_42").is_none());
    }

    /// v0.7 C4 + #859 acceptance gate: the trimmed `tools/list`
    /// payload (the shape an MCP host actually receives by default)
    /// must be materially smaller than the verbose ceiling AND must
    /// stay under the post-#859 5000-token wire-form budget.
    ///
    /// **History.** Pre-#859 baseline: trimmed ≈ 3456 tokens,
    /// verbose ≈ 7416 tokens (~53% saved). The trim dropped every
    /// optional property entry from the wire, hiding the call
    /// surface from MCP clients. #859 (v0.7.0 fix) restored every
    /// property entry on the wire (keeping per-property `description`
    /// prose stripped + the top-level tool description compacted to
    /// the first sentence) so NHI agents can discover what knobs
    /// exist. Post-#859: trimmed ≈ 4500-4700 tokens, verbose ≈ 9500.
    ///
    /// The savings now sit at ~50% (down from ~53%) because the
    /// property metadata that pre-fix lived only in the verbose
    /// catalog now also appears on the wire. The 5000 ceiling pins
    /// the post-#859 floor with ~300 tokens of headroom for future
    /// tool additions; the 25% lower bound on `saved_pct` keeps the
    /// trim itself honest (a regression that re-bloated the wire
    /// path with docs / per-property prose would still trip).
    #[test]
    fn trimmed_full_profile_total_under_post_859_ceiling() {
        let trimmed = trimmed_full_profile_total_tokens();
        let verbose = full_profile_total_tokens();
        assert!(
            trimmed < verbose,
            "trimmed total ({trimmed}) must be strictly smaller than verbose ({verbose})"
        );
        let saved_pct = (verbose - trimmed) as f64 / verbose as f64 * 100.0;
        assert!(
            saved_pct >= 25.0,
            "trim should save >=25% of full-profile tokens; got {saved_pct:.1}% \
             (verbose={verbose}, trimmed={trimmed}). Audit `strip_docs_from_tools` and \
             `wire_compact_descriptions` — if those broke the trim itself regressed."
        );
        assert!(
            trimmed <= 11_000,
            "post-#987 D1.6 trimmed full-profile total {trimmed} > 11000-token ceiling. \
             The #859 fix preserves every property entry on the wire. The post-D1.6 \
             ceiling rose from 5000 to 11000 because schemars-derived schemas carry \
             additional metadata (`additionalProperties: false`, `default: null`, \
             `$schema`, `title`, request-struct `description`) that the legacy \
             hand-coded `tool_definitions()` macro did not emit. If trimmed grew \
             beyond 11000, audit per-property `description` prose (must be stripped \
             by `strip_docs_from_tools`) and consider routing the new tool to \
             `family=power` instead of the always-on core."
        );
    }

    /// Trim must shrink at least one optional from at least one tool;
    /// otherwise the wiring is broken (e.g. `trim_optional_params` got
    /// short-circuited or the keep-list went global).
    #[test]
    fn trimmed_table_strictly_smaller_per_tool_where_optionals_existed() {
        let verbose: std::collections::HashMap<&str, usize> = tool_sizes()
            .iter()
            .map(|t| (t.name.as_str(), t.total_tokens))
            .collect();
        let mut at_least_one_smaller = false;
        for trimmed_tool in trimmed_tool_sizes() {
            let v = verbose
                .get(trimmed_tool.name.as_str())
                .copied()
                .unwrap_or(0);
            assert!(
                trimmed_tool.total_tokens <= v,
                "{} grew under trim ({} > {})",
                trimmed_tool.name,
                trimmed_tool.total_tokens,
                v
            );
            if trimmed_tool.total_tokens < v {
                at_least_one_smaller = true;
            }
        }
        assert!(
            at_least_one_smaller,
            "trim should shrink at least one tool; none did — wiring is broken"
        );
    }
}