supercode/

tokens.rs

//! The shared token estimator (SPEC.md C9). The repo has no tokenizer — only
//! provider-reported `Usage` (`provider.rs:153-164`) and the turn/output
//! counters `agent.rs` already tracks. Every other UX figure (the banner,
//! `/status`, `/tokens`, `show-reductions`, notices) derives from one
//! documented heuristic here, and is always printed with a `~` prefix so it
//! reads as an estimate, never a measurement. Acceptance criteria never
//! assert against these numbers directly — ACs assert exact byte counts
//! (`SessionInfo::full_bytes`/`view_bytes`, C9/D15) and only *derive* a
//! token/dollar figure from them for the test log.
//!
//! Provider-reported figures (`Usage.completion_tokens`, B7's optional
//! `cached_tokens`, `agent.total_output_tokens()`) are real counts and are
//! never routed through this module — they print un-tilded.

use crate::message::ChatMessage;
use crate::provider::ToolSchema;
use crate::reduce::format_commas;

/// Deterministic token estimate: `ceil(utf8_bytes / 4)`. Documented
/// heuristic — all UX figures derived from it are printed with a `~`
/// prefix (see [`fmt_approx_tokens`]). Never used in acceptance criteria
/// (ACs assert exact byte counts).
pub fn estimate_tokens(s: &str) -> u64 {
    let bytes = s.len() as u64;
    bytes.div_ceil(4)
}

/// Estimate for a serialized message list: the wire form (`serde_json`)
/// used to build the provider request, not any in-memory/debug
/// representation.
pub fn estimate_view_tokens(msgs: &[ChatMessage]) -> u64 {
    let mut total = 0u64;
    for m in msgs {
        total += estimate_message_tokens(m);
    }
    total
}

/// PARITY-18 D2 — per-message estimate used by [`estimate_view_tokens`].
/// `ChatMessage`'s hand-written `Serialize` impl never actually fails in
/// practice (every field is a plain string or an already-validated
/// `serde_json::Value`), but `serde_json::to_string` still returns a
/// `Result` — and the previous `.unwrap_or_default()` here silently turned
/// any failure into `""`, i.e. **0 tokens counted for a message the
/// provider request still carries**, which is exactly the kind of gap a
/// context guard must never have. Falls back to the `Debug` rendering's
/// byte length instead: `Debug` includes every field verbatim (plus
/// `metadata`, which the wire form omits), so it is never a SMALLER
/// estimate than the wire form would have produced — conservative, never
/// zero, on the one path that can't normally be exercised.
fn estimate_message_tokens(m: &ChatMessage) -> u64 {
    match serde_json::to_string(m) {
        Ok(wire) => estimate_tokens(&wire),
        Err(_) => estimate_tokens(&format!("{m:?}")),
    }
}

/// PARITY-18 D2 — conservative safety margin folded into the context-guard
/// boundary only (never into the plain `~`-prefixed UX estimates
/// themselves, which stay the documented `ceil(bytes/4)` heuristic
/// unmodified). `ceil(utf8_bytes/4)` under-counts real tokenizer output on
/// CJK text, base64/binary-ish blobs, and dense code — informally by 25%+
/// against common tokenizers for those corpora, since multi-byte UTF-8
/// sequences and non-whitespace-delimited runs pack more real tokens per
/// byte than the heuristic assumes. 25% is a round number comfortably above
/// that observed skew: applying it can only make the guard MORE
/// conservative (refuse sooner), never let an over-context request through
/// that a real tokenizer would also have refused.
const GUARD_MARGIN_NUM: u64 = 5;
const GUARD_MARGIN_DEN: u64 = 4;

/// Apply the [`GUARD_MARGIN_NUM`]/[`GUARD_MARGIN_DEN`] context-guard safety
/// margin to a raw token estimate. Rounds up (`div_ceil`), never down —
/// the margin only ever pushes the boundary check to be more cautious.
pub fn with_guard_margin(tokens: u64) -> u64 {
    tokens
        .saturating_mul(GUARD_MARGIN_NUM)
        .div_ceil(GUARD_MARGIN_DEN)
}

/// PARITY-18 — headroom reserved for the model's own completion, folded
/// into the context-guard boundary alongside [`with_guard_margin`]. Neither
/// [`estimate_view_tokens`] nor [`estimate_request_tokens`] counts anything
/// for the reply the model is about to generate — this is a flat token
/// budget carved out of the model's context window for it, since the
/// completion shares the same window as the request on every provider this
/// crate targets.
///
/// PARITY-18 v3 NOTE-AND-DECIDE (NF6, owner-recorded, not fixed here): this
/// is a FLAT reserve — it does not read `Config::max_tokens`
/// (`crates/core/src/config.rs`, user-settable via `--max-tokens`, wired
/// into the actual provider request at `provider.rs`'s
/// `ChatRequest::max_tokens`). A user who passes `--max-tokens` greater than
/// 16,384 can still pass this guard (`projected + 16_384 <= context_limit`)
/// and then draw a provider-side `input_tokens + max_tokens > context_window`
/// rejection the guard never anticipated — i.e. the guard's margin can be
/// smaller than what the user actually asked the provider to reserve for the
/// completion. Making the reserve `max(CONTEXT_RESPONSE_RESERVE_TOKENS,
/// config.max_tokens)` would close this, but `context_guard` doesn't
/// currently receive `Config` at all (only `messages`/`tools`/
/// `context_limit`) — threading it through is a small but real signature
/// change touching every call site (`resume_cmd`, `Agent::run_loop`, and
/// this pass's new `reduce_to_fit` `fits` closures) that's out of scope for
/// this pass's reducer/guard boundary fix. Recorded for the owner.
pub const CONTEXT_RESPONSE_RESERVE_TOKENS: u64 = 16_384;

/// PARITY-18 D1 — the full wire-request token estimate: every message in
/// `messages` (callers pass what [`crate::Agent::build_request_messages`]
/// actually returns, which already carries the system prompt at index 0)
/// PLUS the serialized `tools` schema array, which is a real part of the
/// request body ([`crate::provider::build_request_body`]) but — before
/// PARITY-18's re-fix — was never counted by the preflight guard at all.
/// A session whose messages alone fit comfortably could still carry a fat
/// builtin/MCP tool-schema array that blows the real wire request; this is
/// the fix.
pub fn estimate_request_tokens(messages: &[ChatMessage], tools: &[ToolSchema]) -> u64 {
    let tools_wire = serde_json::to_string(tools).unwrap_or_default();
    estimate_view_tokens(messages).saturating_add(estimate_tokens(&tools_wire))
}

/// PARITY-18 D1/D2/D4 — the single context-guard decision, shared by every
/// call site that must decide whether a request is safe to send: the CLI's
/// `resume_cmd` preflight check AND `Agent::run_loop`'s per-send check
/// (D4 — the guard is a session invariant, not a one-shot preflight, so
/// turn 2+ and `/expand all` are covered too). Because both call through
/// this one function, a request can never pass one gate and fail the
/// other — there is only one formula.
///
/// `fits` is true iff the [`with_guard_margin`]-adjusted
/// [`estimate_request_tokens`] estimate, plus the
/// [`CONTEXT_RESPONSE_RESERVE_TOKENS`] completion reserve, is still within
/// `context_limit` — i.e. gates on the reduce TARGET
/// (`context_limit - CONTEXT_RESPONSE_RESERVE_TOKENS`), not the raw limit,
/// closing the "blind band between reduce target and pass/fail boundary"
/// gap. Returns the margin-adjusted projected total either way so callers
/// can report it (dev/02) regardless of verdict.
pub fn context_guard(
    messages: &[ChatMessage],
    tools: &[ToolSchema],
    context_limit: u64,
) -> (bool, u64) {
    let raw = estimate_request_tokens(messages, tools);
    let projected = with_guard_margin(raw);
    let fits = projected.saturating_add(CONTEXT_RESPONSE_RESERVE_TOKENS) <= context_limit;
    (fits, projected)
}

/// Schema-token estimate for the B6 "tools" banner line: the estimate over
/// the serialized `full` [`ToolSchema`] list minus the estimate over the
/// serialized `advertised` list — i.e. the token cost of what's currently
/// deferred (hidden behind `tool_search`) rather than eagerly advertised.
/// Saturates to `0` rather than underflow if `advertised` somehow estimates
/// larger than `full` (e.g. formatting differences), since "negative
/// deferred tokens" has no meaning for the banner.
pub fn estimate_deferred_schema_tokens(full: &[ToolSchema], advertised: &[ToolSchema]) -> u64 {
    let full_tokens = estimate_tokens(&serde_json::to_string(full).unwrap_or_default());
    let advertised_tokens = estimate_tokens(&serde_json::to_string(advertised).unwrap_or_default());
    full_tokens.saturating_sub(advertised_tokens)
}

/// Render an estimated token count in the shared UX style: `~21,904 tok`
/// (tilde prefix + comma-grouped thousands, matching the stub-line comma
/// style in `reduce.rs`). Every figure that flows through
/// [`estimate_tokens`]/[`estimate_view_tokens`]/[`estimate_deferred_schema_tokens`]
/// should be rendered through this helper so the `~` discipline (D11) is
/// applied uniformly rather than ad hoc at each call site.
pub fn fmt_approx_tokens(n: u64) -> String {
    format!("~{} tok", format_commas(n as usize))
}

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

    #[test]
    fn empty_string_is_zero_tokens() {
        assert_eq!(estimate_tokens(""), 0);
    }

    #[test]
    fn four_byte_ascii_is_one_token() {
        assert_eq!(estimate_tokens("abcd"), 1);
    }

    #[test]
    fn ceil_behavior_rounds_up() {
        // 5 bytes / 4 = 1.25 -> ceil to 2.
        assert_eq!(estimate_tokens("abcde"), 2);
    }

    #[test]
    fn deterministic_same_input_same_output() {
        let s = "the quick brown fox jumps over the lazy dog";
        assert_eq!(estimate_tokens(s), estimate_tokens(s));
    }

    #[test]
    fn monotone_under_concatenation() {
        // Property-style: a few dozen generated strings, checking
        // est(a+b) >= est(a) for each.
        let words = [
            "a",
            "ab",
            "abc",
            "hello",
            "world",
            "",
            "x",
            "supercode",
            "token",
            "estimate",
            "!",
            "  ",
            "\n",
            "quick brown fox",
            "z",
            "1234567890",
            "the",
            "lazy",
            "dog",
            "jumps",
            "over",
            "sidecar",
            "reduction",
            "archive",
            "delete",
            "session",
            "store",
            "family",
            "meta",
            "json",
        ];
        for a in &words {
            for b in &words {
                let combined = format!("{a}{b}");
                assert!(
                    estimate_tokens(&combined) >= estimate_tokens(a),
                    "est({a:?}+{b:?}) = {} should be >= est({a:?}) = {}",
                    estimate_tokens(&combined),
                    estimate_tokens(a)
                );
            }
        }
    }

    #[test]
    fn estimate_view_tokens_matches_serde_json_wire_form() {
        let msgs = vec![
            ChatMessage::user("hello there, this is a test message"),
            ChatMessage::assistant("sure, here's a longer reply with more bytes in it"),
        ];
        let expected: u64 = msgs
            .iter()
            .map(|m| estimate_tokens(&serde_json::to_string(m).unwrap()))
            .sum();
        assert_eq!(estimate_view_tokens(&msgs), expected);
    }

    #[test]
    fn estimate_view_tokens_empty_slice_is_zero() {
        assert_eq!(estimate_view_tokens(&[]), 0);
    }

    fn fat_schema(name: &str, filler_len: usize) -> ToolSchema {
        ToolSchema {
            name: name.to_string(),
            description: "x".repeat(filler_len),
            parameters: serde_json::json!({"type": "object", "properties": {}}),
        }
    }

    #[test]
    fn deferred_schema_tokens_full_equals_advertised_is_zero() {
        let full = vec![fat_schema("shell", 500), fat_schema("read_file", 200)];
        let advertised = full.clone();
        assert_eq!(estimate_deferred_schema_tokens(&full, &advertised), 0);
    }

    #[test]
    fn deferred_schema_tokens_measures_the_difference() {
        let builtin = fat_schema("shell", 50);
        let mcp_fat = fat_schema("mcp__github__search_issues", 4000);
        let full = vec![builtin.clone(), mcp_fat];
        let advertised = vec![builtin];
        let deferred = estimate_deferred_schema_tokens(&full, &advertised);
        let expected = estimate_tokens(&serde_json::to_string(&full).unwrap()).saturating_sub(
            estimate_tokens(&serde_json::to_string(&advertised).unwrap()),
        );
        assert_eq!(deferred, expected);
        assert!(deferred > 0, "a fat deferred MCP schema should cost tokens");
    }

    #[test]
    fn deferred_schema_tokens_advertised_larger_than_full_saturates_to_zero() {
        // Pathological input (advertised isn't actually a subset of full) —
        // must not panic or underflow.
        let full = vec![fat_schema("a", 1)];
        let advertised = vec![fat_schema("a", 1000)];
        assert_eq!(estimate_deferred_schema_tokens(&full, &advertised), 0);
    }

    #[test]
    fn fmt_approx_tokens_style() {
        assert_eq!(fmt_approx_tokens(0), "~0 tok");
        assert_eq!(fmt_approx_tokens(21904), "~21,904 tok");
        assert_eq!(fmt_approx_tokens(1000000), "~1,000,000 tok");
    }

    // ---- PARITY-18 D2: guard margin ----

    #[test]
    fn guard_margin_adds_25_percent_and_rounds_up() {
        assert_eq!(with_guard_margin(0), 0);
        assert_eq!(with_guard_margin(4), 5); // 4 * 1.25 = 5.0 exact
        assert_eq!(with_guard_margin(100), 125);
        // 101 * 5 / 4 = 505/4 = 126.25 -> ceil to 127.
        assert_eq!(with_guard_margin(101), 127);
    }

    #[test]
    fn guard_margin_never_decreases() {
        for n in [0u64, 1, 3, 4, 17, 1_000, 1_048_576] {
            assert!(
                with_guard_margin(n) >= n,
                "margin must never make the estimate smaller: {n} -> {}",
                with_guard_margin(n)
            );
        }
    }

    #[test]
    fn guard_margin_saturates_instead_of_overflowing() {
        // The important property: no panic/wraparound on the largest
        // possible input — `saturating_mul` must clamp to `u64::MAX`
        // rather than wrap around to something small (which would defeat
        // the whole point of a "conservative" margin).
        let result = with_guard_margin(u64::MAX);
        assert_eq!(result, u64::MAX.div_ceil(4));
    }

    // ---- PARITY-18 D2: unserializable-message fallback never counts zero ----

    #[test]
    fn message_with_content_never_estimates_to_zero_tokens() {
        let m = ChatMessage::user("hello world, this has real content in it");
        assert!(estimate_message_tokens(&m) > 0);
    }

    #[test]
    fn debug_fallback_is_conservative_not_smaller_than_wire_form() {
        // We can't force `serde_json::to_string` to fail on a real
        // `ChatMessage` (its Serialize impl is infallible in practice), so
        // this test instead pins the CONTRACT the fallback must uphold:
        // the Debug rendering of a message is never a smaller byte count
        // than its wire JSON form, which is what makes it safe to use as
        // the "serialization failed" fallback (D2: conservative, not zero).
        let m = ChatMessage::user("x".repeat(500));
        let wire = serde_json::to_string(&m).unwrap();
        let debug = format!("{m:?}");
        assert!(
            debug.len() >= wire.len(),
            "Debug fallback ({} bytes) must be >= wire form ({} bytes) to stay conservative",
            debug.len(),
            wire.len()
        );
    }

    // ---- PARITY-18 D1: full request-token accounting ----

    fn schema(name: &str, desc_len: usize) -> ToolSchema {
        ToolSchema::new(
            name,
            "x".repeat(desc_len),
            serde_json::json!({"type":"object"}),
        )
    }

    #[test]
    fn estimate_request_tokens_includes_tool_schemas() {
        let messages = vec![ChatMessage::user("hi")];
        let no_tools = estimate_request_tokens(&messages, &[]);
        let with_tools = estimate_request_tokens(&messages, &[schema("shell", 2000)]);
        assert!(
            with_tools > no_tools,
            "a fat tool-schema array must increase the request-token estimate"
        );
    }

    #[test]
    fn estimate_request_tokens_matches_messages_plus_schema_sum() {
        let messages = vec![
            ChatMessage::system("system prompt"),
            ChatMessage::user("user turn"),
        ];
        let tools = vec![schema("read_file", 100), schema("edit", 100)];
        let expected = estimate_view_tokens(&messages)
            + estimate_tokens(&serde_json::to_string(&tools).unwrap());
        assert_eq!(estimate_request_tokens(&messages, &tools), expected);
    }

    // ---- PARITY-18 D1/D4: the shared context_guard decision ----

    #[test]
    fn context_guard_passes_a_small_request() {
        let messages = vec![ChatMessage::user("hi")];
        let (fits, projected) = context_guard(&messages, &[], 200_000);
        assert!(fits, "a tiny request must fit a 200k-token limit");
        assert!(projected < 200_000);
    }

    #[test]
    fn context_guard_refuses_when_reserve_alone_exceeds_limit() {
        // A limit smaller than the completion reserve can never be
        // satisfied, no matter how small the request is.
        let messages = vec![ChatMessage::user("hi")];
        let (fits, _) = context_guard(&messages, &[], CONTEXT_RESPONSE_RESERVE_TOKENS - 1);
        assert!(!fits);
    }

    #[test]
    fn context_guard_messages_only_fit_but_overhead_pushes_over() {
        // D1's exact defect shape: a messages-only estimate comfortably
        // under `context_limit`, but once the (margin-adjusted) tool-schema
        // overhead and completion reserve are added, the true projected
        // request no longer fits. The OLD guard (`view_tokens >
        // context_limit`) would have let this through.
        let context_limit = 10_000u64;
        // ~9,000 raw content tokens (36,000 bytes) — well under
        // `context_limit` on a messages-only basis.
        let big_text = "x".repeat(36_000);
        let messages = vec![ChatMessage::user(big_text)];
        let messages_only = estimate_view_tokens(&messages);
        assert!(
            messages_only < context_limit,
            "fixture must fit messages-only for this test to prove anything: {messages_only} vs {context_limit}"
        );
        // A modest tool-schema array on top.
        let tools = vec![
            schema("shell", 200),
            schema("read_file", 200),
            schema("edit", 200),
        ];
        let (fits, projected) = context_guard(&messages, &tools, context_limit);
        assert!(
            !fits,
            "messages alone fit but margin + schema overhead + reserve should push this over: projected={projected} limit={context_limit}"
        );
    }

    #[test]
    fn context_guard_is_the_single_formula_both_call_sites_share() {
        // Sanity-pin the exact relationship documented on `context_guard`:
        // fits iff with_guard_margin(estimate_request_tokens(..)) + reserve
        // <= limit. If this ever drifts from the implementation, every
        // caller (CLI preflight + Agent::run_loop per-turn guard) drifts
        // silently apart with it.
        let messages = vec![
            ChatMessage::user("hello"),
            ChatMessage::assistant("hi there"),
        ];
        let tools = vec![schema("shell", 50)];
        let limit = 1_000u64;
        let raw = estimate_request_tokens(&messages, &tools);
        let expected_projected = with_guard_margin(raw);
        let expected_fits =
            expected_projected.saturating_add(CONTEXT_RESPONSE_RESERVE_TOKENS) <= limit;
        let (fits, projected) = context_guard(&messages, &tools, limit);
        assert_eq!(projected, expected_projected);
        assert_eq!(fits, expected_fits);
    }
}