feat: implemented native ctok tokenization engine with model scopes
- Implemented the `ctok` Rust native tokenization engine with offline support for Claude V3, V47, V5, and V5Sonnet families. - Replaced global token estimation with model-scoped `Tokenizer` instances and provider-anchored transcript accounting across packages. - Added vocabulary generation scripts, test fixtures, and comprehensive unit tests for tokenizer routing and matching modes.
This commit is contained in:
can1357
Generated
+2
@@ -5106,6 +5106,8 @@ dependencies = [
|
||||
"tokio-util",
|
||||
"toml",
|
||||
"uiautomation",
|
||||
"unicode-normalization",
|
||||
"unicode-properties",
|
||||
"unicode-segmentation",
|
||||
"unicode-width 0.2.2",
|
||||
"windows-sys 0.61.2",
|
||||
|
||||
@@ -233,6 +233,8 @@ pdf-inspector = "1"
|
||||
regex = "1"
|
||||
similar = "3.1.0"
|
||||
unicode-segmentation = "1.13"
|
||||
unicode-normalization = "0.1"
|
||||
unicode-properties = "0.1"
|
||||
unicode-width = "0.2"
|
||||
fontdue = { version = "0.9", default-features = false }
|
||||
|
||||
|
||||
@@ -18,6 +18,8 @@ rust_shared_library(
|
||||
compile_data = glob([
|
||||
"src/syntaxes/*.sublime-syntax",
|
||||
"src/fonts/*",
|
||||
"src/ctok/data/*.bin",
|
||||
"src/ctok/testdata/*.json",
|
||||
]),
|
||||
crate_features = [],
|
||||
edition = "2024",
|
||||
|
||||
@@ -65,6 +65,8 @@ tokio.workspace = true
|
||||
tokio-util.workspace = true
|
||||
toml.workspace = true
|
||||
unicode-segmentation.workspace = true
|
||||
unicode-normalization.workspace = true
|
||||
unicode-properties.workspace = true
|
||||
unicode-width.workspace = true
|
||||
xxhash-rust.workspace = true
|
||||
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
//! Marker notation and the measured character tables.
|
||||
//!
|
||||
//! Port of ctok's `constants.py`. The tables were measured against Anthropic's
|
||||
//! `count_tokens`; tables noted as enumerated must stay enumerated — no Unicode
|
||||
//! category rule reproduces them.
|
||||
|
||||
/// ⟨bow⟩ — begin-of-word marker byte.
|
||||
pub const BOW: u8 = 0x01;
|
||||
/// ⟨eow⟩ — end-of-word marker byte.
|
||||
pub const EOW: u8 = 0x02;
|
||||
// 0x03 is ⟨pad⟩, the message-frame marker: it never stands in the marked
|
||||
// stream, only inside vocabulary pieces.
|
||||
/// ⟨shift⟩ — title-case marker byte.
|
||||
pub const SHIFT: u8 = 0x04;
|
||||
/// ⟨caps⟩ — all-caps marker byte.
|
||||
pub const CAPS: u8 = 0x05;
|
||||
|
||||
/// Whether `b` is one of the five marker bytes, ⟨bow⟩ through ⟨caps⟩.
|
||||
///
|
||||
/// ctok writes its markers as the noncharacters U+FDD0..=U+FDD4; this port and
|
||||
/// its vocabulary generator write one C0 byte each instead, three times shorter
|
||||
/// in the stream and in the matching automaton. Text cannot collide with them:
|
||||
/// [`super::normalize::nfc`] folds NUL to a space and strips every other C0
|
||||
/// control before a stream is written, so a marker byte in a stream is always a
|
||||
/// marker — where ctok's spelling has to escape literal noncharacters out of
|
||||
/// the way, this needs nothing.
|
||||
pub const fn is_marker_byte(b: u8) -> bool {
|
||||
matches!(b, BOW..=CAPS)
|
||||
}
|
||||
|
||||
/// Non-ASCII symbols/punctuation that tile over the punct vocabulary rather
|
||||
/// than standing alone. Enumerated: the behaviour splits per codepoint with no
|
||||
/// categorical rule (`（` is punct but `）` is hard). Anything unlisted stays
|
||||
/// HARD.
|
||||
pub const fn is_punct_sym(c: char) -> bool {
|
||||
matches!(
|
||||
c,
|
||||
'—' | '»'
|
||||
| '«' | '•'
|
||||
| '°' | '„'
|
||||
| '–' | '−'
|
||||
| '£' | '§'
|
||||
| '€' | '…'
|
||||
| '√' | '→'
|
||||
| '（' | '№'
|
||||
| '†' | '└'
|
||||
| '│' | '།'
|
||||
| '·' | '─'
|
||||
| '═' | '█'
|
||||
)
|
||||
}
|
||||
|
||||
/// Symbol-letters measured to take the full word model exactly like Latin
|
||||
/// letters. Enumerated blocks: category Nl/So/Lu splits both ways, and the
|
||||
/// Hangzhou numerals (also Nl) measured markerless, so it is the block that
|
||||
/// predicts, not the category.
|
||||
pub const fn is_symbol_letter(o: u32) -> bool {
|
||||
matches!(o,
|
||||
0x16ee..=0x16f0 // Runic golden numbers (Nl, caseless)
|
||||
| 0x2160..=0x2188 // Roman numerals (Nl/Lu/Ll, cased pairs)
|
||||
| 0x24b6..=0x24e9 // circled letters (So, cased)
|
||||
| 0xa6e6..=0xa6ef // Bamum number-letters (Nl, caseless)
|
||||
)
|
||||
}
|
||||
|
||||
/// Variation selectors are gc=Mn but take no word model. The supplementary
|
||||
/// selectors (U+E0100..) are astral and already HARD.
|
||||
pub const fn is_variation_selector(c: char) -> bool {
|
||||
matches!(c, '\u{fe00}'..='\u{fe0f}')
|
||||
}
|
||||
|
||||
/// The one canonical-combining-class-9 character that does not separate word
|
||||
/// runs: U+0E3A THAI CHARACTER PHINTHU.
|
||||
pub const NON_SEPARATOR: char = '\u{0e3a}';
|
||||
|
||||
/// The suffixes an apostrophe binds into the word ahead of it, deleting that
|
||||
/// word's ⟨bow⟩. Standard English contraction set, lowercase and whole-word
|
||||
/// only; measured per member.
|
||||
pub const fn is_contraction_suffix(body: &[u8]) -> bool {
|
||||
matches!(body, b"s" | b"t" | b"d" | b"m" | b"ll" | b"re" | b"ve")
|
||||
}
|
||||
|
||||
/// C0/C1 controls the API strips before tokenizing (cost 0): every gc=Cc
|
||||
/// except TAB, LF and NUL.
|
||||
pub const fn is_stripped_control(c: char) -> bool {
|
||||
matches!(c, '\u{01}'..='\u{08}' | '\u{0b}'..='\u{1f}' | '\u{7f}'..='\u{9f}')
|
||||
}
|
||||
|
||||
/// BMP private use is stripped the same way; its two neighbours join into one
|
||||
/// word. The supplementary private-use planes are unprobed and deliberately
|
||||
/// left out.
|
||||
pub const fn is_stripped_private(c: char) -> bool {
|
||||
matches!(c, '\u{e000}'..='\u{f8ff}')
|
||||
}
|
||||
|
||||
/// Space separators the tokenizer treats identically to U+0020: all Zs except
|
||||
/// U+3000 (ideographic space), plus Zl/Zp. TAB, LF and U+3000 each have their
|
||||
/// own cost.
|
||||
pub const fn is_funny_space(c: char) -> bool {
|
||||
matches!(
|
||||
c,
|
||||
'\u{a0}' | '\u{1680}' | '\u{2000}'
|
||||
..='\u{200a}' | '\u{2028}' | '\u{2029}' | '\u{202f}' | '\u{205f}'
|
||||
)
|
||||
}
|
||||
|
||||
/// The four standard curly quotes fold to their ASCII forms (v3 only; NFC does
|
||||
/// not do this). The low-9 mark U+201E is a different token and is deliberately
|
||||
/// not folded.
|
||||
pub const fn fold_quote(c: char) -> char {
|
||||
match c {
|
||||
'\u{2018}' | '\u{2019}' => '\'',
|
||||
'\u{201c}' | '\u{201d}' => '"',
|
||||
_ => c,
|
||||
}
|
||||
}
|
||||
|
||||
/// Marks that terminate the orthographic syllable and so separate word runs,
|
||||
/// beyond the ccc-9 viramas: the measured U+0300 combining ranges, the swept
|
||||
/// combining-block annotation ranges, and the enumerated separator signs
|
||||
/// (Thai/Lao tone marks, nukta, Khmer consonant shifters, …). Merged into one
|
||||
/// sorted range table; provenance per group lives in ctok's `constants.py`.
|
||||
static SEPARATOR_RANGES: &[(u32, u32)] = &[
|
||||
(0x0300, 0x0344),
|
||||
(0x0346, 0x0362),
|
||||
(0x0483, 0x0489),
|
||||
(0x0591, 0x05af),
|
||||
(0x0658, 0x0658),
|
||||
(0x06df, 0x06e0),
|
||||
(0x06ea, 0x06ec),
|
||||
(0x0740, 0x074a),
|
||||
(0x07eb, 0x07f3),
|
||||
(0x07fd, 0x07fd),
|
||||
(0x0818, 0x0819),
|
||||
(0x082d, 0x082d),
|
||||
(0x0859, 0x085b),
|
||||
(0x0898, 0x089f),
|
||||
(0x08ca, 0x08d3),
|
||||
(0x08e0, 0x08e1),
|
||||
(0x08ea, 0x08ef),
|
||||
(0x093c, 0x093c),
|
||||
(0x0951, 0x0954),
|
||||
(0x09bc, 0x09bc),
|
||||
(0x09fe, 0x09fe),
|
||||
(0x0a3c, 0x0a3c),
|
||||
(0x0abc, 0x0abc),
|
||||
(0x0afd, 0x0aff),
|
||||
(0x0b3c, 0x0b3c),
|
||||
(0x0b55, 0x0b55),
|
||||
(0x0c3c, 0x0c3c),
|
||||
(0x0cbc, 0x0cbc),
|
||||
(0x0e47, 0x0e4c),
|
||||
(0x0e4e, 0x0e4e),
|
||||
(0x0ec8, 0x0ecc),
|
||||
(0x0ece, 0x0ece),
|
||||
(0x0f18, 0x0f19),
|
||||
(0x0f35, 0x0f35),
|
||||
(0x0f37, 0x0f37),
|
||||
(0x0f39, 0x0f39),
|
||||
(0x0f3e, 0x0f3f),
|
||||
(0x0f86, 0x0f87),
|
||||
(0x0fc6, 0x0fc6),
|
||||
(0x1037, 0x1037),
|
||||
(0x135d, 0x135f),
|
||||
(0x17b4, 0x17b5),
|
||||
(0x17c9, 0x17d1),
|
||||
(0x17d3, 0x17d3),
|
||||
(0x17dd, 0x17dd),
|
||||
(0x180b, 0x180d),
|
||||
(0x180f, 0x180f),
|
||||
(0x1939, 0x193b),
|
||||
(0x1a75, 0x1a7c),
|
||||
(0x1a7f, 0x1a7f),
|
||||
(0x1ab0, 0x1abe),
|
||||
(0x1ac1, 0x1acb),
|
||||
(0x1b34, 0x1b34),
|
||||
(0x1b6b, 0x1b73),
|
||||
(0x1be6, 0x1be6),
|
||||
(0x1c37, 0x1c37),
|
||||
(0x1cd0, 0x1ce8),
|
||||
(0x1ced, 0x1ced),
|
||||
(0x1cf4, 0x1cf4),
|
||||
(0x1cf7, 0x1cf7),
|
||||
(0x1cf8, 0x1cf9),
|
||||
(0x1dc0, 0x1dd2),
|
||||
(0x1df5, 0x1dff),
|
||||
(0x20d0, 0x20f0),
|
||||
(0x2cef, 0x2cf1),
|
||||
(0x302a, 0x302f),
|
||||
(0x3099, 0x309a),
|
||||
(0xa66f, 0xa672),
|
||||
(0xa67c, 0xa67d),
|
||||
(0xa6f0, 0xa6f1),
|
||||
(0xa8e0, 0xa8f1),
|
||||
(0xa92b, 0xa92d),
|
||||
(0xa9b3, 0xa9b3),
|
||||
(0xaabf, 0xaabf),
|
||||
(0xaac1, 0xaac1),
|
||||
(0xabec, 0xabec),
|
||||
(0xfe20, 0xfe2f),
|
||||
];
|
||||
|
||||
/// Whether `o` falls in one of the measured separator ranges (excluding the
|
||||
/// ccc-9 virama population, which [`super::normalize::is_separator`] handles).
|
||||
pub fn in_separator_ranges(o: u32) -> bool {
|
||||
let idx = SEPARATOR_RANGES.partition_point(|&(lo, _)| lo <= o);
|
||||
idx > 0 && o <= SEPARATOR_RANGES[idx - 1].1
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 Sander Land
|
||||
(measured tokenizer vocabulary data, reconstruction model, and reference
|
||||
implementation: https://github.com/sanderland/ctok)
|
||||
Copyright (c) 2026 Can Bölük and the Oh My Pi contributors
|
||||
(Rust implementation and the compact binary vocabulary encoding in
|
||||
crates/pi-natives/src/ctok)
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,27 @@
|
||||
# ctok vocabulary data
|
||||
|
||||
`ctok_v3.bin` and `ctok_v4_7.bin` are **generated** — do not hand-edit.
|
||||
They are compacted from the measured vocabulary files of
|
||||
[sanderland/ctok](https://github.com/sanderland/ctok) v1.0.0 (revision
|
||||
`df3b59b5e645289a5eadc8e24036b99d39c333c4`), MIT licensed — see
|
||||
`LICENSE.ctok`. The vocabulary data is Sander Land's measurement work
|
||||
("On the biology of Claude's tokenizer",
|
||||
<https://tokencontributions.substack.com/p/on-the-biology-of-claudes-tokenizer>);
|
||||
the Rust implementation in the parent directory is this repository's own.
|
||||
|
||||
Upstream ships every piece with a `count_tokens` witness probe; compaction
|
||||
drops that metadata, parses the public `⟨bow⟩the⟨eow⟩` key notation into the
|
||||
internal marked form — one byte per marker, a third shorter than ctok's
|
||||
noncharacter spelling in both the pieces and the stream they tile — adds the
|
||||
glued contraction spellings, and front-codes the sorted piece list into the
|
||||
binary format documented in `packages/natives/scripts/gen-ctok-vocab.ts`
|
||||
(~4.7 MB of upstream JSON → ~254 KB embedded).
|
||||
|
||||
Regenerate with:
|
||||
|
||||
```sh
|
||||
bun --cwd packages/natives run gen:ctok
|
||||
```
|
||||
|
||||
If the upstream pin moves, also regenerate `../testdata/fixtures.json`
|
||||
against the same ctok release (see the fixture doc in `../mod.rs`).
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,557 @@
|
||||
//! The min-cost tiling, with the marked-stream vocabulary as an Aho-Corasick
|
||||
//! automaton over UTF-8 bytes.
|
||||
//!
|
||||
//! Port of ctok's `engine.py`, count-only: marker atoms are in the vocabulary
|
||||
//! as cost-1 tokens, so a marker no piece absorbs tiles as itself; the count
|
||||
//! is the number of tiles. Token materialization (`tokenize()`) is not ported
|
||||
//! — every consumer here needs the count alone.
|
||||
//!
|
||||
//! The vocabulary is matched with one Aho-Corasick transition per stream byte:
|
||||
//! the state a byte lands in names, through its dictionary-link chain, every
|
||||
//! piece ending at that byte, so the DP scores exactly the candidate tiles that
|
||||
//! exist. A trie of reversed pieces answers the same question with a fresh
|
||||
//! backwards descent per position, re-walking the vocabulary for every byte.
|
||||
//!
|
||||
//! Only pieces that end where a tile ends are candidates, and all of them are
|
||||
//! scored: `best` is *not* monotone, because a long piece can cover a prefix
|
||||
//! that is expensive to tile, so the longest match is not always the cheapest.
|
||||
//!
|
||||
//! Stream and vocabulary are both UTF-8 byte strings, never decoded character
|
||||
//! sequences: a tile boundary is a byte offset, a transition is one byte, and a
|
||||
//! character is decoded only where no piece covers it. Byte matching is exact
|
||||
//! because a substring of valid UTF-8 that is itself valid UTF-8 must start on
|
||||
//! a character boundary — a continuation byte cannot open a sequence — so no
|
||||
//! piece can match across a character seam, and no match can be shorter than
|
||||
//! its final character.
|
||||
|
||||
use std::collections::{HashSet, VecDeque};
|
||||
|
||||
use super::{constants::is_marker_byte, normalize::FrameParams};
|
||||
|
||||
/// Whether `b` continues a multi-byte UTF-8 sequence, so it is not the start of
|
||||
/// a character and no tile may begin or end there.
|
||||
#[inline]
|
||||
const fn is_continuation(b: u8) -> bool {
|
||||
b & 0xc0 == 0x80
|
||||
}
|
||||
|
||||
/// The vocabulary as an Aho-Corasick automaton: [`PieceMatcher::advance`] walks
|
||||
/// the stream one byte at a time and the dictionary-link chain from the state
|
||||
/// it lands in enumerates every piece ending at that byte, longest first.
|
||||
pub struct PieceMatcher {
|
||||
states: Vec<State>,
|
||||
edge_bytes: Vec<u8>,
|
||||
edge_targets: Vec<u32>,
|
||||
/// Dense transitions out of the root, which is the hottest state by far.
|
||||
/// Zero (the root itself) where no piece starts with that byte.
|
||||
root_goto: [u32; 256],
|
||||
}
|
||||
|
||||
/// One automaton state: goto edges, the fail link, and the piece chain the DP
|
||||
/// reads. Sixteen bytes, so a transition and the match it reports share a cache
|
||||
/// line.
|
||||
struct State {
|
||||
edge_start: u32,
|
||||
/// Longest proper suffix of this state's path that is also a state.
|
||||
fail: u32,
|
||||
/// Next state along the fail chain at which a piece ends, 0 for none — the
|
||||
/// root can serve as that sentinel because no piece is empty.
|
||||
dict: u32,
|
||||
edge_count: u16,
|
||||
/// Byte length of the piece ending in this state, 0 if none ends here.
|
||||
out_len: u16,
|
||||
}
|
||||
|
||||
/// Construction form: the vocabulary arrives sorted by bytes (the blob is
|
||||
/// front-coded), so the trie grows along one path stack without any hashing,
|
||||
/// and each parent's edges are created in increasing byte order.
|
||||
struct Builder {
|
||||
terminal: Vec<bool>,
|
||||
depth: Vec<u32>,
|
||||
edges: Vec<(u32, u8, u32)>,
|
||||
/// States along the piece last inserted, indexed by prefix length.
|
||||
path: Vec<u32>,
|
||||
}
|
||||
|
||||
impl Builder {
|
||||
fn new(piece_count: usize) -> Self {
|
||||
Self {
|
||||
terminal: vec![false],
|
||||
depth: vec![0],
|
||||
edges: Vec::with_capacity(piece_count * 2),
|
||||
path: vec![0],
|
||||
}
|
||||
}
|
||||
|
||||
/// Insert the next piece of the sorted list: `shared` leading bytes the
|
||||
/// previous piece already walked, then `suffix`.
|
||||
fn push_piece(&mut self, shared: usize, suffix: &[u8]) {
|
||||
self.path.truncate(shared + 1);
|
||||
let mut at = *self.path.last().expect("root state");
|
||||
for &b in suffix {
|
||||
let child = self.terminal.len() as u32;
|
||||
self.terminal.push(false);
|
||||
self.depth.push(self.depth[at as usize] + 1);
|
||||
self.edges.push((at, b, child));
|
||||
self.path.push(child);
|
||||
at = child;
|
||||
}
|
||||
self.terminal[at as usize] = true;
|
||||
}
|
||||
|
||||
/// Group the edges by parent and link the automaton.
|
||||
fn freeze(self) -> PieceMatcher {
|
||||
let count = self.terminal.len();
|
||||
// Counting sort by parent; within a parent the creation order is
|
||||
// already byte-sorted, so the groups come out sorted.
|
||||
let mut offsets = vec![0u32; count + 1];
|
||||
for &(parent, ..) in &self.edges {
|
||||
offsets[parent as usize + 1] += 1;
|
||||
}
|
||||
for i in 0..count {
|
||||
offsets[i + 1] += offsets[i];
|
||||
}
|
||||
let mut edge_bytes = vec![0u8; self.edges.len()];
|
||||
let mut edge_targets = vec![0u32; self.edges.len()];
|
||||
let mut cursor = offsets.clone();
|
||||
for &(parent, b, child) in &self.edges {
|
||||
let at = cursor[parent as usize] as usize;
|
||||
cursor[parent as usize] += 1;
|
||||
edge_bytes[at] = b;
|
||||
edge_targets[at] = child;
|
||||
}
|
||||
|
||||
let states = (0..count)
|
||||
.map(|i| State {
|
||||
edge_start: offsets[i],
|
||||
edge_count: u16::try_from(offsets[i + 1] - offsets[i]).expect("byte fanout"),
|
||||
out_len: 0,
|
||||
fail: 0,
|
||||
dict: 0,
|
||||
})
|
||||
.collect();
|
||||
let mut matcher = PieceMatcher { states, edge_bytes, edge_targets, root_goto: [0; 256] };
|
||||
matcher.link(&self.terminal, &self.depth);
|
||||
matcher
|
||||
}
|
||||
}
|
||||
|
||||
impl PieceMatcher {
|
||||
/// Fill the fail links, piece lengths and dictionary links breadth-first: a
|
||||
/// state's fail link points at a strictly shallower state, so every lookup
|
||||
/// below reads a state this loop has already linked.
|
||||
fn link(&mut self, terminal: &[bool], depth: &[u32]) {
|
||||
let mut queue: VecDeque<u32> = VecDeque::new();
|
||||
let root = &self.states[0];
|
||||
let (start, count) = (root.edge_start as usize, usize::from(root.edge_count));
|
||||
for i in start..start + count {
|
||||
let child = self.edge_targets[i];
|
||||
self.root_goto[usize::from(self.edge_bytes[i])] = child;
|
||||
queue.push_back(child);
|
||||
}
|
||||
while let Some(u) = queue.pop_front() {
|
||||
let fail = self.states[u as usize].fail;
|
||||
self.states[u as usize].out_len = if terminal[u as usize] {
|
||||
u16::try_from(depth[u as usize]).expect("piece length fits u16")
|
||||
} else {
|
||||
0
|
||||
};
|
||||
// Shortcut straight to the next shorter piece, so the DP walks
|
||||
// matches only, never the whole fail chain.
|
||||
self.states[u as usize].dict = if self.states[fail as usize].out_len > 0 {
|
||||
fail
|
||||
} else {
|
||||
self.states[fail as usize].dict
|
||||
};
|
||||
let st = &self.states[u as usize];
|
||||
let (start, count) = (st.edge_start as usize, usize::from(st.edge_count));
|
||||
for i in start..start + count {
|
||||
let child = self.edge_targets[i];
|
||||
// `fail` is shallower than `u`, so this transition only reads
|
||||
// linked states and can never land on `child` itself.
|
||||
self.states[child as usize].fail = self.advance(fail, self.edge_bytes[i]);
|
||||
queue.push_back(child);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Follow one goto edge, or `None` when this state has no piece continuing
|
||||
/// with `b`.
|
||||
#[inline]
|
||||
fn goto(&self, state: u32, b: u8) -> Option<u32> {
|
||||
let st = &self.states[state as usize];
|
||||
let start = st.edge_start as usize;
|
||||
let bytes = &self.edge_bytes[start..start + usize::from(st.edge_count)];
|
||||
// Fanout is one or two bytes for most states, where scanning the sorted
|
||||
// run beats a binary search's branch chain; the wide states sit near
|
||||
// the root, and the root itself never gets here.
|
||||
let hit = if bytes.len() <= 8 {
|
||||
bytes
|
||||
.iter()
|
||||
.position(|&e| e >= b)
|
||||
.filter(|&i| bytes[i] == b)
|
||||
} else {
|
||||
bytes.binary_search(&b).ok()
|
||||
}?;
|
||||
Some(self.edge_targets[start + hit])
|
||||
}
|
||||
|
||||
/// Consume one stream byte. Each fail step drops at least one level and a
|
||||
/// transition adds at most one, so the whole scan is linear in the stream.
|
||||
#[inline]
|
||||
fn advance(&self, state: u32, b: u8) -> u32 {
|
||||
let mut at = state;
|
||||
loop {
|
||||
if at == 0 {
|
||||
return self.root_goto[usize::from(b)];
|
||||
}
|
||||
if let Some(next) = self.goto(at, b) {
|
||||
return next;
|
||||
}
|
||||
at = self.states[at as usize].fail;
|
||||
}
|
||||
}
|
||||
|
||||
/// The pieces ending at the byte this state consumed, longest first: the
|
||||
/// state itself when one ends here, then its dictionary chain.
|
||||
#[inline]
|
||||
const fn matches(&self, state: u32) -> Matches<'_> {
|
||||
Matches { vocab: self, at: state }
|
||||
}
|
||||
}
|
||||
|
||||
/// Iterator over the byte lengths of the pieces ending at one stream position,
|
||||
/// longest first.
|
||||
struct Matches<'a> {
|
||||
vocab: &'a PieceMatcher,
|
||||
at: u32,
|
||||
}
|
||||
|
||||
impl Iterator for Matches<'_> {
|
||||
type Item = usize;
|
||||
|
||||
#[inline]
|
||||
fn next(&mut self) -> Option<usize> {
|
||||
loop {
|
||||
let st = &self.vocab.states[self.at as usize];
|
||||
let len = st.out_len;
|
||||
self.at = st.dict;
|
||||
if len > 0 {
|
||||
return Some(usize::from(len));
|
||||
}
|
||||
// Only the entry state can lack a piece; the chain holds nothing
|
||||
// but pieces after that.
|
||||
if self.at == 0 {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Min-cost tiling over the cost-1 vocabulary plus a guaranteed one-character
|
||||
/// floor. `unit_cost(start, end)` prices the character in `s[start..end]` where
|
||||
/// no piece covers it (may be > 1: a 4-byte letter with no piece costs its byte
|
||||
/// tiling).
|
||||
///
|
||||
/// Long pieces make `best` non-monotone — covering a longer prefix can cost
|
||||
/// fewer tokens than covering a shorter one — so every piece ending at a
|
||||
/// position is scored, not just the longest. The one-character floor is needed
|
||||
/// only when no piece spells that character alone: any piece ending here covers
|
||||
/// at least the final character (a piece cannot start mid-character), and the
|
||||
/// character's own piece costs 1, never more than its floor.
|
||||
pub fn min_vocab_tile(
|
||||
s: &[u8],
|
||||
vocab: &PieceMatcher,
|
||||
mut unit_cost: impl FnMut(usize, usize) -> u32,
|
||||
) -> u32 {
|
||||
let n = s.len();
|
||||
if n == 0 {
|
||||
return 0;
|
||||
}
|
||||
let mut best = vec![0u32; n + 1];
|
||||
let mut state = 0u32;
|
||||
for end in 1..=n {
|
||||
state = vocab.advance(state, s[end - 1]);
|
||||
// Tiles start and end on character boundaries only; interior byte
|
||||
// positions of a character are never a DP state.
|
||||
if end != n && is_continuation(s[end]) {
|
||||
continue;
|
||||
}
|
||||
let mut start = end - 1;
|
||||
while is_continuation(s[start]) {
|
||||
start -= 1;
|
||||
}
|
||||
let single = end - start;
|
||||
let mut cost = u32::MAX;
|
||||
let mut spelled = false;
|
||||
for len in vocab.matches(state) {
|
||||
cost = cost.min(best[end - len] + 1);
|
||||
spelled |= len == single;
|
||||
}
|
||||
if !spelled {
|
||||
cost = cost.min(best[start] + unit_cost(start, end));
|
||||
}
|
||||
best[end] = cost;
|
||||
}
|
||||
best[n]
|
||||
}
|
||||
|
||||
/// What a codepoint costs when no piece covers it: a min-cost tiling of its
|
||||
/// UTF-8 bytes over the partial byte-prefix tokens, every single byte
|
||||
/// costing 1.
|
||||
struct ByteFloor {
|
||||
/// Sorted sentinel-packed byte strings (≤ 4 bytes): membership is all that
|
||||
/// is needed, since every token costs 1.
|
||||
tokens: Vec<u64>,
|
||||
max_len: usize,
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn pack_bytes(bs: &[u8]) -> u64 {
|
||||
bs.iter().fold(1u64, |acc, &b| (acc << 8) | u64::from(b))
|
||||
}
|
||||
|
||||
impl ByteFloor {
|
||||
fn cost_bytes(&self, bs: &[u8]) -> u32 {
|
||||
let n = bs.len();
|
||||
let mut best = [u32::MAX; 5];
|
||||
best[0] = 0;
|
||||
for i in 1..=n {
|
||||
for j in i.saturating_sub(self.max_len)..i {
|
||||
if best[j] != u32::MAX
|
||||
&& (i - j == 1 || self.tokens.binary_search(&pack_bytes(&bs[j..i])).is_ok())
|
||||
&& best[j] + 1 < best[i]
|
||||
{
|
||||
best[i] = best[j] + 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
best[n]
|
||||
}
|
||||
|
||||
fn cost_char(&self, c: char) -> u32 {
|
||||
let mut buf = [0u8; 4];
|
||||
self.cost_bytes(c.encode_utf8(&mut buf).as_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
/// The first character of `bytes`, whose length the lead byte gives — valid
|
||||
/// UTF-8 is guaranteed by both callers (vocabulary pieces and the encoder's own
|
||||
/// output).
|
||||
#[inline]
|
||||
fn decode_char(bytes: &[u8]) -> char {
|
||||
#[inline]
|
||||
const fn cont(b: u8) -> u32 {
|
||||
(b & 0x3f) as u32
|
||||
}
|
||||
let lead = u32::from(bytes[0]);
|
||||
let cp = if lead < 0x80 {
|
||||
lead
|
||||
} else if lead < 0xe0 {
|
||||
(lead & 0x1f) << 6 | cont(bytes[1])
|
||||
} else if lead < 0xf0 {
|
||||
(lead & 0x0f) << 12 | cont(bytes[1]) << 6 | cont(bytes[2])
|
||||
} else {
|
||||
(lead & 0x07) << 18 | cont(bytes[1]) << 12 | cont(bytes[2]) << 6 | cont(bytes[3])
|
||||
};
|
||||
char::from_u32(cp).expect("valid UTF-8")
|
||||
}
|
||||
|
||||
/// Byte cursor over the embedded vocabulary blob. The data is generated and
|
||||
/// committed alongside the code, so malformed input is a build defect:
|
||||
/// readers panic with context rather than propagating errors.
|
||||
struct Cursor<'a> {
|
||||
data: &'a [u8],
|
||||
pos: usize,
|
||||
}
|
||||
|
||||
impl<'a> Cursor<'a> {
|
||||
fn take(&mut self, n: usize) -> &'a [u8] {
|
||||
let slice = &self.data[self.pos..self.pos + n];
|
||||
self.pos += n;
|
||||
slice
|
||||
}
|
||||
|
||||
fn u8(&mut self) -> u8 {
|
||||
self.take(1)[0]
|
||||
}
|
||||
|
||||
fn u16(&mut self) -> u16 {
|
||||
u16::from_le_bytes(self.take(2).try_into().expect("two bytes"))
|
||||
}
|
||||
|
||||
fn u32(&mut self) -> u32 {
|
||||
u32::from_le_bytes(self.take(4).try_into().expect("four bytes"))
|
||||
}
|
||||
|
||||
fn varint(&mut self) -> usize {
|
||||
let mut value = 0usize;
|
||||
let mut shift = 0u32;
|
||||
loop {
|
||||
let byte = self.u8();
|
||||
value |= usize::from(byte & 0x7f) << shift;
|
||||
if byte & 0x80 == 0 {
|
||||
return value;
|
||||
}
|
||||
shift += 7;
|
||||
assert!(shift < 32, "ctok varint overflow");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The loaded vocabulary plus the scalars the encoder and tiler read from it.
|
||||
/// One core serves every family that borrows its vocabulary file (v5 reuses
|
||||
/// v4.7's with different frame scalars).
|
||||
pub struct VocabCore {
|
||||
/// Measured metadata of the vocabulary file (v5 overrides parts of it).
|
||||
pub message_overhead: u32,
|
||||
/// See [`FrameParams::fold_quotes`].
|
||||
pub fold_quotes: bool,
|
||||
/// See [`FrameParams::allcaps_min`].
|
||||
pub allcaps_min: Option<usize>,
|
||||
vocab: PieceMatcher,
|
||||
/// Sorted codepoints of the cost-1 whole-character pieces.
|
||||
unit_pieces: Vec<u32>,
|
||||
/// Sorted lengths of the pure-newline pieces; coins for the
|
||||
/// trailing-newline ladder, where nothing else can match.
|
||||
newline_ladder: Vec<u32>,
|
||||
floor: ByteFloor,
|
||||
}
|
||||
|
||||
impl VocabCore {
|
||||
/// Parse one front-coded binary vocabulary blob produced by
|
||||
/// `packages/natives/scripts/gen-ctok-vocab.ts` (format documented there;
|
||||
/// pieces arrive in the compact marker alphabet, sorted by those bytes).
|
||||
/// Pieces stream straight into the automaton builder; nothing is buffered
|
||||
/// beyond the front-coding scratch.
|
||||
pub fn parse(blob: &[u8]) -> Self {
|
||||
let mut cur = Cursor { data: blob, pos: 0 };
|
||||
assert_eq!(cur.take(4), b"CTOK", "bad ctok vocabulary magic");
|
||||
assert_eq!(cur.u8(), 2, "unsupported ctok vocabulary version");
|
||||
let fold_quotes = cur.u8() & 1 != 0;
|
||||
let message_overhead = u32::from(cur.u8());
|
||||
let allcaps_min = match cur.u8() {
|
||||
0 => None,
|
||||
n => Some(usize::from(n)),
|
||||
};
|
||||
let byte_token_count = usize::from(cur.u16());
|
||||
let piece_count = cur.u32();
|
||||
|
||||
let mut tokens = HashSet::with_capacity(byte_token_count + 512);
|
||||
let mut max_len = 1usize;
|
||||
for _ in 0..byte_token_count {
|
||||
let len = usize::from(cur.u8());
|
||||
assert!((1..=4).contains(&len), "byte token out of range: {len}");
|
||||
max_len = max_len.max(len);
|
||||
tokens.insert(pack_bytes(cur.take(len)));
|
||||
}
|
||||
|
||||
let mut builder = Builder::new(piece_count as usize);
|
||||
let mut unit_pieces = Vec::new();
|
||||
let mut newline_ladder = Vec::new();
|
||||
let mut scratch: Vec<u8> = Vec::with_capacity(64);
|
||||
for _ in 0..piece_count {
|
||||
let shared = cur.varint();
|
||||
let suffix_len = cur.varint();
|
||||
assert!(shared <= scratch.len(), "ctok pieces should be front-coded in order");
|
||||
scratch.truncate(shared);
|
||||
let suffix = cur.take(suffix_len);
|
||||
scratch.extend_from_slice(suffix);
|
||||
builder.push_piece(shared, suffix);
|
||||
assert!(std::str::from_utf8(&scratch).is_ok(), "ctok piece should be UTF-8");
|
||||
if scratch.iter().all(|&b| b == b'\n') {
|
||||
newline_ladder.push(scratch.len() as u32);
|
||||
}
|
||||
// Cost-1 whole-character pieces fold into the byte floor's
|
||||
// membership set, so an uncovered character still prices at 1.
|
||||
// Marker atoms are pieces too, but they spell no character.
|
||||
if scratch.len() == 1 && is_marker_byte(scratch[0]) {
|
||||
continue;
|
||||
}
|
||||
let c = decode_char(&scratch);
|
||||
if scratch.len() == c.len_utf8() {
|
||||
unit_pieces.push(c as u32);
|
||||
max_len = max_len.max(scratch.len());
|
||||
tokens.insert(pack_bytes(&scratch));
|
||||
}
|
||||
}
|
||||
assert_eq!(cur.pos, blob.len(), "trailing ctok vocabulary bytes");
|
||||
newline_ladder.sort_unstable();
|
||||
unit_pieces.sort_unstable();
|
||||
let mut tokens: Vec<u64> = tokens.into_iter().collect();
|
||||
tokens.sort_unstable();
|
||||
|
||||
Self {
|
||||
message_overhead,
|
||||
fold_quotes,
|
||||
allcaps_min,
|
||||
vocab: builder.freeze(),
|
||||
unit_pieces,
|
||||
newline_ladder,
|
||||
floor: ByteFloor { tokens, max_len },
|
||||
}
|
||||
}
|
||||
|
||||
/// One character standing where the vocabulary covers nothing: markers and
|
||||
/// whole-character pieces cost one token, everything else falls to the byte
|
||||
/// floor.
|
||||
fn uncovered_cost(&self, bytes: &[u8]) -> u32 {
|
||||
if bytes.len() == 1 && is_marker_byte(bytes[0]) {
|
||||
return 1; // a marker no piece absorbs tiles as itself
|
||||
}
|
||||
let c = decode_char(bytes);
|
||||
if self.unit_pieces.binary_search(&(c as u32)).is_ok() {
|
||||
1
|
||||
} else {
|
||||
self.floor.cost_char(c)
|
||||
}
|
||||
}
|
||||
|
||||
/// Tile the marked stream: minimum token count over the vocabulary, with
|
||||
/// markers costing one and uncovered characters falling to the byte floor.
|
||||
pub fn tile_cost(&self, stream: &[u8]) -> u32 {
|
||||
min_vocab_tile(stream, &self.vocab, |start, end| self.uncovered_cost(&stream[start..end]))
|
||||
}
|
||||
|
||||
/// What a content-final run of `n_tail` frame-absorbed newlines costs
|
||||
/// beyond the frame's own trailing token (ladder families only). The
|
||||
/// v3/v4.7 frames append ⏎⏎ after the content and one token can span
|
||||
/// into them, so the run the tokenizer sees is `n_tail + 2`; the
|
||||
/// sonnet/fable-5 frame appends nothing (`appended = 0`). In both cases
|
||||
/// the frame already pays for one token of the run.
|
||||
pub fn ladder_tail_cost(&self, n_tail: usize, appended: usize) -> u32 {
|
||||
if n_tail == 0 {
|
||||
return 0;
|
||||
}
|
||||
let m = n_tail + appended;
|
||||
// Min-token cover of a homogeneous run: only the pure-newline
|
||||
// pieces can match (each costing one), and a newline left over
|
||||
// costs one. The lengths are not contiguous (v4.7 jumps 16 → 24),
|
||||
// so this is a coin DP, not a division by the longest piece.
|
||||
let mut best = vec![0u32; m + 1];
|
||||
for end in 1..=m {
|
||||
let mut cost = best[end - 1] + 1;
|
||||
for &len in &self.newline_ladder {
|
||||
let len = len as usize;
|
||||
if len > end {
|
||||
break;
|
||||
}
|
||||
cost = cost.min(best[end - len] + 1);
|
||||
}
|
||||
best[end] = cost;
|
||||
}
|
||||
best[m] - 1
|
||||
}
|
||||
|
||||
/// The frame scalars this vocabulary file measures (v3/v4.7 defaults:
|
||||
/// the frame ends in ⟨bow⟩ and pays the trailing-newline ladder).
|
||||
pub const fn frame_params(&self) -> FrameParams {
|
||||
FrameParams {
|
||||
message_overhead: self.message_overhead,
|
||||
fold_quotes: self.fold_quotes,
|
||||
allcaps_min: self.allcaps_min,
|
||||
frame_bow: true,
|
||||
ladder: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
//! Offline reconstruction of Claude's tokenizer token counts.
|
||||
//!
|
||||
//! Ground-up Rust implementation of the tokenizer-count model reconstructed
|
||||
//! by [ctok](https://github.com/sanderland/ctok): the algorithm port and its
|
||||
//! optimizations are this repository's; the measured vocabulary *data* is
|
||||
//! Sander Land's (MIT — see `data/LICENSE.ctok`), pinned at upstream revision
|
||||
//! `df3b59b` (v1.0.0) and embedded in the front-coded binary form produced by
|
||||
//! `packages/natives/scripts/gen-ctok-vocab.ts`. The research behind the
|
||||
//! model is described in "On the biology of Claude's tokenizer"
|
||||
//! (<https://tokencontributions.substack.com/p/on-the-biology-of-claudes-tokenizer>).
|
||||
//!
|
||||
//! The reconstruction targets counts, not boundaries, and this port keeps
|
||||
//! only the counting surface (ctok's `tokenize()` token list, witness
|
||||
//! verification, and CLI are not ported). For one user message the pipeline
|
||||
//! is:
|
||||
//!
|
||||
//! 1. normalize the text (NFC plus family-specific folding), borrowing it whole
|
||||
//! where no rule can fire;
|
||||
//! 2. rewrite it into a marked byte stream, with word, case and boundary
|
||||
//! markers written in as single bytes;
|
||||
//! 3. min-cost tile the stream over the measured vocabulary and UTF-8 byte
|
||||
//! fallback, matching pieces with one Aho-Corasick transition per byte;
|
||||
//! 4. add the measured message frame.
|
||||
//!
|
||||
//! Nothing in the pipeline materializes decoded characters: the stream, the
|
||||
//! vocabulary and the tiling are all byte-level, and Unicode tables are read
|
||||
//! only for the non-ASCII, non-ideograph characters whose class needs them.
|
||||
//!
|
||||
//! Exactness inherited from upstream: 0 mismatches on ~3.4 M recorded
|
||||
//! `count_tokens` responses across the v3 and v4.7 corpora. The port is
|
||||
//! validated against Python ctok fixtures in `testdata/fixtures.json`.
|
||||
|
||||
mod constants;
|
||||
mod engine;
|
||||
mod normalize;
|
||||
|
||||
use std::sync::LazyLock;
|
||||
|
||||
use engine::VocabCore;
|
||||
use normalize::{FrameParams, nfc, raw_head_space, stream_norm};
|
||||
|
||||
/// One reconstructed tokenizer generation.
|
||||
///
|
||||
/// | family | model generation |
|
||||
/// |---|---|
|
||||
/// | `V3` | Claude 3 through Opus 4.6 (and every non-opus Claude < 5) |
|
||||
/// | `V47` | Opus 4.7 through 4.9 |
|
||||
/// | `V5` | Opus 5 |
|
||||
/// | `V5Sonnet` | Sonnet 5, Fable 5 (non-opus 5-series) |
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
||||
pub enum Family {
|
||||
/// Claude 3 … Opus 4.5/4.6 vocabulary, curly quotes folded, ⟨caps⟩ at 4+.
|
||||
V3,
|
||||
/// Opus 4.7+ vocabulary: quotes literal, ⟨caps⟩ disabled.
|
||||
V47,
|
||||
/// v4.7 vocabulary with the opus-5 message frame: no frame ⟨bow⟩,
|
||||
/// trailing ASCII whitespace absorbed for free, overhead 6.
|
||||
V5,
|
||||
/// v4.7 vocabulary with the sonnet/fable-5 frame: no frame ⟨bow⟩,
|
||||
/// overhead 6, and a trailing-newline ladder — the frame appends nothing
|
||||
/// but absorbs one token of the content-final newline run
|
||||
/// (`tile(run) - 1`). Measured live against `claude-sonnet-5` /
|
||||
/// `claude-fable-5` (2026-08-19; ladder dips pinned at run = 32);
|
||||
/// upstream ctok's "Sonnet 5 counts like Opus 5" does not hold on
|
||||
/// trailing whitespace.
|
||||
V5Sonnet,
|
||||
}
|
||||
|
||||
static CORE_V3: LazyLock<VocabCore> =
|
||||
LazyLock::new(|| VocabCore::parse(include_bytes!("data/ctok_v3.bin")));
|
||||
|
||||
static CORE_V47: LazyLock<VocabCore> =
|
||||
LazyLock::new(|| VocabCore::parse(include_bytes!("data/ctok_v4_7.bin")));
|
||||
|
||||
impl Family {
|
||||
fn core(self) -> &'static VocabCore {
|
||||
match self {
|
||||
Self::V3 => &CORE_V3,
|
||||
Self::V47 | Self::V5 | Self::V5Sonnet => &CORE_V47,
|
||||
}
|
||||
}
|
||||
|
||||
fn params(self) -> FrameParams {
|
||||
let mut p = self.core().frame_params();
|
||||
if matches!(self, Self::V5 | Self::V5Sonnet) {
|
||||
// The 5-series reuses v4.7's vocabulary; only the message frame
|
||||
// differs (ctok's `FAMILIES` overrides plus the live-measured
|
||||
// sonnet/fable split).
|
||||
p.message_overhead = 6;
|
||||
p.frame_bow = false;
|
||||
p.ladder = self == Self::V5Sonnet;
|
||||
}
|
||||
p
|
||||
}
|
||||
|
||||
/// How many newlines the message frame appends after the content on the
|
||||
/// ladder families: one token can span into them, so a content-final
|
||||
/// newline run of `n` is priced as `tile(n + appended) - 1`. v3/v4.7
|
||||
/// frames append ⏎⏎ (free tail dips where `n + 2` is a ladder piece,
|
||||
/// e.g. n = 30); the sonnet/fable-5 frame appends nothing (dips at
|
||||
/// n = 32 itself), both pinned live against the API.
|
||||
const fn appended_newlines(self) -> usize {
|
||||
match self {
|
||||
Self::V5Sonnet => 0,
|
||||
_ => 2,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Token count of `text` as message *content*: the min-cost tiling of the
|
||||
/// marked stream, without the fixed per-message frame. This is the right
|
||||
/// quantity for budget estimates that sum fragments.
|
||||
pub fn content_token_count(text: &str, family: Family) -> u32 {
|
||||
let core = family.core();
|
||||
let p = family.params();
|
||||
if p.ladder {
|
||||
let norm = nfc(text, p.fold_quotes);
|
||||
// The frame appends newline(s) and one token can span into them: read
|
||||
// the content-final newline run before `stream_norm` strips it.
|
||||
let n_tail = norm.bytes().rev().take_while(|&b| b == b'\n').count();
|
||||
let stream = stream_norm(&norm, &p, raw_head_space(text));
|
||||
let tail = core.ladder_tail_cost(n_tail, family.appended_newlines());
|
||||
if stream.is_empty() {
|
||||
return tail;
|
||||
}
|
||||
core.tile_cost(&stream) + tail
|
||||
} else {
|
||||
// The v5 frame absorbs raw ASCII whitespace, so strip before NFC:
|
||||
// NFC folds NBSP etc. to U+0020, and those are not free at the end.
|
||||
let stripped = text.trim_end_matches([' ', '\t', '\n', '\r', '\u{0b}', '\u{0c}']);
|
||||
let norm = nfc(stripped, p.fold_quotes);
|
||||
let stream = stream_norm(&norm, &p, raw_head_space(text));
|
||||
if stream.is_empty() {
|
||||
0
|
||||
} else {
|
||||
core.tile_cost(&stream)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Reconstructed `count_tokens` value for `text` as a single user message:
|
||||
/// content tiling plus the measured message frame (ctok's `token_count`).
|
||||
pub fn message_token_count(text: &str, family: Family) -> u32 {
|
||||
content_token_count(text, family) + family.params().message_overhead
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::{Family, content_token_count, message_token_count};
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Fixture {
|
||||
text: String,
|
||||
v3: u32,
|
||||
v4_7: u32,
|
||||
v5: u32,
|
||||
}
|
||||
|
||||
/// Ground truth recorded from Python ctok 1.0.0 (`token_count(text, v)`
|
||||
/// for v in 3.0/4.7/5.0). Regenerate with a `uv run --with ctok` sweep if
|
||||
/// the vendored vocabulary files are updated.
|
||||
fn fixtures() -> Vec<Fixture> {
|
||||
serde_json::from_str(include_str!("testdata/fixtures.json")).expect("fixtures parse")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_reference_counts() {
|
||||
let mut checked = 0usize;
|
||||
for f in fixtures() {
|
||||
for (family, want) in [(Family::V3, f.v3), (Family::V47, f.v4_7), (Family::V5, f.v5)] {
|
||||
let got = message_token_count(&f.text, family);
|
||||
assert_eq!(got, want, "family {family:?} text {:?}", f.text);
|
||||
checked += 1;
|
||||
}
|
||||
}
|
||||
assert!(checked >= 250, "fixture corpus unexpectedly small: {checked}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_live_sonnet5_counts() {
|
||||
// The sonnet/fable-5 frame diverges from opus-5 on trailing
|
||||
// whitespace and is not modeled by upstream ctok: these rows are raw
|
||||
// `count_tokens` responses recorded from `claude-sonnet-5`
|
||||
// (2026-08-19), whitespace-edge heavy, ladder dips included.
|
||||
#[derive(Deserialize)]
|
||||
struct LiveRow {
|
||||
text: String,
|
||||
count: u32,
|
||||
}
|
||||
let rows: Vec<LiveRow> =
|
||||
serde_json::from_str(include_str!("testdata/sonnet5_live.json")).expect("rows parse");
|
||||
assert!(rows.len() >= 50, "live corpus unexpectedly small: {}", rows.len());
|
||||
for row in rows {
|
||||
assert_eq!(
|
||||
message_token_count(&row.text, Family::V5Sonnet),
|
||||
row.count,
|
||||
"text {:?}",
|
||||
row.text
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn content_count_is_message_minus_frame() {
|
||||
// The public split every consumer relies on: summing fragments must
|
||||
// never include per-message frame overhead.
|
||||
assert_eq!(content_token_count("", Family::V5), 0);
|
||||
for (family, overhead) in
|
||||
[(Family::V3, 7), (Family::V47, 11), (Family::V5, 6), (Family::V5Sonnet, 6)]
|
||||
{
|
||||
let text = "hello, world";
|
||||
assert_eq!(
|
||||
message_token_count(text, family),
|
||||
content_token_count(text, family) + overhead,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// The stream spells markers as bytes no text can produce (`nfc` strips
|
||||
/// every C0 control and folds NUL), so text that contains ctok's own marker
|
||||
/// noncharacters, private-use codepoints, or raw controls tiles as the text
|
||||
/// it is — no escaping stage in between. These rows were verified against
|
||||
/// the previous noncharacter-marker encoder, which did escape them.
|
||||
#[test]
|
||||
fn marker_lookalikes_count_as_text() {
|
||||
let rows: &[(&str, [u32; 4])] = &[
|
||||
("a\u{fdd0}b", [12, 16, 11, 11]),
|
||||
("a\u{fdd1}b", [12, 16, 11, 11]),
|
||||
("\u{fdd4}WORD", [13, 18, 12, 12]),
|
||||
("\u{fdd0}\u{fdd1}\u{fdd2}\u{fdd3}\u{fdd4}", [23, 27, 21, 21]),
|
||||
("hello \u{fdd0}world\u{fdd1} bye", [18, 24, 19, 19]),
|
||||
// Private use is stripped, so both neighbours join one word.
|
||||
("a\u{e000}b\u{e004}c", [8, 13, 8, 8]),
|
||||
("\u{e000}\u{e001}\u{e002}", [8, 12, 6, 6]),
|
||||
// C0 controls are stripped; NUL folds to a space instead.
|
||||
("a\u{01}\u{08}b\u{7f}c", [8, 13, 8, 8]),
|
||||
("a\0b", [9, 13, 8, 8]),
|
||||
];
|
||||
for (text, want) in rows {
|
||||
let families = [Family::V3, Family::V47, Family::V5, Family::V5Sonnet];
|
||||
for (family, &expected) in families.into_iter().zip(want) {
|
||||
assert_eq!(
|
||||
message_token_count(text, family),
|
||||
expected,
|
||||
"family {family:?} text {text:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,912 @@
|
||||
//! Text → the marked stream: everything that happens before the tiling.
|
||||
//!
|
||||
//! Port of ctok's `normalize.py`:
|
||||
//!
|
||||
//! ```text
|
||||
//! NFC + quote fold -> class split -> case marking -> boundary markers written in
|
||||
//! ```
|
||||
//!
|
||||
//! [`stream_norm`] is the output: one byte string of UTF-8 text with word
|
||||
//! boundaries, case and absorbed spaces written in as marker bytes, which
|
||||
//! `engine.rs` then tiles. Every rule here is either a designed rewrite of the
|
||||
//! text or a measured fact about the `count_tokens` oracle. No costs live in
|
||||
//! this module.
|
||||
//!
|
||||
//! Nothing buffers decoded characters: [`nfc`] hands back its input borrowed
|
||||
//! whenever no rule can fire on it (clean ASCII, the common case), runs are
|
||||
//! byte ranges into that text, and a marker costs one byte. Unicode tables are
|
||||
//! consulted only for characters that are neither ASCII nor ideographs.
|
||||
|
||||
use std::borrow::Cow;
|
||||
|
||||
use unicode_normalization::{UnicodeNormalization, char::canonical_combining_class};
|
||||
use unicode_properties::{GeneralCategory, GeneralCategoryGroup, UnicodeGeneralCategory};
|
||||
|
||||
use super::constants::{
|
||||
BOW, CAPS, EOW, NON_SEPARATOR, SHIFT, fold_quote, in_separator_ranges, is_contraction_suffix,
|
||||
is_funny_space, is_punct_sym, is_stripped_control, is_stripped_private, is_symbol_letter,
|
||||
is_variation_selector,
|
||||
};
|
||||
|
||||
/// The stream class of one codepoint.
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
||||
pub enum Class {
|
||||
/// Letters and marks that take the full word model.
|
||||
Wordy,
|
||||
/// Isolated-character path: CJK, astral, and everything unlisted.
|
||||
Hard,
|
||||
/// ASCII and Arabic-Indic decimal digits.
|
||||
Digit,
|
||||
/// Punctuation and symbols that tile over the punct vocabulary.
|
||||
Punct,
|
||||
/// Whitespace.
|
||||
Space,
|
||||
/// A combining mark with no letter in front of it to be its base.
|
||||
StrayMark,
|
||||
}
|
||||
|
||||
/// One maximal same-class run, as a byte range into the normalized text.
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct Run {
|
||||
pub cls: Class,
|
||||
pub start: usize,
|
||||
pub end: usize,
|
||||
}
|
||||
|
||||
/// Frame scalars a family contributes to the encoder (ctok's `TokenizerModel`
|
||||
/// metadata; the vocabulary itself lives in [`super::engine::VocabCore`]).
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct FrameParams {
|
||||
/// Fixed cost of a single user message before its content.
|
||||
pub message_overhead: u32,
|
||||
/// v3 alone folds curly quotes to their ASCII forms.
|
||||
pub fold_quotes: bool,
|
||||
/// Minimum span length for the ⟨caps⟩ marker; `None` disables it (v4.7+).
|
||||
pub allcaps_min: Option<usize>,
|
||||
/// Whether the frame ends in a ⟨bow⟩ that absorbs one leading space.
|
||||
pub frame_bow: bool,
|
||||
/// `true` for the measured trailing-newline ladder (v3/v4.7); `false` when
|
||||
/// the frame absorbs all trailing ASCII whitespace (v5).
|
||||
pub ladder: bool,
|
||||
}
|
||||
|
||||
// Python 3.13 ships Unicode 15.1 data; the source models know these Unicode
|
||||
// 16.0 case pairs. Kept explicit so the port does not depend on the std
|
||||
// Unicode revision.
|
||||
const NEW_CASE_PAIRS: [(char, char); 2] = [('\u{1c89}', '\u{1c8a}'), ('\u{a7cb}', '\u{0264}')];
|
||||
|
||||
fn new_case_lower(c: char) -> Option<char> {
|
||||
NEW_CASE_PAIRS
|
||||
.iter()
|
||||
.find(|&&(u, _)| u == c)
|
||||
.map(|&(_, l)| l)
|
||||
}
|
||||
|
||||
fn is_new_cased(c: char) -> bool {
|
||||
NEW_CASE_PAIRS.iter().any(|&(u, l)| u == c || l == c)
|
||||
}
|
||||
|
||||
fn is_upper_x(c: char) -> bool {
|
||||
new_case_lower(c).is_some() || c.is_uppercase()
|
||||
}
|
||||
|
||||
fn is_lower_x(c: char) -> bool {
|
||||
NEW_CASE_PAIRS.iter().any(|&(_, l)| l == c) || c.is_lowercase()
|
||||
}
|
||||
|
||||
/// Whether the single-character full lowering of `c` is `c` itself.
|
||||
fn lowers_to_self(c: char) -> bool {
|
||||
if new_case_lower(c).is_some() {
|
||||
return false;
|
||||
}
|
||||
let mut it = c.to_lowercase();
|
||||
it.next() == Some(c) && it.next().is_none()
|
||||
}
|
||||
|
||||
/// Append one character's UTF-8 bytes.
|
||||
#[inline]
|
||||
fn push_char(c: char, out: &mut Vec<u8>) {
|
||||
let mut buf = [0u8; 4];
|
||||
out.extend_from_slice(c.encode_utf8(&mut buf).as_bytes());
|
||||
}
|
||||
|
||||
fn push_lower(c: char, out: &mut Vec<u8>) {
|
||||
match new_case_lower(c) {
|
||||
Some(l) => push_char(l, out),
|
||||
None => {
|
||||
for l in c.to_lowercase() {
|
||||
push_char(l, out);
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether [`nfc`] rewrites this ASCII byte: the stripped controls, plus NUL,
|
||||
/// which folds to a space.
|
||||
#[inline]
|
||||
const fn ascii_needs_fold(b: u8) -> bool {
|
||||
matches!(b, 0x00..=0x08 | 0x0b..=0x1f | 0x7f)
|
||||
}
|
||||
|
||||
/// Apply the measured text normalization: NFC, control/private-use stripping,
|
||||
/// NUL and exotic-space folding, Thai SARA AM composition, and (v3) quote
|
||||
/// folding. Text no rule touches is returned borrowed. Lone surrogates cannot
|
||||
/// exist in a Rust `str`; the JS boundary has already folded them to U+FFFD
|
||||
/// exactly as the Python port does.
|
||||
pub fn nfc(text: &str, fold_quotes: bool) -> Cow<'_, str> {
|
||||
if text.is_ascii() {
|
||||
// ASCII is closed under every rule here except control stripping, and
|
||||
// is always in NFC.
|
||||
if !text.bytes().any(ascii_needs_fold) {
|
||||
return Cow::Borrowed(text);
|
||||
}
|
||||
let mut out = String::with_capacity(text.len());
|
||||
for b in text.bytes() {
|
||||
match b {
|
||||
0x00 => out.push(' '),
|
||||
b if ascii_needs_fold(b) => {},
|
||||
b => out.push(char::from(b)),
|
||||
}
|
||||
}
|
||||
return Cow::Owned(out);
|
||||
}
|
||||
let mut out = String::with_capacity(text.len());
|
||||
for c in text.chars().nfc() {
|
||||
if is_stripped_control(c) {
|
||||
continue;
|
||||
}
|
||||
let c = if c == '\0' { ' ' } else { c };
|
||||
// Claude composes decomposed Thai SARA AM, whose compatibility
|
||||
// decomposition NFC leaves alone. Lao SARA AM does not fold.
|
||||
if c == '\u{0e32}' && out.ends_with('\u{0e4d}') {
|
||||
out.pop();
|
||||
out.push('\u{0e33}');
|
||||
continue;
|
||||
}
|
||||
if is_stripped_private(c) {
|
||||
continue;
|
||||
}
|
||||
let c = if fold_quotes { fold_quote(c) } else { c };
|
||||
out.push(if is_funny_space(c) { ' ' } else { c });
|
||||
}
|
||||
Cow::Owned(out)
|
||||
}
|
||||
|
||||
/// Whether a codepoint uses the isolated character path for letters.
|
||||
fn is_hard_cp(o: u32) -> bool {
|
||||
o >= 0x10000 // all astral (CJK ext, astral scripts, emoji)
|
||||
|| (0x4e00..=0x9fff).contains(&o) // CJK Unified
|
||||
|| (0x3400..=0x4dbf).contains(&o) // CJK Ext A
|
||||
|| (0xf900..=0xfaff).contains(&o) // CJK Compatibility
|
||||
|| (0xac00..=0xd7a3).contains(&o) // Hangul syllables
|
||||
// The ideographic iteration and closing marks (々 〆) continue the Han
|
||||
// run they follow despite being letters by category.
|
||||
|| o == 0x3005 || o == 0x3006
|
||||
// Quranic annotation signs: measured as unattached marks.
|
||||
|| (0x06dd..=0x06e0).contains(&o)
|
||||
|| (0x06e9..=0x06ec).contains(&o)
|
||||
}
|
||||
|
||||
/// A mark that terminates the orthographic syllable, and so separates word
|
||||
/// runs: viramas (ccc 9, minus THAI PHINTHU) plus the measured ranges and
|
||||
/// enumerated signs in `constants.rs`.
|
||||
pub fn is_separator(c: char) -> bool {
|
||||
// No ASCII codepoint combines, and the measured ranges start at U+0300.
|
||||
!c.is_ascii() && is_separator_general(c)
|
||||
}
|
||||
|
||||
fn is_separator_general(c: char) -> bool {
|
||||
(canonical_combining_class(c) == 9 && c != NON_SEPARATOR) || in_separator_ranges(c as u32)
|
||||
}
|
||||
|
||||
/// Stream class of every ASCII codepoint: the fast path of [`classify`], since
|
||||
/// no rule below the 0x80 boundary needs a Unicode table. Pinned against the
|
||||
/// general path by `ascii_class_table_matches_general`.
|
||||
static ASCII_CLASS: [Class; 128] = {
|
||||
let mut table = [Class::Hard; 128];
|
||||
let mut i = 0usize;
|
||||
while i < 128 {
|
||||
table[i] = match i as u8 {
|
||||
b'\t' | b'\n' | 0x0b | 0x0c | b'\r' | b' ' => Class::Space,
|
||||
b'0'..=b'9' => Class::Digit,
|
||||
b'A'..=b'Z' | b'a'..=b'z' => Class::Wordy,
|
||||
// Every ASCII printable that is neither alphanumeric nor a space
|
||||
// is gc=P* or gc=S*, and both tile over the merged punct
|
||||
// vocabulary, so operators (`==`, `=>`, `});`) form one PUNCT run.
|
||||
0x21..=0x2f | 0x3a..=0x40 | 0x5b..=0x60 | 0x7b..=0x7e => Class::Punct,
|
||||
// C0 controls and DEL: stripped before classification ever sees
|
||||
// them, and HARD by the general path in any case.
|
||||
_ => Class::Hard,
|
||||
};
|
||||
i += 1;
|
||||
}
|
||||
table
|
||||
};
|
||||
|
||||
pub fn classify(c: char) -> Class {
|
||||
if c.is_ascii() {
|
||||
return ASCII_CLASS[c as usize];
|
||||
}
|
||||
if is_ideograph(c as u32) {
|
||||
return Class::Hard;
|
||||
}
|
||||
classify_nonascii(c)
|
||||
}
|
||||
|
||||
/// Han ideographs and Hangul syllables, which the isolated-character path
|
||||
/// claims wholesale: gc=Lo throughout, so they take the HARD class, no border
|
||||
/// marker and no digit border, with no Unicode lookup needed. Pinned against
|
||||
/// the general path by `ideograph_shortcut_matches_general`.
|
||||
const fn is_ideograph(o: u32) -> bool {
|
||||
matches!(o, 0x3400..=0x4dbf | 0x4e00..=0x9fff | 0xac00..=0xd7a3)
|
||||
}
|
||||
|
||||
/// The stream class of one codepoint, derived from Unicode data and the
|
||||
/// measured tables.
|
||||
fn classify_nonascii(c: char) -> Class {
|
||||
if is_separator(c) {
|
||||
return Class::Hard; // separates word runs; `marks_like_punct` claims its borders
|
||||
}
|
||||
if is_new_cased(c) {
|
||||
return Class::Wordy;
|
||||
}
|
||||
let o = c as u32;
|
||||
let group = c.general_category_group();
|
||||
if group == GeneralCategoryGroup::Separator
|
||||
|| matches!(c, '\t' | '\n' | '\r' | '\u{0b}' | '\u{0c}')
|
||||
{
|
||||
return Class::Space;
|
||||
}
|
||||
if is_symbol_letter(o) {
|
||||
return Class::Wordy;
|
||||
}
|
||||
let cat = c.general_category();
|
||||
if cat == GeneralCategory::DecimalNumber
|
||||
&& (o < 0x80 || (0x0660..=0x0669).contains(&o) || (0x06f0..=0x06f9).contains(&o))
|
||||
{
|
||||
return Class::Digit;
|
||||
}
|
||||
if o < 0x80 && matches!(group, GeneralCategoryGroup::Punctuation | GeneralCategoryGroup::Symbol)
|
||||
{
|
||||
return Class::Punct;
|
||||
}
|
||||
if is_punct_sym(c) {
|
||||
return Class::Punct;
|
||||
}
|
||||
if is_variation_selector(c) {
|
||||
return Class::Hard; // gc=Mn, but they take no word model at all
|
||||
}
|
||||
if o == 0x0cf3 {
|
||||
// KANNADA SIGN COMBINING ANUSVARA ABOVE RIGHT: measured as plain word
|
||||
// material; pinned so older Unicode tables cannot drop it to HARD.
|
||||
return Class::Wordy;
|
||||
}
|
||||
if matches!(group, GeneralCategoryGroup::Letter | GeneralCategoryGroup::Mark) && !is_hard_cp(o) {
|
||||
return Class::Wordy;
|
||||
}
|
||||
Class::Hard
|
||||
}
|
||||
|
||||
/// A Syriac vowel point or superscript alaph that acts as a word-forming
|
||||
/// letter wherever no base can hold it.
|
||||
fn is_syriac_vowel(c: char) -> bool {
|
||||
let o = c as u32;
|
||||
o == 0x0711 || (0x0730..=0x073f).contains(&o)
|
||||
}
|
||||
|
||||
/// A combining mark, asked at a position where nothing before it can be its
|
||||
/// base. Fires only where a mark's base is not a letter.
|
||||
fn is_stray_mark(c: char) -> bool {
|
||||
!c.is_ascii() && is_stray_mark_general(c) // no ASCII codepoint combines
|
||||
}
|
||||
|
||||
fn is_stray_mark_general(c: char) -> bool {
|
||||
if is_syriac_vowel(c) {
|
||||
return false; // a baseless Syriac vowel is a word-forming letter instead
|
||||
}
|
||||
canonical_combining_class(c) != 0 && !is_separator(c)
|
||||
}
|
||||
|
||||
/// Whether a digit receives a border marker: ASCII digits take none, every
|
||||
/// other BMP decimal or other-number digit does. Astral digits take none.
|
||||
fn digit_border(c: char) -> bool {
|
||||
// ASCII digits take no border, and ASCII holds no gc=No.
|
||||
!c.is_ascii() && digit_border_general(c)
|
||||
}
|
||||
|
||||
fn digit_border_general(c: char) -> bool {
|
||||
let cat = c.general_category();
|
||||
(cat == GeneralCategory::OtherNumber || (cat == GeneralCategory::DecimalNumber && !c.is_ascii()))
|
||||
&& (c as u32) < 0x10000
|
||||
}
|
||||
|
||||
/// A uniform run of decimal digits (Nd) or other numbers (No).
|
||||
fn is_digit_run(body: &str) -> bool {
|
||||
if body.is_ascii() {
|
||||
// Nd is the only number category in ASCII, so a uniform run there is a
|
||||
// run of ASCII digits.
|
||||
return !body.is_empty() && body.bytes().all(|b| b.is_ascii_digit());
|
||||
}
|
||||
is_digit_run_general(body)
|
||||
}
|
||||
|
||||
fn is_digit_run_general(body: &str) -> bool {
|
||||
let mut chars = body.chars();
|
||||
let Some(first) = chars.next() else {
|
||||
return false;
|
||||
};
|
||||
let category = first.general_category();
|
||||
if !matches!(category, GeneralCategory::DecimalNumber | GeneralCategory::OtherNumber) {
|
||||
return false;
|
||||
}
|
||||
chars.all(|c| c.general_category() == category)
|
||||
}
|
||||
|
||||
/// Does a digit run write ⟨bow⟩ where it borders a space on the left? The
|
||||
/// run's first character decides.
|
||||
fn digit_bow(body: &str) -> bool {
|
||||
is_digit_run(body) && body.chars().next().is_some_and(digit_border)
|
||||
}
|
||||
|
||||
/// Does a digit run write ⟨eow⟩ where it borders a single space on the right?
|
||||
/// The run's last character decides.
|
||||
fn digit_eow(body: &str) -> bool {
|
||||
is_digit_run(body) && body.chars().next_back().is_some_and(digit_border)
|
||||
}
|
||||
|
||||
/// Whether the border-marker branch can claim this character: punctuation,
|
||||
/// symbols, format characters, unassigned codepoints, and separator marks —
|
||||
/// excluding ideographic punctuation and everything astral.
|
||||
fn marks_like_punct(c: char) -> bool {
|
||||
if c.is_ascii() {
|
||||
// ASCII holds no separator mark, no gc=Cf and no unassigned codepoint,
|
||||
// so this is exactly the ASCII half of the PUNCT class.
|
||||
return ASCII_CLASS[c as usize] == Class::Punct;
|
||||
}
|
||||
if is_ideograph(c as u32) {
|
||||
return false; // gc=Lo: no border marker of its own
|
||||
}
|
||||
marks_like_punct_general(c)
|
||||
}
|
||||
|
||||
fn marks_like_punct_general(c: char) -> bool {
|
||||
let o = c as u32;
|
||||
if o >= 0x10000 {
|
||||
return false;
|
||||
}
|
||||
if is_separator(c) {
|
||||
return true;
|
||||
}
|
||||
// gc=L*, gc=M*, gc=N* and gc=Z* are never gc=Cf or unassigned, so the
|
||||
// second category lookup is needed for gc=C* alone.
|
||||
match c.general_category_group() {
|
||||
GeneralCategoryGroup::Punctuation => !(0x3001..=0x303f).contains(&o),
|
||||
GeneralCategoryGroup::Symbol => true,
|
||||
GeneralCategoryGroup::Other => {
|
||||
matches!(c.general_category(), GeneralCategory::Format | GeneralCategory::Unassigned)
|
||||
},
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether the first character of a hard run takes a left border marker.
|
||||
fn hard_bow(body: &str) -> bool {
|
||||
body
|
||||
.chars()
|
||||
.next()
|
||||
.is_some_and(|c| is_variation_selector(c) || marks_like_punct(c))
|
||||
}
|
||||
|
||||
/// Whether the last character of a hard run takes a right border marker.
|
||||
fn hard_eow(body: &str) -> bool {
|
||||
body
|
||||
.chars()
|
||||
.next_back()
|
||||
.is_some_and(|c| is_variation_selector(c) || marks_like_punct(c))
|
||||
}
|
||||
|
||||
/// Which pretoken alternative a character of a HARD run belongs to. A HARD run
|
||||
/// is a run of our own class, not a pretoken: each sub-run takes its own
|
||||
/// border markers.
|
||||
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
|
||||
enum HardKind {
|
||||
Punct,
|
||||
Number,
|
||||
Letter,
|
||||
}
|
||||
|
||||
fn hard_kind(c: char) -> HardKind {
|
||||
if is_ideograph(c as u32) {
|
||||
// The common HARD run: every character is a letter, so the sub-run
|
||||
// never breaks and no Unicode table is consulted.
|
||||
return HardKind::Letter;
|
||||
}
|
||||
hard_kind_general(c)
|
||||
}
|
||||
|
||||
fn hard_kind_general(c: char) -> HardKind {
|
||||
if marks_like_punct_general(c) {
|
||||
HardKind::Punct
|
||||
} else if digit_border_general(c) {
|
||||
HardKind::Number
|
||||
} else {
|
||||
HardKind::Letter
|
||||
}
|
||||
}
|
||||
|
||||
/// Close one class run, splitting a HARD run where the pretoken kind changes.
|
||||
/// A variation selector never opens a sub-run — it rides its base's sub-run,
|
||||
/// or `⚖️` would sever at the selector and lose its ⟨eow⟩.
|
||||
fn push_run(runs: &mut Vec<Run>, s: &str, cls: Class, start: usize, end: usize) {
|
||||
if cls != Class::Hard {
|
||||
runs.push(Run { cls, start, end });
|
||||
return;
|
||||
}
|
||||
let body = &s[start..end];
|
||||
let mut chars = body.char_indices();
|
||||
let Some((_, first)) = chars.next() else {
|
||||
return;
|
||||
};
|
||||
let mut sub_start = start;
|
||||
let mut kind = hard_kind(first);
|
||||
for (off, ch) in chars {
|
||||
if is_variation_selector(ch) || hard_kind(ch) == kind {
|
||||
continue;
|
||||
}
|
||||
runs.push(Run { cls: Class::Hard, start: sub_start, end: start + off });
|
||||
sub_start = start + off;
|
||||
kind = hard_kind(ch);
|
||||
}
|
||||
runs.push(Run { cls: Class::Hard, start: sub_start, end });
|
||||
}
|
||||
|
||||
/// The text split into maximal same-class runs, with terminal marks as
|
||||
/// unmarked separators and HARD runs split where the pretoken kind changes.
|
||||
fn split_runs(s: &str) -> Vec<Run> {
|
||||
// Runs average a few bytes each: one allocation covers the typical text.
|
||||
let mut runs: Vec<Run> = Vec::with_capacity(s.len() / 3 + 8);
|
||||
let mut chars = s.char_indices();
|
||||
let Some((_, first)) = chars.next() else {
|
||||
return runs;
|
||||
};
|
||||
let mut start = 0usize;
|
||||
let mut cls = classify(first);
|
||||
if cls == Class::Wordy && is_stray_mark(first) {
|
||||
cls = Class::StrayMark; // nothing in front of it, so no letter can be its base
|
||||
}
|
||||
for (i, ch) in chars {
|
||||
let c = classify(ch);
|
||||
// Consecutive unattached marks are one regex-style run; an accent
|
||||
// riding a separator run opens a stray word like any other.
|
||||
if cls == Class::StrayMark && c == Class::Wordy && is_stray_mark(ch) {
|
||||
continue;
|
||||
}
|
||||
if c == cls {
|
||||
continue;
|
||||
}
|
||||
if c == Class::Wordy && cls != Class::Wordy && is_stray_mark(ch) {
|
||||
push_run(&mut runs, s, cls, start, i);
|
||||
start = i;
|
||||
cls = Class::StrayMark;
|
||||
continue;
|
||||
}
|
||||
push_run(&mut runs, s, cls, start, i);
|
||||
start = i;
|
||||
cls = c;
|
||||
}
|
||||
push_run(&mut runs, s, cls, start, s.len());
|
||||
runs
|
||||
}
|
||||
|
||||
/// How a cased span is written into the stream (ctok's `mark_case`).
|
||||
enum CaseForm {
|
||||
/// The span keeps its bytes.
|
||||
Literal,
|
||||
/// ⟨shift⟩ + fully lowered span.
|
||||
Shift,
|
||||
/// ⟨shift⟩ + lowered span with İ kept literal (İ is transparent to the
|
||||
/// title-case test and literal in the lowered body).
|
||||
ShiftKeepDotted,
|
||||
/// ⟨caps⟩ + fully lowered span.
|
||||
Caps,
|
||||
}
|
||||
|
||||
fn span_is_upper(span: &str) -> bool {
|
||||
let mut saw_cased = false;
|
||||
for c in span.chars() {
|
||||
if is_upper_x(c) {
|
||||
saw_cased = true;
|
||||
} else if is_lower_x(c) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
saw_cased
|
||||
}
|
||||
|
||||
/// Cased span → its marked form. A case marker fires only on a whole span:
|
||||
/// pure all-caps of length ≥ `allcaps_min` becomes ⟨caps⟩ + lowercase, pure
|
||||
/// title-case becomes ⟨shift⟩ + lowercase; everything else stays literal
|
||||
/// (`GaN`/`WiFi`/`QQ` keep their bytes). `head_mark` means an unattached mark
|
||||
/// opened the word, so neither marker can assert its lowered first letter.
|
||||
fn case_form(span: &str, allcaps_min: Option<usize>, head_mark: bool) -> CaseForm {
|
||||
if head_mark {
|
||||
return CaseForm::Literal;
|
||||
}
|
||||
if span.is_ascii() {
|
||||
// A WORDY run of ASCII is letters only: ẞ and İ cannot appear, and
|
||||
// every ASCII letter has a distinct lowering, so no character can
|
||||
// block a marker. One pass settles both span shapes.
|
||||
let bytes = span.as_bytes();
|
||||
let mut any_upper = false;
|
||||
let mut any_lower = false;
|
||||
let mut upper_after_first = false;
|
||||
for (i, &b) in bytes.iter().enumerate() {
|
||||
if b.is_ascii_uppercase() {
|
||||
any_upper = true;
|
||||
upper_after_first |= i > 0;
|
||||
} else if b.is_ascii_lowercase() {
|
||||
any_lower = true;
|
||||
}
|
||||
}
|
||||
if allcaps_min.is_some_and(|min| span.len() >= min && any_upper && !any_lower) {
|
||||
return CaseForm::Caps;
|
||||
}
|
||||
if any_upper && !upper_after_first && bytes[0].is_ascii_uppercase() {
|
||||
return CaseForm::Shift;
|
||||
}
|
||||
return CaseForm::Literal;
|
||||
}
|
||||
if span.contains('ẞ') {
|
||||
return CaseForm::Literal;
|
||||
}
|
||||
let first = span.chars().next().expect("run bodies are non-empty");
|
||||
if span.contains('İ') {
|
||||
if is_upper_x(first)
|
||||
&& first != 'İ'
|
||||
&& !span.chars().skip(1).any(|c| c != 'İ' && is_upper_x(c))
|
||||
{
|
||||
return CaseForm::ShiftKeepDotted;
|
||||
}
|
||||
return CaseForm::Literal;
|
||||
}
|
||||
// A character with no lowered form blocks either marker: a marker in
|
||||
// front of an unchanged body over-counts.
|
||||
let unlowerable = span.chars().any(|c| {
|
||||
(is_new_cased(c)
|
||||
|| matches!(
|
||||
c.general_category_group(),
|
||||
GeneralCategoryGroup::Letter | GeneralCategoryGroup::Mark
|
||||
)) && !is_lower_x(c)
|
||||
&& (!is_upper_x(c) || lowers_to_self(c))
|
||||
});
|
||||
if let Some(min) = allcaps_min
|
||||
&& span.chars().count() >= min
|
||||
&& span_is_upper(span)
|
||||
&& !unlowerable
|
||||
{
|
||||
return CaseForm::Caps;
|
||||
}
|
||||
if lowers_to_self(first) {
|
||||
return CaseForm::Literal;
|
||||
}
|
||||
if is_upper_x(first) && !span.chars().skip(1).any(is_upper_x) {
|
||||
return CaseForm::Shift;
|
||||
}
|
||||
CaseForm::Literal
|
||||
}
|
||||
|
||||
fn emit_case_body(span: &str, form: &CaseForm, out: &mut Vec<u8>) {
|
||||
match form {
|
||||
CaseForm::Literal => out.extend_from_slice(span.as_bytes()),
|
||||
// Per-char lowering never applies Final_Sigma, so Σ lowers to σ
|
||||
// everywhere — exactly the oracle's ⟨caps⟩ body spelling.
|
||||
CaseForm::Shift | CaseForm::Caps => {
|
||||
if span.is_ascii() {
|
||||
out.extend(span.bytes().map(|b| b.to_ascii_lowercase()));
|
||||
} else {
|
||||
for c in span.chars() {
|
||||
push_lower(c, out);
|
||||
}
|
||||
}
|
||||
},
|
||||
CaseForm::ShiftKeepDotted => {
|
||||
let mut chars = span.chars();
|
||||
push_lower(chars.next().expect("run bodies are non-empty"), out);
|
||||
for c in chars {
|
||||
if c == 'İ' {
|
||||
push_char(c, out);
|
||||
} else {
|
||||
push_lower(c, out);
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether run `i` is a lone `'` that opens the word after it (`a 'b`,
|
||||
/// `'First`, `x 'REXX`). Only a punct run that is exactly `'` qualifies.
|
||||
const fn opens_word(s: &str, runs: &[Run], i: usize) -> bool {
|
||||
let r = &runs[i];
|
||||
r.end - r.start == 1
|
||||
&& s.as_bytes()[r.start] == b'\''
|
||||
&& i + 1 < runs.len()
|
||||
&& matches!(runs[i + 1].cls, Class::Wordy | Class::StrayMark)
|
||||
}
|
||||
|
||||
/// Does this run write a boundary marker of its own on its right edge?
|
||||
fn takes_right_border(s: &str, run: &Run) -> bool {
|
||||
let body = &s[run.start..run.end];
|
||||
run.cls == Class::Punct
|
||||
|| hard_eow(body)
|
||||
|| (matches!(run.cls, Class::Digit | Class::Hard) && is_digit_run(body) && digit_eow(body))
|
||||
}
|
||||
|
||||
/// Does a lone apostrophe immediately left of wordy run `i` supply that word's
|
||||
/// ⟨bow⟩? `it's` is one word boundary, not two. Three conditions, all
|
||||
/// measured: the suffix is a contraction suffix (whole-word, lowercase); the
|
||||
/// apostrophe is a punct run of its own; and the run on the far side of the
|
||||
/// apostrophe writes no right-hand border marker of its own.
|
||||
fn contraction_seam(s: &str, runs: &[Run], i: usize) -> bool {
|
||||
if i == 0 {
|
||||
return false;
|
||||
}
|
||||
let r = &runs[i];
|
||||
if !is_contraction_suffix(&s.as_bytes()[r.start..r.end]) {
|
||||
return false;
|
||||
}
|
||||
let prev = &runs[i - 1];
|
||||
if prev.cls != Class::Punct || prev.end - prev.start != 1 || s.as_bytes()[prev.start] != b'\'' {
|
||||
return false;
|
||||
}
|
||||
i < 2 || !takes_right_border(s, &runs[i - 2])
|
||||
}
|
||||
|
||||
/// Whether raw (pre-normalization) text supplies the leading space the frame
|
||||
/// absorbs. A space a fold produced or exposed is not absorbed.
|
||||
pub fn raw_head_space(text: &str) -> bool {
|
||||
text.starts_with(' ')
|
||||
}
|
||||
|
||||
/// Byte offset where the character before `at` starts (`at` is a character
|
||||
/// boundary with at least one character before it).
|
||||
#[inline]
|
||||
const fn prev_char_start(out: &[u8], at: usize) -> usize {
|
||||
let mut start = at - 1;
|
||||
while out[start] & 0xc0 == 0x80 {
|
||||
start -= 1;
|
||||
}
|
||||
start
|
||||
}
|
||||
|
||||
/// Write ⟨bow⟩, applying the seam law: in `⟨eow⟩ ' ' [case markers] ⟨bow⟩` the
|
||||
/// space is not written as a character. `guard` is the byte length of `out`
|
||||
/// just past the last dropped seam's ⟨bow⟩; the character before the ⟨eow⟩ has
|
||||
/// to start at or after it, which is what keeps the rewrite non-overlapping
|
||||
/// and left-to-right, exactly as `re.sub` scans.
|
||||
fn push_bow(out: &mut Vec<u8>, guard: &mut usize) {
|
||||
let mut case_at = out.len();
|
||||
while case_at >= 1 && matches!(out[case_at - 1], SHIFT | CAPS) {
|
||||
case_at -= 1;
|
||||
}
|
||||
// `case_at` is where this run's case markers begin: a seam needs ⟨eow⟩ and
|
||||
// a space in front of them, plus a character before the ⟨eow⟩ that no
|
||||
// earlier seam has consumed.
|
||||
let seam = case_at >= 3
|
||||
&& out[case_at - 1] == b' '
|
||||
&& out[case_at - 2] == EOW
|
||||
&& prev_char_start(out, case_at - 2) >= *guard;
|
||||
if seam {
|
||||
// Drop the space, sliding this run's case markers down over it.
|
||||
out.copy_within(case_at.., case_at - 1);
|
||||
out.truncate(out.len() - 1);
|
||||
}
|
||||
out.push(BOW);
|
||||
if seam {
|
||||
*guard = out.len();
|
||||
}
|
||||
}
|
||||
|
||||
/// The marked stream over already-normalized text, in the internal marked
|
||||
/// form. A WORDY run is bracketed by ⟨bow⟩/⟨eow⟩ and case-normalized; the
|
||||
/// ⟨eow⟩⟨bow⟩ seam encodes a single space between two such runs; punct, digit,
|
||||
/// separator and stray-mark runs receive their measured boundary markers.
|
||||
pub fn stream_norm(norm: &str, p: &FrameParams, raw_head_space: bool) -> Vec<u8> {
|
||||
let mut s = norm;
|
||||
// The frame's tail for a ladder family; a "free" family was stripped on
|
||||
// the raw text instead (see `content_token_count`).
|
||||
if p.ladder {
|
||||
s = s.trim_end_matches('\n');
|
||||
}
|
||||
// A single leading space is dropped where the frame ends in ⟨bow⟩: that
|
||||
// ⟨bow⟩ is the space (' a' = 1). Two or more are a whitespace-run token
|
||||
// and stay.
|
||||
if p.frame_bow
|
||||
&& raw_head_space
|
||||
&& s.as_bytes().first() == Some(&b' ')
|
||||
&& s.as_bytes().get(1) != Some(&b' ')
|
||||
{
|
||||
s = &s[1..];
|
||||
}
|
||||
|
||||
let runs = split_runs(s);
|
||||
if runs.is_empty() {
|
||||
// Content that normalizes away entirely still pays for the frame's
|
||||
// ⟨bow⟩; with nothing to attach to it, it tiles as itself.
|
||||
return if p.frame_bow { vec![BOW] } else { Vec::new() };
|
||||
}
|
||||
let bytes = s.as_bytes();
|
||||
|
||||
let borders_space = |i: usize, side: isize| -> bool {
|
||||
let at = i as isize + side;
|
||||
if at < 0 {
|
||||
return p.frame_bow; // message start counts: the frame ends in ⟨bow⟩, which is a space
|
||||
}
|
||||
let at = at as usize;
|
||||
if at >= runs.len() {
|
||||
return false; // message end does not: the trailing frame is not a space
|
||||
}
|
||||
let r = &runs[at];
|
||||
if r.cls != Class::Space {
|
||||
return false;
|
||||
}
|
||||
if side < 0 {
|
||||
bytes[r.end - 1] == b' '
|
||||
} else {
|
||||
// Run-kills-marker: a right-hand marker is written for the seam
|
||||
// space only, never before a run of two or more spaces.
|
||||
bytes[r.start] == b' ' && (r.end - r.start < 2 || bytes[r.start + 1] != b' ')
|
||||
}
|
||||
};
|
||||
|
||||
let mut out: Vec<u8> = Vec::with_capacity(s.len() + s.len() / 2 + 16);
|
||||
// Byte length of `out` just past the last seam-dropped space's ⟨bow⟩.
|
||||
let mut guard = 0usize;
|
||||
let head_quote = opens_word(s, &runs, 0);
|
||||
let first = &runs[0];
|
||||
let first_body = &s[first.start..first.end];
|
||||
let has_own_bow = !head_quote
|
||||
&& (matches!(first.cls, Class::Wordy | Class::Punct | Class::StrayMark)
|
||||
|| hard_bow(first_body)
|
||||
|| (matches!(first.cls, Class::Digit | Class::Hard) && digit_bow(first_body))
|
||||
|| (first.cls == Class::Space && bytes[first.start] == b' '));
|
||||
if !has_own_bow && p.frame_bow {
|
||||
// The frame ends in ⟨bow⟩, always; a run that supplies none leaves it
|
||||
// to be written here, where it tiles as itself. A word-opening `'`
|
||||
// receives the frame's space as the character it is.
|
||||
if head_quote {
|
||||
out.push(b' ');
|
||||
} else {
|
||||
out.push(BOW);
|
||||
}
|
||||
}
|
||||
|
||||
for i in 0..runs.len() {
|
||||
let r = &runs[i];
|
||||
let body = &s[r.start..r.end];
|
||||
match r.cls {
|
||||
Class::Wordy => {
|
||||
// Flanked on both sides, except where a contraction apostrophe
|
||||
// is already its opening boundary or an unattached mark run
|
||||
// already opened this word.
|
||||
let fused = i > 0 && runs[i - 1].cls == Class::StrayMark;
|
||||
let form = case_form(body, p.allcaps_min, fused);
|
||||
match form {
|
||||
CaseForm::Shift | CaseForm::ShiftKeepDotted => out.push(SHIFT),
|
||||
CaseForm::Caps => out.push(CAPS),
|
||||
CaseForm::Literal => {},
|
||||
}
|
||||
if !(fused || contraction_seam(s, &runs, i)) {
|
||||
push_bow(&mut out, &mut guard);
|
||||
}
|
||||
emit_case_body(body, &form, &mut out);
|
||||
out.push(EOW);
|
||||
},
|
||||
Class::StrayMark => {
|
||||
// A stray-mark pretoken is a word: ⟨bow⟩ on the left always,
|
||||
// ⟨eow⟩ on the right against everything except a letter, which
|
||||
// is the rest of the same word.
|
||||
push_bow(&mut out, &mut guard);
|
||||
out.extend_from_slice(body.as_bytes());
|
||||
let letter_follows = i + 1 < runs.len() && runs[i + 1].cls == Class::Wordy;
|
||||
if !letter_follows {
|
||||
out.push(EOW);
|
||||
}
|
||||
},
|
||||
_ if r.cls == Class::Punct || hard_bow(body) || hard_eow(body) => {
|
||||
// A punct span is marked only on the side that borders
|
||||
// whitespace: `a! b` gets `!⟨eow⟩`, `a!b` a bare `!`.
|
||||
let takes_bow = borders_space(i, -1)
|
||||
&& !opens_word(s, &runs, i)
|
||||
&& (r.cls == Class::Punct || hard_bow(body));
|
||||
if takes_bow {
|
||||
push_bow(&mut out, &mut guard);
|
||||
}
|
||||
out.extend_from_slice(body.as_bytes());
|
||||
if borders_space(i, 1) && (r.cls == Class::Punct || hard_eow(body)) {
|
||||
out.push(EOW);
|
||||
}
|
||||
},
|
||||
Class::Digit | Class::Hard if is_digit_run(body) => {
|
||||
// Same border markers as punctuation, decided per border
|
||||
// character; deliberately no lookback across the space.
|
||||
if digit_bow(body) && borders_space(i, -1) {
|
||||
push_bow(&mut out, &mut guard);
|
||||
}
|
||||
out.extend_from_slice(body.as_bytes());
|
||||
if digit_eow(body) && borders_space(i, 1) {
|
||||
out.push(EOW);
|
||||
}
|
||||
},
|
||||
_ => out.extend_from_slice(body.as_bytes()), // HARD letter scripts and whitespace
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
ASCII_CLASS, Class, classify, classify_nonascii, digit_border, digit_border_general,
|
||||
hard_kind, hard_kind_general, is_digit_run, is_digit_run_general, is_separator,
|
||||
is_separator_general, is_stray_mark, is_stray_mark_general, marks_like_punct,
|
||||
marks_like_punct_general,
|
||||
};
|
||||
|
||||
/// Every ASCII shortcut has to answer exactly what the Unicode-table path
|
||||
/// answers. These decide run splitting and boundary markers, so one wrong
|
||||
/// answer silently changes counts for the most common input there is —
|
||||
/// and the shortcuts are justified by claims about Unicode data
|
||||
/// (no ASCII codepoint combines, ASCII has no gc=No/gc=Cf, …) that only an
|
||||
/// exhaustive comparison can hold to account.
|
||||
#[test]
|
||||
fn ascii_fast_paths_match_general() {
|
||||
for i in 0u8..128 {
|
||||
let c = char::from(i);
|
||||
let at = format!("ascii {i:#04x}");
|
||||
assert_eq!(ASCII_CLASS[i as usize], classify_nonascii(c), "class {at}");
|
||||
assert_eq!(is_separator(c), is_separator_general(c), "separator {at}");
|
||||
assert_eq!(is_stray_mark(c), is_stray_mark_general(c), "stray mark {at}");
|
||||
assert_eq!(digit_border(c), digit_border_general(c), "digit border {at}");
|
||||
assert_eq!(marks_like_punct(c), marks_like_punct_general(c), "punct-like {at}");
|
||||
let body = c.to_string();
|
||||
assert_eq!(is_digit_run(&body), is_digit_run_general(&body), "digit run {at}");
|
||||
}
|
||||
// The ASCII digit-run shortcut also has to agree on runs, not just on
|
||||
// single characters: uniform, mixed, and empty.
|
||||
for body in ["", "0", "12", "123456789", "1a", "a1", "1 2", "٣", "1٣", "½", "12½"] {
|
||||
assert_eq!(is_digit_run(body), is_digit_run_general(body), "digit run {body:?}");
|
||||
}
|
||||
}
|
||||
|
||||
/// PUNCT is the class `marks_like_punct` mirrors on ASCII; if that ever
|
||||
/// stops holding, the shortcut above is reading the wrong table.
|
||||
#[test]
|
||||
fn ascii_punct_class_is_punct_like() {
|
||||
for i in 0u8..128 {
|
||||
let c = char::from(i);
|
||||
assert_eq!(
|
||||
ASCII_CLASS[i as usize] == Class::Punct,
|
||||
marks_like_punct_general(c),
|
||||
"ascii {i:#04x}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// The ideograph shortcut answers for ~38 k codepoints without a Unicode
|
||||
/// lookup; if any of them is not gc=Lo in the vendored tables (an
|
||||
/// unassigned hole, a future reassignment), the shortcut would silently
|
||||
/// mark borders differently, so every codepoint in the ranges is compared.
|
||||
#[test]
|
||||
fn ideograph_shortcut_matches_general() {
|
||||
let ranges = [0x3400u32..=0x4dbf, 0x4e00..=0x9fff, 0xac00..=0xd7a3];
|
||||
let mut checked = 0usize;
|
||||
for cp in ranges.into_iter().flatten() {
|
||||
let c = char::from_u32(cp).expect("ideograph ranges hold no surrogate");
|
||||
assert_eq!(classify(c), classify_nonascii(c), "class {cp:#06x}");
|
||||
assert_eq!(marks_like_punct(c), marks_like_punct_general(c), "punct-like {cp:#06x}");
|
||||
assert_eq!(hard_kind(c), hard_kind_general(c), "hard kind {cp:#06x}");
|
||||
checked += 1;
|
||||
}
|
||||
assert_eq!(checked, 0x4dbf - 0x3400 + 0x9fff - 0x4e00 + 0xd7a3 - 0xac00 + 3);
|
||||
}
|
||||
}
|
||||
+2960
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,254 @@
|
||||
[
|
||||
{
|
||||
"text": "a ",
|
||||
"count": 8
|
||||
},
|
||||
{
|
||||
"text": "a\t",
|
||||
"count": 8
|
||||
},
|
||||
{
|
||||
"text": "a \n",
|
||||
"count": 8
|
||||
},
|
||||
{
|
||||
"text": "a\n ",
|
||||
"count": 9
|
||||
},
|
||||
{
|
||||
"text": "a \t \n ",
|
||||
"count": 12
|
||||
},
|
||||
{
|
||||
"text": " a ",
|
||||
"count": 9
|
||||
},
|
||||
{
|
||||
"text": " 'REXX ",
|
||||
"count": 12
|
||||
},
|
||||
{
|
||||
"text": "日本語 text ",
|
||||
"count": 12
|
||||
},
|
||||
{
|
||||
"text": " hello ",
|
||||
"count": 10
|
||||
},
|
||||
{
|
||||
"text": "'REXX\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
|
||||
"count": 11
|
||||
},
|
||||
{
|
||||
"text": "a\n",
|
||||
"count": 7
|
||||
},
|
||||
{
|
||||
"text": "a\n\n",
|
||||
"count": 7
|
||||
},
|
||||
{
|
||||
"text": "a\n\n\n\n\n",
|
||||
"count": 7
|
||||
},
|
||||
{
|
||||
"text": "a\n\n\n\n\n\n\n\n\n\n\n\n",
|
||||
"count": 7
|
||||
},
|
||||
{
|
||||
"text": "a\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
|
||||
"count": 8
|
||||
},
|
||||
{
|
||||
"text": "a\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
|
||||
"count": 8
|
||||
},
|
||||
{
|
||||
"text": "a\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
|
||||
"count": 8
|
||||
},
|
||||
{
|
||||
"text": "a\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
|
||||
"count": 7
|
||||
},
|
||||
{
|
||||
"text": "a\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n",
|
||||
"count": 8
|
||||
},
|
||||
{
|
||||
"text": "para\n\npara\n",
|
||||
"count": 9
|
||||
},
|
||||
{
|
||||
"text": "\tindent\n\tindent\n",
|
||||
"count": 13
|
||||
},
|
||||
{
|
||||
"text": "trailing \r\n",
|
||||
"count": 10
|
||||
},
|
||||
{
|
||||
"text": "mix ",
|
||||
"count": 9
|
||||
},
|
||||
{
|
||||
"text": "<.;)кл ж漢仮三本仮中平文一名三二`#+/आअक्ंआेईिइघी—‘’。”",
|
||||
"count": 52
|
||||
},
|
||||
{
|
||||
"text": " if err != nil {\n\treturn err\n}\n\n",
|
||||
"count": 19
|
||||
},
|
||||
{
|
||||
"text": " SELECT * FROM users WHERE id = 42;\n\n",
|
||||
"count": 25
|
||||
},
|
||||
{
|
||||
"text": "다다라마나한่̧️््̧́️่",
|
||||
"count": 28
|
||||
},
|
||||
{
|
||||
"text": "XOIOTV््่่่่̧́́dpldvnzgdazจา์กคงัีुआकूअंेआििの語👍🏽✨🚀✨🏽⚖",
|
||||
"count": 69
|
||||
},
|
||||
{
|
||||
"text": "ΒΖкааабжблгв9,614€…！…！jkrm마가라다어나마한UEòáßñá",
|
||||
"count": 45
|
||||
},
|
||||
{
|
||||
"text": "اختثثóçúßáвбкг гаебле\n\n \t \t \n\t",
|
||||
"count": 31
|
||||
},
|
||||
{
|
||||
"text": " SELECT * FROM users WHERE id = 42;\n\n",
|
||||
"count": 25
|
||||
},
|
||||
{
|
||||
"text": ">/,(,^%.+@%",
|
||||
"count": 17
|
||||
},
|
||||
{
|
||||
"text": "jumps naïve the CamelCase jumps fox brown naïve NASA HTTP2 CamelCase naïve",
|
||||
"count": 45
|
||||
},
|
||||
{
|
||||
"text": "NASA we'll HTTP2 don't CamelCase brown",
|
||||
"count": 27
|
||||
},
|
||||
{
|
||||
"text": "ANUअ्इीअिाçìééßúàèóßù́́️่̧̧́éàóòàééáç0315,",
|
||||
"count": 49
|
||||
},
|
||||
{
|
||||
"text": "they've CamelCase quick the the brown NASA jumps brown naïve don't the it's naïve quick they've",
|
||||
"count": 42
|
||||
},
|
||||
{
|
||||
"text": "://]&تب تحجحبحاлзкйдг",
|
||||
"count": 24
|
||||
},
|
||||
{
|
||||
"text": "्่̧́́-가어어국다다라어국🏽😀🎉⚖✨\t\n\n \n\t\n çìòá ",
|
||||
"count": 51
|
||||
},
|
||||
{
|
||||
"text": "CamelCase CamelCase naïve fox HTTP2 don't jumps fox fox NASA we'll İstanbul quick they've don't İstanbul brown",
|
||||
"count": 60
|
||||
},
|
||||
{
|
||||
"text": "ंुुईीค่éóíóجج خاخبت्̧̧́️̧̧́",
|
||||
"count": 45
|
||||
},
|
||||
{
|
||||
"text": "SELECT * FROM users WHERE id = 42;",
|
||||
"count": 24
|
||||
},
|
||||
{
|
||||
"text": " SELECT * FROM users WHERE id = 42;",
|
||||
"count": 25
|
||||
},
|
||||
{
|
||||
"text": " function f(a, b) { return a === b || a <= b; }\n",
|
||||
"count": 26
|
||||
},
|
||||
{
|
||||
"text": " SELECT * FROM users WHERE id = 42;\n",
|
||||
"count": 25
|
||||
},
|
||||
{
|
||||
"text": "const x = { a: 1, b: [2, 3] };\n\n",
|
||||
"count": 24
|
||||
},
|
||||
{
|
||||
"text": " SELECT * FROM users WHERE id = 42;\n\n",
|
||||
"count": 25
|
||||
},
|
||||
{
|
||||
"text": "İstanbul NASA jumps brown jumps jumps HTTP2 CamelCase NASA we'll brown brown quick jumps İstanbul NASA NASA",
|
||||
"count": 62
|
||||
},
|
||||
{
|
||||
"text": "खकईआूुिईثاخخ اجبااث \n\t\n \t\n",
|
||||
"count": 30
|
||||
},
|
||||
{
|
||||
"text": "длкддкагзLQZIFZEWY 名日平本の平漢εΔ나국 라라나 나국나라",
|
||||
"count": 45
|
||||
},
|
||||
{
|
||||
"text": " const x = { a: 1, b: [2, 3] };\n\n\n\n",
|
||||
"count": 25
|
||||
},
|
||||
{
|
||||
"text": "don't they've brown NASA fox fox HTTP2 the don't",
|
||||
"count": 28
|
||||
},
|
||||
{
|
||||
"text": "if err != nil {\n\treturn err\n}\n",
|
||||
"count": 18
|
||||
},
|
||||
{
|
||||
"text": "가어한나ाुअकीघคิา๊๋ัข็ค่๊ีحجا حتث خبختิีิ์ זהבג דדבהא平語中本二字語ß",
|
||||
"count": 69
|
||||
},
|
||||
{
|
||||
"text": " \t\t \t\n\t ́️्่̧",
|
||||
"count": 26
|
||||
},
|
||||
{
|
||||
"text": " SELECT * FROM users WHERE id = 42;\n",
|
||||
"count": 25
|
||||
},
|
||||
{
|
||||
"text": " if err != nil {\n\treturn err\n}",
|
||||
"count": 19
|
||||
},
|
||||
{
|
||||
"text": "naïve NASA fox jumps quick quick we'll we'll HTTP2 quick quick it's HTTP2 we'll brown it's they've don't",
|
||||
"count": 56
|
||||
},
|
||||
{
|
||||
"text": " function f(a, b) { return a === b || a <= b; }\n",
|
||||
"count": 26
|
||||
},
|
||||
{
|
||||
"text": " \t \n \n\n\t\t🚀🏽🎉✨😀️️✨️🚀่️्่́́️्््ζΓ्่́️्่́́",
|
||||
"count": 66
|
||||
},
|
||||
{
|
||||
"text": " const x = { a: 1, b: [2, 3] };\n\n\n\n",
|
||||
"count": 25
|
||||
},
|
||||
{
|
||||
"text": " if err != nil {\n\treturn err\n}\n\n",
|
||||
"count": 19
|
||||
},
|
||||
{
|
||||
"text": " for i in range(10):\n print(i ** 2)\n\n\n\n",
|
||||
"count": 24
|
||||
},
|
||||
{
|
||||
"text": "it's we'll it's brown we'll the quick we'll CamelCase CamelCase jumps HTTP2 NASA",
|
||||
"count": 46
|
||||
}
|
||||
]
|
||||
@@ -28,6 +28,7 @@ pub mod audio;
|
||||
pub mod block;
|
||||
pub mod clipboard;
|
||||
pub mod crash_handler;
|
||||
pub mod ctok;
|
||||
pub mod desktop;
|
||||
pub mod devicecheck;
|
||||
pub mod diff;
|
||||
|
||||
@@ -1,17 +1,19 @@
|
||||
//! Token counting via tiktoken-rs.
|
||||
//! Token counting via tiktoken-rs and the ctok Claude reconstruction.
|
||||
//!
|
||||
//! Two encodings are exposed:
|
||||
//! Encodings:
|
||||
//!
|
||||
//! - `O200kBase` — GPT-4o / o1 / GPT-5 (the modern `OpenAI` default).
|
||||
//! - `Cl100kBase` — GPT-3.5 / GPT-4 / older models.
|
||||
//! - `ClaudeV3` / `ClaudeV47` / `ClaudeV5` — offline reconstructions of
|
||||
//! Anthropic's `count_tokens` (see [`crate::ctok`]): v3 serves Claude 3
|
||||
//! through Opus 4.6, v4.7 serves Opus 4.7–4.9, v5 serves the 5-series.
|
||||
//!
|
||||
//! `o200k_base` is the default. Anthropic doesn't publish their tokenizer, so
|
||||
//! either of these is an approximation for Claude (within ~5–10% across
|
||||
//! English/code text). `o200k_base` is closer to current frontier models'
|
||||
//! actual segmentation and is the right default for budget estimates.
|
||||
//! `o200k_base` is the default. For Claude models the ctok encodings count
|
||||
//! exactly (message content, excluding the fixed per-message frame), so they
|
||||
//! are the right choice wherever the model is known to be Claude.
|
||||
//!
|
||||
//! Both BPE tables are embedded in the binary; encoders are built once on
|
||||
//! first use and reused thereafter.
|
||||
//! BPE tables and the ctok vocabularies are embedded in the binary; encoders
|
||||
//! are built once on first use and reused thereafter.
|
||||
|
||||
use std::sync::LazyLock;
|
||||
|
||||
@@ -21,6 +23,8 @@ use pi_shell::rayon_global_pool_available;
|
||||
use rayon::prelude::*;
|
||||
use tiktoken_rs::{CoreBPE, cl100k_base, o200k_base};
|
||||
|
||||
use crate::ctok;
|
||||
|
||||
/// Tokenizer encoding to use.
|
||||
#[napi(string_enum)]
|
||||
pub enum Encoding {
|
||||
@@ -28,6 +32,14 @@ pub enum Encoding {
|
||||
O200kBase,
|
||||
/// GPT-3.5 / GPT-4 / older.
|
||||
Cl100kBase,
|
||||
/// Claude 3 … Opus 4.6 (ctok v3 reconstruction).
|
||||
ClaudeV3,
|
||||
/// Claude Opus 4.7–4.9 (ctok v4.7 reconstruction).
|
||||
ClaudeV47,
|
||||
/// Claude Opus 5+ (ctok v5 reconstruction).
|
||||
ClaudeV5,
|
||||
/// Claude Sonnet/Fable 5+ (live-measured non-opus v5 frame).
|
||||
ClaudeV5Sonnet,
|
||||
}
|
||||
|
||||
static O200K: LazyLock<CoreBPE> =
|
||||
@@ -36,10 +48,29 @@ static O200K: LazyLock<CoreBPE> =
|
||||
static CL100K: LazyLock<CoreBPE> =
|
||||
LazyLock::new(|| cl100k_base().expect("failed to initialize cl100k_base BPE tables"));
|
||||
|
||||
fn encoder(encoding: Option<Encoding>) -> &'static CoreBPE {
|
||||
/// A resolved counting backend: a BPE encoder or a ctok family.
|
||||
enum Counter {
|
||||
Bpe(&'static CoreBPE),
|
||||
Claude(ctok::Family),
|
||||
}
|
||||
|
||||
impl Counter {
|
||||
fn resolve(encoding: Option<Encoding>) -> Self {
|
||||
match encoding.unwrap_or(Encoding::O200kBase) {
|
||||
Encoding::O200kBase => &O200K,
|
||||
Encoding::Cl100kBase => &CL100K,
|
||||
Encoding::O200kBase => Self::Bpe(&O200K),
|
||||
Encoding::Cl100kBase => Self::Bpe(&CL100K),
|
||||
Encoding::ClaudeV3 => Self::Claude(ctok::Family::V3),
|
||||
Encoding::ClaudeV47 => Self::Claude(ctok::Family::V47),
|
||||
Encoding::ClaudeV5 => Self::Claude(ctok::Family::V5),
|
||||
Encoding::ClaudeV5Sonnet => Self::Claude(ctok::Family::V5Sonnet),
|
||||
}
|
||||
}
|
||||
|
||||
fn count(&self, text: &str) -> u32 {
|
||||
match self {
|
||||
Self::Bpe(bpe) => bpe.encode_ordinary(text).len() as u32,
|
||||
Self::Claude(family) => ctok::content_token_count(text, *family),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -50,21 +81,18 @@ fn encoder(encoding: Option<Encoding>) -> &'static CoreBPE {
|
||||
/// pool is available). Always returns a single token total — use this for any
|
||||
/// aggregate budget question without paying a per-element napi crossing.
|
||||
///
|
||||
/// Uses ordinary encoding (no special-token handling), which is the right
|
||||
/// choice for measuring user/model content rather than wire-protocol tokens.
|
||||
/// Defaults to `o200k_base`; pass `Cl100kBase` for older `OpenAI` models.
|
||||
/// Measures user/model content, not wire-protocol tokens: BPE encodings use
|
||||
/// ordinary encoding (no special-token handling) and the Claude encodings
|
||||
/// count message content without the fixed per-message frame. Defaults to
|
||||
/// `o200k_base`; pass a `Claude*` encoding for exact Claude counts.
|
||||
#[napi]
|
||||
pub fn count_tokens(input: Either<String, Vec<String>>, encoding: Option<Encoding>) -> u32 {
|
||||
let bpe = encoder(encoding);
|
||||
let counter = Counter::resolve(encoding);
|
||||
match input {
|
||||
Either::A(text) => bpe.encode_ordinary(&text).len() as u32,
|
||||
Either::B(texts) if rayon_global_pool_available() => texts
|
||||
.par_iter()
|
||||
.map(|s| bpe.encode_ordinary(s).len() as u32)
|
||||
.sum(),
|
||||
Either::B(texts) => texts
|
||||
.iter()
|
||||
.map(|s| bpe.encode_ordinary(s).len() as u32)
|
||||
.sum(),
|
||||
Either::A(text) => counter.count(&text),
|
||||
Either::B(texts) if rayon_global_pool_available() => {
|
||||
texts.par_iter().map(|s| counter.count(s)).sum()
|
||||
},
|
||||
Either::B(texts) => texts.iter().map(|s| counter.count(s)).sum(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,6 +2,20 @@
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Breaking Changes
|
||||
|
||||
- Local token counting is now an immutable, model-scoped `Tokenizer` instance instead of process-global functions. The free `countTokens`/`countTokensConservatively` exports and the `setTokenizerModel()` global setter are gone; construct a `Tokenizer(modelId)` — the encoding is fixed at construction, there is no `setModel` — and call `tokenizer.countTokens(text, mode?)`. An `Agent` owns one for its active model (exposed as `agent.tokenizer`) and replaces the instance when the active model's encoding changes, so don't cache it across model switches. `countTokensConservatively` collapsed into that one method as `mode: "upperbound"`; the modes are `"strict"` (always exact native), `"approximate"` (default, fast byte estimate when no exact tokenizer applies), and `"upperbound"` (raw byte length, never undercounts).
|
||||
- The free `estimateTokens(message, tokenizer, options?)` is gone; use `tokenizer.countMessage(message, options?)` or `tokenizer.countMessages(messages, options?)`. Per-message estimates are memoized per tokenizer instance (keyed by message identity) and invalidated across every instance by `invalidateMessageCache(message)`, which bumps a shared version tag. `findCutPoint`, `prepareBranchEntries`, `collectShakeRegions`, `pruneToolOutputs`, `pruneSupersededToolResults`, and `trimRemoteCompactionInputToContextWindow` take an explicit `Tokenizer`, and `prepareCompaction` accepts the caller's (warm) tokenizer as an optional trailing parameter. Token math is scoped to the model that will be billed for it rather than to whichever `Agent` last constructed itself.
|
||||
|
||||
### Added
|
||||
|
||||
- Added `Tokenizer.checkTokenBudget(text, budget)`: a cheap-first budget probe. Byte length is a hard upper bound on token count, so text whose raw bytes already fit answers "fits" without tokenizing at all; only text that busts the bound pays for an exact count (and that count is returned, so a proportional clamp gets the denominator it needs). Since the bound overshoots ~4x on prose, the common "comfortably under budget" answer is free. Compaction's summary-window fit check and OpenAI remote-compaction trimming now route through it.
|
||||
- Added provider-anchored transcript accounting (`findTranscriptUsageAnchor`, `isTranscriptUsageAnchor`, `estimateTranscriptTokens`). Every settled assistant turn carries `usage` covering the exact prompt it was sent, so transcript sizing charges that report for the prefix and tokenizes only the tail appended after it — counting proportional to one turn instead of the whole history, every turn. The four hand-rolled copies of the anchor trust rules (session stats ×3, shake) now share one predicate, and the deliberately provider-independent compaction floor counts every message locally via `tokenizer.countMessages`.
|
||||
|
||||
### Changed
|
||||
|
||||
- Claude models get exact native ctok counts instead of the bytes/4 estimate: the tokenizer family is resolved per model (v3 for Claude 3 … Opus 4.6, v4.7 for Opus 4.7–4.9, v5 for Opus 5+, the sonnet-5 frame variant for the non-opus 5-series). Non-Claude models keep the fast estimate (or o200k with `PI_TOKENIZER_ACCURATE=1`).
|
||||
|
||||
## [17.3.8] - 2026-08-19
|
||||
|
||||
### Fixed
|
||||
|
||||
@@ -37,6 +37,7 @@ import {
|
||||
} from "./agent-loop";
|
||||
import type { AppendOnlyContextManager } from "./append-only-context";
|
||||
import { isProviderRefusalMessage } from "./replay-policy";
|
||||
import { claudeEncodingForModel, Tokenizer } from "./tokenizer";
|
||||
import type {
|
||||
AgentBeforeModelCall,
|
||||
AgentContext,
|
||||
@@ -363,7 +364,7 @@ export class Agent {
|
||||
pendingToolCalls: new Set<string>(),
|
||||
error: undefined,
|
||||
};
|
||||
|
||||
#tokenizer = new Tokenizer(this.#state.model?.id);
|
||||
#listeners = new Set<(e: AgentEvent) => void>();
|
||||
#abortController?: AbortController;
|
||||
#convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
|
||||
@@ -459,6 +460,7 @@ export class Agent {
|
||||
if (opts.initialState?.messages) this.#state.messages = opts.initialState.messages.slice();
|
||||
if (opts.initialState?.pendingToolCalls)
|
||||
this.#state.pendingToolCalls = new Set(opts.initialState.pendingToolCalls);
|
||||
this.#syncTokenizer(this.#state.model?.id);
|
||||
this.#convertToLlm = opts.convertToLlm || defaultConvertToLlm;
|
||||
this.#transformContext = opts.transformContext;
|
||||
this.#steeringMode = opts.steeringMode || "one-at-a-time";
|
||||
@@ -720,11 +722,29 @@ export class Agent {
|
||||
set maxRetryDelayMs(value: number | undefined) {
|
||||
this.#maxRetryDelayMs = value;
|
||||
}
|
||||
|
||||
get state(): AgentState {
|
||||
return this.#state;
|
||||
}
|
||||
|
||||
/**
|
||||
* Tokenizer for the active model. The instance is replaced whenever the
|
||||
* active model's encoding changes (see {@link setModel}), so callers must
|
||||
* not cache it across model switches.
|
||||
*/
|
||||
get tokenizer(): Tokenizer {
|
||||
return this.#tokenizer;
|
||||
}
|
||||
|
||||
/**
|
||||
* Swap the tokenizer only when the encoding actually changes, so the warm
|
||||
* per-message memo survives same-encoding model switches.
|
||||
*/
|
||||
#syncTokenizer(modelId: string | null | undefined): void {
|
||||
if ((modelId ? claudeEncodingForModel(modelId) : null) !== this.#tokenizer.encoding) {
|
||||
this.#tokenizer = new Tokenizer(modelId);
|
||||
}
|
||||
}
|
||||
|
||||
get appendOnlyContext(): AppendOnlyContextManager | undefined {
|
||||
return this.#appendOnlyContext;
|
||||
}
|
||||
@@ -904,6 +924,7 @@ export class Agent {
|
||||
|
||||
setModel(m: Model) {
|
||||
this.#state.model = m;
|
||||
this.#syncTokenizer(m?.id);
|
||||
}
|
||||
|
||||
setThinkingLevel(l: Effort | undefined) {
|
||||
|
||||
@@ -9,8 +9,8 @@ import type { Api, ApiKey, AssistantMessage, Context, Model, SimpleStreamOptions
|
||||
import { preferredDialect } from "@oh-my-pi/pi-catalog/identity";
|
||||
import { prompt } from "@oh-my-pi/pi-utils";
|
||||
import { type AgentTelemetry, instrumentedCompleteSimple } from "../telemetry";
|
||||
import { Tokenizer } from "../tokenizer";
|
||||
import type { AgentMessage } from "../types";
|
||||
import { estimateTokens } from "./compaction";
|
||||
import type { ReadonlySessionManager, SessionEntry } from "./entries";
|
||||
import {
|
||||
type ConvertToLlm,
|
||||
@@ -206,14 +206,14 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined {
|
||||
}
|
||||
}
|
||||
|
||||
function estimateBranchSummaryTokens(message: AgentMessage): number {
|
||||
if (message.role !== "toolResult") return estimateTokens(message);
|
||||
function estimateBranchSummaryTokens(message: AgentMessage, tokenizer: Tokenizer): number {
|
||||
if (message.role !== "toolResult") return tokenizer.countMessage(message);
|
||||
const text = message.content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
.map(c => c.text)
|
||||
.join("");
|
||||
if (!text) return 0;
|
||||
return estimateTokens({
|
||||
return tokenizer.countMessage({
|
||||
...message,
|
||||
content: [{ type: "text", text: truncateToolResultForSummary(text) }],
|
||||
});
|
||||
@@ -232,7 +232,11 @@ function estimateBranchSummaryTokens(message: AgentMessage): number {
|
||||
* @param entries - Entries in chronological order
|
||||
* @param tokenBudget - Maximum tokens to include (0 = no limit)
|
||||
*/
|
||||
export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: number = 0): BranchPreparation {
|
||||
export function prepareBranchEntries(
|
||||
entries: SessionEntry[],
|
||||
tokenizer: Tokenizer,
|
||||
tokenBudget: number = 0,
|
||||
): BranchPreparation {
|
||||
const messages: AgentMessage[] = [];
|
||||
const fileOps = createFileOps();
|
||||
let totalTokens = 0;
|
||||
@@ -264,7 +268,7 @@ export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: numbe
|
||||
// Extract file ops from assistant messages (tool calls)
|
||||
extractFileOpsFromMessage(message, fileOps);
|
||||
|
||||
const tokens = estimateBranchSummaryTokens(message);
|
||||
const tokens = estimateBranchSummaryTokens(message, tokenizer);
|
||||
|
||||
// Check budget before adding
|
||||
if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) {
|
||||
@@ -309,8 +313,9 @@ export async function generateBranchSummary(
|
||||
// Token budget = context window minus reserved space for prompt + response
|
||||
const contextWindow = model.contextWindow || 128000;
|
||||
const tokenBudget = contextWindow - reserveTokens;
|
||||
const tokenizer = new Tokenizer(model.id);
|
||||
|
||||
const { messages, fileOps } = prepareBranchEntries(entries, tokenBudget);
|
||||
const { messages, fileOps } = prepareBranchEntries(entries, tokenizer, tokenBudget);
|
||||
|
||||
if (messages.length === 0) {
|
||||
return { summary: "No content to summarize" };
|
||||
|
||||
@@ -31,11 +31,11 @@ import { buildResponsesInput, resolveOpenAICompatPolicy } from "@oh-my-pi/pi-ai/
|
||||
import { stripOpenAIResponsesOutputOnlyStatusesForReplay } from "@oh-my-pi/pi-ai/utils";
|
||||
import { preferredDialect } from "@oh-my-pi/pi-catalog/identity";
|
||||
import { clampThinkingLevelForModel } from "@oh-my-pi/pi-catalog/model-thinking";
|
||||
import { isRecord, logger, prompt, stringifyJson } from "@oh-my-pi/pi-utils";
|
||||
import { isRecord, logger, prompt } from "@oh-my-pi/pi-utils";
|
||||
import * as snapcompact from "@oh-my-pi/snapcompact";
|
||||
import { type AgentTelemetry, instrumentedCompleteSimple } from "../telemetry";
|
||||
import { ThinkingLevel } from "../thinking";
|
||||
import { countTokens } from "../tokenizer";
|
||||
import { Tokenizer } from "../tokenizer";
|
||||
import type { AgentMessage } from "../types";
|
||||
import {
|
||||
buildCompactionV2Request,
|
||||
@@ -47,7 +47,6 @@ import {
|
||||
} from "./compaction-v2-streaming";
|
||||
import type { CompactionEntry, SessionEntry } from "./entries";
|
||||
import { NativeCompactionError } from "./errors";
|
||||
import { isEstimateCacheable, readEstimateCache, writeEstimateCache } from "./message-cache";
|
||||
import { type ConvertToLlm, createBranchSummaryMessage, createCustomMessage, defaultConvertToLlm } from "./messages";
|
||||
import {
|
||||
buildOpenAiNativeHistory,
|
||||
@@ -390,144 +389,17 @@ export function resolveThresholdTokens(contextWindow: number, settings: Compacti
|
||||
// Cut point detection
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* Image content has no tokenizer representation; charge a fixed estimate
|
||||
* matching what providers typically bill for inline images.
|
||||
*/
|
||||
const IMAGE_TOKEN_ESTIMATE = 1200;
|
||||
|
||||
/**
|
||||
* Estimate token count for a message using cl100k_base via the native
|
||||
* tokenizer. This is not Claude's first-party tokenizer (Anthropic doesn't
|
||||
* publish one) but is within ~5–10% across English/code text.
|
||||
*
|
||||
* `excludeEncryptedReasoning` drops opaque provider reasoning payloads
|
||||
* (`thinkingSignature`, `redactedThinking`) from the estimate. Those are billed
|
||||
* by the provider on replay, so the default counts them — but their *local*
|
||||
* byte size can diverge wildly from what the provider charges, so the
|
||||
* compaction floor (which only needs the reliably-countable, on-wire-compressible
|
||||
* content) excludes them to avoid false triggers on thinking-heavy turns.
|
||||
*/
|
||||
export function estimateTokens(message: AgentMessage, options?: { excludeEncryptedReasoning?: boolean }): number {
|
||||
// Settled historical messages are counted once and reused until an owner
|
||||
// (prune/shake/strip-images) invalidates them; streaming assistants bypass
|
||||
// the cache entirely (see message-cache.ts settle-gate invariant).
|
||||
const cacheable = isEstimateCacheable(message);
|
||||
const excludeEncryptedReasoning = options?.excludeEncryptedReasoning === true;
|
||||
if (cacheable) {
|
||||
const cached = readEstimateCache(message, excludeEncryptedReasoning);
|
||||
if (cached !== undefined) return cached;
|
||||
}
|
||||
const result = computeMessageTokens(message, options);
|
||||
if (cacheable) writeEstimateCache(message, excludeEncryptedReasoning, result);
|
||||
return result;
|
||||
}
|
||||
|
||||
function computeMessageTokens(message: AgentMessage, options?: { excludeEncryptedReasoning?: boolean }): number {
|
||||
const fragments: string[] = [];
|
||||
let extra = 0;
|
||||
if ((message as { role?: string }).role === "bashExecution") {
|
||||
const bash = message as { command?: unknown; output?: unknown };
|
||||
if (typeof bash.command === "string") fragments.push(bash.command);
|
||||
if (typeof bash.output === "string") fragments.push(bash.output);
|
||||
return fragments.length === 0 ? 0 : countTokens(fragments);
|
||||
}
|
||||
|
||||
switch (message.role) {
|
||||
case "user": {
|
||||
const content = (message as { content: string | Array<{ type: string; text?: string }> }).content;
|
||||
if (typeof content === "string") {
|
||||
fragments.push(content);
|
||||
} else if (Array.isArray(content)) {
|
||||
for (const block of content) {
|
||||
if (block.type === "text" && block.text) {
|
||||
fragments.push(block.text);
|
||||
}
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "assistant": {
|
||||
const assistant = message as AssistantMessage;
|
||||
for (const block of assistant.content) {
|
||||
if (block.type === "text") {
|
||||
fragments.push(block.text);
|
||||
} else if (block.type === "thinking") {
|
||||
fragments.push(block.thinking);
|
||||
// Providers charge for the opaque signature/reasoning payload that
|
||||
// rides alongside the thinking text (OpenAI Responses encrypted
|
||||
// reasoning items, Anthropic signed thinking blocks, etc.). Without
|
||||
// counting it, this estimator can read ~half of the provider-reported
|
||||
// usage on thinking-heavy turns — see #2275 for the resulting
|
||||
// compaction-trigger / post-check metric divergence. The compaction
|
||||
// floor excludes it (its local byte size diverges from provider billing).
|
||||
if (block.thinkingSignature && !options?.excludeEncryptedReasoning) {
|
||||
fragments.push(block.thinkingSignature);
|
||||
}
|
||||
} else if (block.type === "toolCall") {
|
||||
fragments.push(block.name);
|
||||
fragments.push(stringifyJson(block.arguments) ?? "null");
|
||||
} else if (block.type === "redactedThinking") {
|
||||
// Encrypted reasoning blob the provider still bills for on replay;
|
||||
// excluded from the compaction floor for the same reason as above.
|
||||
if (!options?.excludeEncryptedReasoning) fragments.push(block.data);
|
||||
} else if (block.type === "anthropicServerTool") {
|
||||
// Native Anthropic server-tool call/result replayed verbatim on the
|
||||
// wire (server_tool_use input and opaque result content). This opaque
|
||||
// provider-replay state the provider still
|
||||
// bills for on same-provider replay; excluded from the compaction
|
||||
// floor like other encrypted reasoning because its local byte size
|
||||
// diverges from provider billing.
|
||||
if (!options?.excludeEncryptedReasoning) fragments.push(stringifyJson(block.block) ?? "null");
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "hookMessage":
|
||||
case "toolResult": {
|
||||
if (typeof message.content === "string") {
|
||||
fragments.push(message.content);
|
||||
} else {
|
||||
for (const block of message.content) {
|
||||
if (block.type === "text" && block.text) {
|
||||
fragments.push(block.text);
|
||||
} else if (block.type === "image") {
|
||||
extra += IMAGE_TOKEN_ESTIMATE;
|
||||
}
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "branchSummary":
|
||||
case "compactionSummary": {
|
||||
fragments.push(message.summary);
|
||||
if (message.role === "compactionSummary") {
|
||||
if (message.blocks) {
|
||||
for (const block of message.blocks) {
|
||||
if (block.type === "text") fragments.push(block.text);
|
||||
else extra += snapcompact.FRAME_TOKEN_ESTIMATE;
|
||||
}
|
||||
} else if (message.images) {
|
||||
// Snapcompact frames render at ≥1568px; providers bill the downscaled cap.
|
||||
extra += message.images.length * snapcompact.FRAME_TOKEN_ESTIMATE;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (fragments.length === 0) return extra;
|
||||
return extra + countTokens(fragments);
|
||||
}
|
||||
|
||||
function estimateEntriesTokens(entries: SessionEntry[], startIndex: number, endIndex: number): number {
|
||||
function estimateEntriesTokens(
|
||||
entries: SessionEntry[],
|
||||
tokenizer: Tokenizer,
|
||||
startIndex: number,
|
||||
endIndex: number,
|
||||
): number {
|
||||
let total = 0;
|
||||
for (let i = startIndex; i < endIndex; i++) {
|
||||
const msg = getMessageFromEntry(entries[i]);
|
||||
if (msg) {
|
||||
total += estimateTokens(msg);
|
||||
total += tokenizer.countMessage(msg);
|
||||
}
|
||||
}
|
||||
return total;
|
||||
@@ -626,6 +498,7 @@ export interface CutPointResult {
|
||||
*/
|
||||
export function findCutPoint(
|
||||
entries: SessionEntry[],
|
||||
tokenizer: Tokenizer,
|
||||
startIndex: number,
|
||||
endIndex: number,
|
||||
keepRecentTokens: number,
|
||||
@@ -645,7 +518,7 @@ export function findCutPoint(
|
||||
if (entry.type !== "message") continue;
|
||||
|
||||
// Estimate this message's size
|
||||
const messageTokens = estimateTokens(entry.message);
|
||||
const messageTokens = tokenizer.countMessage(entry.message);
|
||||
accumulatedTokens += messageTokens;
|
||||
|
||||
// Check if we've exceeded the budget
|
||||
@@ -948,12 +821,17 @@ interface SummaryWindow {
|
||||
* on message boundaries. Only called when the whole conversation does not fit —
|
||||
* the common single-window path never pays this per-message sizing pass.
|
||||
*/
|
||||
function planSummaryWindows(messages: Message[], dialect: Dialect | undefined, budgetTokens: number): Message[][] {
|
||||
function planSummaryWindows(
|
||||
messages: Message[],
|
||||
tokenizer: Tokenizer,
|
||||
dialect: Dialect | undefined,
|
||||
budgetTokens: number,
|
||||
): Message[][] {
|
||||
const windows: Message[][] = [];
|
||||
let current: Message[] = [];
|
||||
let currentTokens = 0;
|
||||
for (const message of messages) {
|
||||
const tokens = countTokens(serializeConversationForSummary([message], dialect));
|
||||
const tokens = tokenizer.countTokens(serializeConversationForSummary([message], dialect));
|
||||
if (currentTokens > 0 && currentTokens + tokens > budgetTokens) {
|
||||
windows.push(current);
|
||||
current = [];
|
||||
@@ -982,6 +860,7 @@ export async function generateSummary(
|
||||
// Convert to LLM messages first (handles custom app messages when caller provides a transformer).
|
||||
const llmMessages = (options?.convertToLlm ?? defaultConvertToLlm)(currentMessages);
|
||||
const dialect = preferredDialect(model.id);
|
||||
const tokenizer = new Tokenizer(model.id);
|
||||
const wholeConversation = serializeConversationForSummary(llmMessages, dialect);
|
||||
const budgetTokens = summaryInputBudgetTokens(model, maxTokens);
|
||||
// A span that outgrew the summarizer's window is summarized as a fold: each
|
||||
@@ -991,19 +870,21 @@ export async function generateSummary(
|
||||
// retry can shrink — the state a cross-provider compaction boundary
|
||||
// (see `prepareCompaction`) puts a long session into. One window is the
|
||||
// common case and costs exactly the one call it always did.
|
||||
const pending: SummaryWindow[] =
|
||||
countTokens(wholeConversation) <= budgetTokens
|
||||
const pending: SummaryWindow[] = tokenizer.checkTokenBudget(wholeConversation, budgetTokens).fits
|
||||
? [{ messages: llmMessages, budgetTokens, text: wholeConversation }]
|
||||
: planSummaryWindows(llmMessages, dialect, budgetTokens).map(messages => ({ messages, budgetTokens }));
|
||||
: planSummaryWindows(llmMessages, tokenizer, dialect, budgetTokens).map(messages => ({ messages, budgetTokens }));
|
||||
|
||||
let carriedSummary = previousSummary;
|
||||
while (pending.length > 0) {
|
||||
const window = pending[0];
|
||||
const text = window.text ?? serializeConversationForSummary(window.messages, dialect);
|
||||
const windowTokens = countTokens(text);
|
||||
// A budget probe, not a raw count: a window whose bytes already fit needs
|
||||
// neither an exact count nor the clamp, and the bust path hands back the
|
||||
// exact count the proportional clamp needs as its denominator.
|
||||
const budget = tokenizer.checkTokenBudget(text, window.budgetTokens);
|
||||
try {
|
||||
carriedSummary = await summarizeConversationWindow(
|
||||
clampConversationToBudget(text, window.budgetTokens, windowTokens),
|
||||
budget.fits ? text : clampConversationToBudget(text, window.budgetTokens, budget.tokens),
|
||||
carriedSummary,
|
||||
model,
|
||||
maxTokens,
|
||||
@@ -1020,8 +901,11 @@ export async function generateSummary(
|
||||
// window size only the provider can tell us is wrong.
|
||||
// Halve what was actually SENT, not the budget it was planned against:
|
||||
// the rejection proves the plan was fiction, so converging on the real
|
||||
// cap must not spend a call per level of an imaginary ladder.
|
||||
const halved = Math.floor(Math.min(window.budgetTokens, windowTokens) / 2);
|
||||
// cap must not spend a call per level of an imaginary ladder. The cheap
|
||||
// fit path never counted this window, so pay for the exact size here —
|
||||
// one tokenization is nothing against the provider round trip already lost.
|
||||
const sentTokens = budget.exact ? budget.tokens : tokenizer.countTokens(text, "strict");
|
||||
const halved = Math.floor(Math.min(window.budgetTokens, sentTokens) / 2);
|
||||
if (
|
||||
!AIError.is(AIError.classify(error), AIError.Flag.ContextOverflow) ||
|
||||
halved < minSummaryInputTokens(model)
|
||||
@@ -1031,7 +915,7 @@ export async function generateSummary(
|
||||
pending.splice(
|
||||
0,
|
||||
1,
|
||||
...planSummaryWindows(window.messages, dialect, halved).map(messages => ({
|
||||
...planSummaryWindows(window.messages, tokenizer, dialect, halved).map(messages => ({
|
||||
messages,
|
||||
budgetTokens: halved,
|
||||
})),
|
||||
@@ -1423,10 +1307,16 @@ export function findReadableCompactionIndex(
|
||||
return -1;
|
||||
}
|
||||
|
||||
/**
|
||||
* Pass the caller's warm `tokenizer` (the Agent's for the active model) so the
|
||||
* full-branch estimate walk hits its memo; the cold default is for one-shot
|
||||
* callers that have no live agent.
|
||||
*/
|
||||
export function prepareCompaction(
|
||||
pathEntries: SessionEntry[],
|
||||
settings: CompactionSettings,
|
||||
activeModel?: Model,
|
||||
tokenizer: Tokenizer = new Tokenizer(activeModel?.id),
|
||||
): CompactionPreparation | undefined {
|
||||
if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1].type === "compaction") {
|
||||
return undefined;
|
||||
@@ -1459,7 +1349,7 @@ export function prepareCompaction(
|
||||
const tokensBefore = lastUsage ? calculateContextTokens(lastUsage) : 0;
|
||||
let keepRecentTokens = settings.keepRecentTokens;
|
||||
if (lastUsage) {
|
||||
const estimatedTokens = estimateEntriesTokens(pathEntries, boundaryStart, boundaryEnd);
|
||||
const estimatedTokens = estimateEntriesTokens(pathEntries, tokenizer, boundaryStart, boundaryEnd);
|
||||
const promptTokens = calculatePromptTokens(lastUsage);
|
||||
const ratio = estimatedTokens > 0 ? promptTokens / estimatedTokens : 0;
|
||||
if (Number.isFinite(ratio) && ratio > 1) {
|
||||
@@ -1467,7 +1357,7 @@ export function prepareCompaction(
|
||||
}
|
||||
}
|
||||
|
||||
const cutPoint = findCutPoint(pathEntries, boundaryStart, boundaryEnd, keepRecentTokens);
|
||||
const cutPoint = findCutPoint(pathEntries, tokenizer, boundaryStart, boundaryEnd, keepRecentTokens);
|
||||
|
||||
// Get ID of first kept entry
|
||||
const firstKeptEntry = pathEntries[cutPoint.firstKeptEntryIndex];
|
||||
@@ -1706,6 +1596,7 @@ export async function compact(
|
||||
: undefined;
|
||||
const trimmed = trimRemoteCompactionInputToContextWindow(
|
||||
remoteHistory,
|
||||
new Tokenizer(model.id),
|
||||
model.contextWindow,
|
||||
instructions,
|
||||
tools,
|
||||
|
||||
@@ -11,4 +11,5 @@ export * from "./messages";
|
||||
export * from "./openai";
|
||||
export * from "./pruning";
|
||||
export * from "./shake";
|
||||
export * from "./transcript-tokens";
|
||||
export * from "./utils";
|
||||
|
||||
@@ -1,13 +1,12 @@
|
||||
/**
|
||||
* Per-message memoization for the two hot history walks: token estimation
|
||||
* ({@link estimateTokens}) and LLM conversion (the coding-agent's `convertToLlm`).
|
||||
* Cache-coherence seams for the two hot history walks: token estimation
|
||||
* ({@link Tokenizer.countMessage}) and LLM conversion (the coding-agent's
|
||||
* `convertToLlm`).
|
||||
*
|
||||
* Long sessions re-walk a settled `AgentMessage[]` every turn, re-tokenizing and
|
||||
* re-converting historical objects that only the newest suffix can change. These
|
||||
* caches key on message *identity* so a settled message is counted/converted once
|
||||
* and reused until an owner rewrites it.
|
||||
*
|
||||
* Correctness rests on two invariants:
|
||||
* re-converting historical objects that only the newest suffix can change. Each
|
||||
* `Tokenizer` memoizes estimates per message identity; this module owns the two
|
||||
* invariants that keep those memos (and the cross-package convert memo) honest:
|
||||
*
|
||||
* 1. **Settle gate.** A streaming assistant is mutated under one identity while
|
||||
* its `usage`/`stopReason` are provisional (the seed carries zeroed usage and
|
||||
@@ -19,9 +18,11 @@
|
||||
* 2. **Owner invalidation.** `pruneToolOutputs` / `pruneSupersededToolResults`,
|
||||
* `applyShakeRegion`, and `stripImagesFromMessage` rewrite message content in
|
||||
* place under a stable identity. Each MUST call {@link invalidateMessageCache}
|
||||
* on the mutated message before the next convert/estimate pass so both caches
|
||||
* drop the stale entry. The convert cache lives in another package, so it
|
||||
* subscribes via {@link registerMessageCacheInvalidator}.
|
||||
* on the mutated message before the next convert/estimate pass. Invalidation
|
||||
* bumps a symbol-keyed version tag on the message itself, so every live
|
||||
* `Tokenizer` memo drops its stale entry at once without registering
|
||||
* anywhere; the convert cache lives in another package and subscribes via
|
||||
* {@link registerMessageCacheInvalidator}.
|
||||
*/
|
||||
import type { AssistantMessage } from "@oh-my-pi/pi-ai";
|
||||
import type { AgentMessage } from "../types";
|
||||
@@ -41,18 +42,26 @@ export function registerMessageCacheInvalidator(invalidate: (message: AgentMessa
|
||||
};
|
||||
}
|
||||
|
||||
// Dual option-split estimate caches: the compaction floor passes
|
||||
// `excludeEncryptedReasoning` (dropping opaque provider reasoning), so a message
|
||||
// has two distinct estimates that must not collide in one map.
|
||||
//
|
||||
// These are WeakMaps, not symbol-tagged properties, deliberately: callers spread
|
||||
// messages to derive throwaway variants for counting — `estimateBranchSummaryTokens`
|
||||
// does `estimateTokens({ ...message, content: truncated })`. A symbol-keyed cache
|
||||
// value rides along an object spread, so the truncated clone would inherit (and
|
||||
// return) the full-content estimate. Keying strictly on identity keeps the cache
|
||||
// off spread copies, which get their own fresh count.
|
||||
const estimateCacheDefault = new WeakMap<AgentMessage, number>();
|
||||
const estimateCacheFloored = new WeakMap<AgentMessage, number>();
|
||||
/**
|
||||
* Estimate-version tag riding on the message itself. Symbol-keyed, so JSON
|
||||
* session persistence and default iteration never see it. Object spread copies
|
||||
* the tag onto derived clones — harmless, because estimate memos key on message
|
||||
* *identity* and a fresh clone starts with no memo entries anywhere.
|
||||
*/
|
||||
const kEstimateVersion = Symbol("omp.messageEstimateVersion");
|
||||
|
||||
interface VersionedMessage {
|
||||
[kEstimateVersion]?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Current estimate version of `message` (0 until first invalidation). A
|
||||
* `Tokenizer` memo entry stamped with an older version is stale and must be
|
||||
* recounted.
|
||||
*/
|
||||
export function messageEstimateVersion(message: AgentMessage): number {
|
||||
return (message as VersionedMessage)[kEstimateVersion] ?? 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* True when this message's estimate is safe to cache by identity. Non-assistants
|
||||
@@ -70,23 +79,13 @@ export function isEstimateCacheable(message: AgentMessage): boolean {
|
||||
);
|
||||
}
|
||||
|
||||
/** Read a cached estimate for the given option split, or `undefined` on miss. */
|
||||
export function readEstimateCache(message: AgentMessage, excludeEncryptedReasoning: boolean): number | undefined {
|
||||
return (excludeEncryptedReasoning ? estimateCacheFloored : estimateCacheDefault).get(message);
|
||||
}
|
||||
|
||||
/** Store an estimate for the given option split. */
|
||||
export function writeEstimateCache(message: AgentMessage, excludeEncryptedReasoning: boolean, value: number): void {
|
||||
(excludeEncryptedReasoning ? estimateCacheFloored : estimateCacheDefault).set(message, value);
|
||||
}
|
||||
|
||||
/**
|
||||
* Drop every cached derivation of `message` after an in-place rewrite. Owners of
|
||||
* mutation (prune, shake, strip-images) call this at the mutation seam so the
|
||||
* next convert/estimate pass recomputes from the new content.
|
||||
*/
|
||||
export function invalidateMessageCache(message: AgentMessage): void {
|
||||
estimateCacheDefault.delete(message);
|
||||
estimateCacheFloored.delete(message);
|
||||
const versioned = message as VersionedMessage;
|
||||
versioned[kEstimateVersion] = ((versioned[kEstimateVersion] ?? 0) + 1) | 0;
|
||||
for (const invalidate of externalInvalidators) invalidate(message);
|
||||
}
|
||||
|
||||
@@ -51,7 +51,7 @@ import {
|
||||
OPENAI_HEADERS,
|
||||
} from "@oh-my-pi/pi-catalog/wire/codex";
|
||||
import { $env, isRecord, logger, prompt, stringifyJson, structuredCloneJSON } from "@oh-my-pi/pi-utils";
|
||||
import { countTokensConservatively } from "../tokenizer";
|
||||
import { Tokenizer } from "../tokenizer";
|
||||
import contextWindowTruncatedOutputPrompt from "./prompts/context-window-truncated-output.md" with { type: "text" };
|
||||
|
||||
export * from "./compaction-v2-streaming";
|
||||
@@ -123,14 +123,36 @@ export interface TrimRemoteCompactionInputResult {
|
||||
estimatedTokensAfter: number;
|
||||
}
|
||||
|
||||
function estimateRemoteCompactionInputTokens(
|
||||
/** Verdict for one remote-compaction request measured against the model window. */
|
||||
interface RemoteCompactionBudgetProbe {
|
||||
/** Estimated request tokens; the text part is exact when the cheap bound busted. */
|
||||
tokens: number;
|
||||
/** Whether the request fits the window. Always true when no window is known. */
|
||||
fits: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Cheap-first sizing of a remote-compaction request. Images and the request
|
||||
* frame are charged flat, so they come off the budget rather than through the
|
||||
* tokenizer; the serialized transcript is then probed with
|
||||
* {@link Tokenizer.checkTokenBudget}, which only pays for an exact count when
|
||||
* the byte bound cannot already prove the request fits.
|
||||
*/
|
||||
function probeRemoteCompactionInputBudget(
|
||||
input: Array<Record<string, unknown>>,
|
||||
tokenizer: Tokenizer,
|
||||
instructions: string,
|
||||
tools?: unknown[],
|
||||
): number {
|
||||
tools: unknown[] | undefined,
|
||||
contextWindow: number | null | undefined,
|
||||
): RemoteCompactionBudgetProbe {
|
||||
const normalized = normalizeRemoteCompactionEstimateValue({ instructions, input, ...(tools ? { tools } : {}) });
|
||||
const serialized = stringifyJson(normalized.value) ?? "";
|
||||
return countTokensConservatively(serialized) + normalized.imageTokens + REMOTE_COMPACTION_REQUEST_OVERHEAD_TOKENS;
|
||||
const flatTokens = normalized.imageTokens + REMOTE_COMPACTION_REQUEST_OVERHEAD_TOKENS;
|
||||
if (!contextWindow || contextWindow <= 0) {
|
||||
return { tokens: tokenizer.countTokens(serialized, "upperbound") + flatTokens, fits: true };
|
||||
}
|
||||
const budget = tokenizer.checkTokenBudget(serialized, Math.max(0, contextWindow - flatTokens));
|
||||
return { tokens: budget.tokens + flatTokens, fits: budget.fits };
|
||||
}
|
||||
|
||||
function rewriteToolOutputForContextWindow(item: Record<string, unknown>): Record<string, unknown> | undefined {
|
||||
@@ -164,24 +186,25 @@ function isToolResultImageAttachment(item: Record<string, unknown>): boolean {
|
||||
*/
|
||||
export function trimRemoteCompactionInputToContextWindow(
|
||||
input: Array<Record<string, unknown>>,
|
||||
tokenizer: Tokenizer,
|
||||
contextWindow: number | null | undefined,
|
||||
instructions: string,
|
||||
tools?: unknown[],
|
||||
): TrimRemoteCompactionInputResult {
|
||||
const estimatedTokensBefore = estimateRemoteCompactionInputTokens(input, instructions, tools);
|
||||
if (!contextWindow || contextWindow <= 0 || estimatedTokensBefore <= contextWindow) {
|
||||
const before = probeRemoteCompactionInputBudget(input, tokenizer, instructions, tools, contextWindow);
|
||||
if (before.fits) {
|
||||
return {
|
||||
input,
|
||||
rewrittenOutputs: 0,
|
||||
estimatedTokensBefore,
|
||||
estimatedTokensAfter: estimatedTokensBefore,
|
||||
estimatedTokensBefore: before.tokens,
|
||||
estimatedTokensAfter: before.tokens,
|
||||
};
|
||||
}
|
||||
|
||||
let rewrittenInput: Array<Record<string, unknown>> | undefined;
|
||||
let estimatedTokensAfter = estimatedTokensBefore;
|
||||
let after = before;
|
||||
let rewrittenOutputs = 0;
|
||||
for (let index = input.length - 1; index >= 0 && estimatedTokensAfter > contextWindow; index--) {
|
||||
for (let index = input.length - 1; index >= 0 && !after.fits; index--) {
|
||||
const item = input[index];
|
||||
if (isToolResultImageAttachment(item)) continue;
|
||||
const rewritten = rewriteToolOutputForContextWindow(item);
|
||||
@@ -189,23 +212,23 @@ export function trimRemoteCompactionInputToContextWindow(
|
||||
rewrittenInput ??= input.slice();
|
||||
rewrittenInput[index] = rewritten;
|
||||
rewrittenOutputs++;
|
||||
estimatedTokensAfter = estimateRemoteCompactionInputTokens(rewrittenInput, instructions, tools);
|
||||
after = probeRemoteCompactionInputBudget(rewrittenInput, tokenizer, instructions, tools, contextWindow);
|
||||
}
|
||||
|
||||
if (!rewrittenInput || estimatedTokensAfter > contextWindow) {
|
||||
if (!rewrittenInput || !after.fits) {
|
||||
return {
|
||||
input,
|
||||
rewrittenOutputs: 0,
|
||||
estimatedTokensBefore,
|
||||
estimatedTokensAfter: estimatedTokensBefore,
|
||||
estimatedTokensBefore: before.tokens,
|
||||
estimatedTokensAfter: before.tokens,
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
input: rewrittenInput,
|
||||
rewrittenOutputs,
|
||||
estimatedTokensBefore,
|
||||
estimatedTokensAfter,
|
||||
estimatedTokensBefore: before.tokens,
|
||||
estimatedTokensAfter: after.tokens,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -766,7 +789,12 @@ export async function requestOpenAiRemoteCompaction(
|
||||
): Promise<OpenAiRemoteCompactionResponse> {
|
||||
const endpoint = resolveOpenAiCompactEndpoint(model);
|
||||
const requestModel = resolveOpenAiCompactModel(model);
|
||||
const trimmed = trimRemoteCompactionInputToContextWindow(compactInput, model.contextWindow, instructions);
|
||||
const trimmed = trimRemoteCompactionInputToContextWindow(
|
||||
compactInput,
|
||||
new Tokenizer(model.id),
|
||||
model.contextWindow,
|
||||
instructions,
|
||||
);
|
||||
if (trimmed.rewrittenOutputs > 0) {
|
||||
logger.info("Rewrote trailing tool outputs before OpenAI remote compaction", {
|
||||
model: model.id,
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
*/
|
||||
|
||||
import type { ToolResultMessage } from "@oh-my-pi/pi-ai";
|
||||
import type { Tokenizer } from "../tokenizer";
|
||||
import type { AgentMessage, AgentToolCall } from "../types";
|
||||
import { estimateTokens } from "./compaction";
|
||||
import type { SessionEntry, SessionMessageEntry } from "./entries";
|
||||
import { invalidateMessageCache } from "./message-cache";
|
||||
import {
|
||||
@@ -140,13 +140,13 @@ function estimatePrunedSavings(tokens: number, notice: string): number {
|
||||
* (cacheWrite premium) if that entry is mutated in place. Used to keep prune
|
||||
* mutations inside the cheap-to-recache tail.
|
||||
*/
|
||||
function computeMessageSuffixTokens(entries: readonly SessionEntry[]): number[] {
|
||||
function computeMessageSuffixTokens(entries: readonly SessionEntry[], tokenizer: Tokenizer): number[] {
|
||||
const suffix = new Array<number>(entries.length);
|
||||
let accumulated = 0;
|
||||
for (let i = entries.length - 1; i >= 0; i--) {
|
||||
suffix[i] = accumulated;
|
||||
const entry = entries[i];
|
||||
if (entry.type === "message") accumulated += estimateTokens(entry.message as AgentMessage);
|
||||
if (entry.type === "message") accumulated += tokenizer.countMessage(entry.message as AgentMessage);
|
||||
}
|
||||
return suffix;
|
||||
}
|
||||
@@ -181,6 +181,7 @@ interface SupersedeCandidate {
|
||||
*/
|
||||
function collectSupersededResults(
|
||||
entries: readonly SessionEntry[],
|
||||
tokenizer: Tokenizer,
|
||||
toolCallsById: ReadonlyMap<string, AgentToolCall>,
|
||||
supersedeKey: SupersedeKeyFn,
|
||||
protectedTools: readonly ProtectedToolMatcher[],
|
||||
@@ -204,7 +205,7 @@ function collectSupersededResults(
|
||||
entry: entry as SessionMessageEntry,
|
||||
message,
|
||||
index: i,
|
||||
tokens: estimateTokens(message as AgentMessage),
|
||||
tokens: tokenizer.countMessage(message as AgentMessage),
|
||||
notice: SUPERSEDED_NOTICE,
|
||||
});
|
||||
}
|
||||
@@ -219,6 +220,7 @@ function collectSupersededResults(
|
||||
*/
|
||||
function collectUselessResults(
|
||||
entries: readonly SessionEntry[],
|
||||
tokenizer: Tokenizer,
|
||||
toolCallsById: ReadonlyMap<string, AgentToolCall>,
|
||||
protectedTools: readonly ProtectedToolMatcher[],
|
||||
exclude: ReadonlySet<ToolResultMessage>,
|
||||
@@ -230,7 +232,7 @@ function collectUselessResults(
|
||||
if (message?.useless !== true || message.prunedAt !== undefined || message.isError === true) continue;
|
||||
if (exclude.has(message)) continue;
|
||||
if (isProtectedToolResult(message, toolCallsById.get(message.toolCallId), protectedTools)) continue;
|
||||
const tokens = estimateTokens(message as AgentMessage);
|
||||
const tokens = tokenizer.countMessage(message as AgentMessage);
|
||||
if (estimatePrunedSavings(tokens, USELESS_NOTICE) <= 0) continue;
|
||||
candidates.push({ entry: entry as SessionMessageEntry, message, index: i, tokens, notice: USELESS_NOTICE });
|
||||
}
|
||||
@@ -246,14 +248,18 @@ function collectUselessResults(
|
||||
* the provider cache is cold anyway (then all still-sent candidates flush).
|
||||
* Never mutates entries before `keepBoundaryId` (summarized away — not sent).
|
||||
*/
|
||||
export function pruneSupersededToolResults(entries: SessionEntry[], config: SupersedePruneConfig): PruneResult {
|
||||
export function pruneSupersededToolResults(
|
||||
entries: SessionEntry[],
|
||||
tokenizer: Tokenizer,
|
||||
config: SupersedePruneConfig,
|
||||
): PruneResult {
|
||||
const toolCallsById = collectToolCallsById(entries);
|
||||
const candidates = config.supersedeKey
|
||||
? collectSupersededResults(entries, toolCallsById, config.supersedeKey, config.protectedTools)
|
||||
? collectSupersededResults(entries, tokenizer, toolCallsById, config.supersedeKey, config.protectedTools)
|
||||
: [];
|
||||
if (config.pruneUseless) {
|
||||
const exclude = new Set(candidates.map(candidate => candidate.message));
|
||||
candidates.push(...collectUselessResults(entries, toolCallsById, config.protectedTools, exclude));
|
||||
candidates.push(...collectUselessResults(entries, tokenizer, toolCallsById, config.protectedTools, exclude));
|
||||
candidates.sort((a, b) => a.index - b.index);
|
||||
}
|
||||
if (candidates.length === 0) return { prunedCount: 0, tokensSaved: 0 };
|
||||
@@ -284,7 +290,7 @@ export function pruneSupersededToolResults(entries: SessionEntry[], config: Supe
|
||||
// Mutating a candidate re-writes its suffix in the warm cache, so prune only
|
||||
// when that suffix is small (cheap-to-recache tail) and the candidate sits
|
||||
// at/after the compaction boundary.
|
||||
const suffixTokens = computeMessageSuffixTokens(entries);
|
||||
const suffixTokens = computeMessageSuffixTokens(entries, tokenizer);
|
||||
toPrune = candidates.filter(
|
||||
candidate => candidate.index >= boundaryIndex && suffixTokens[candidate.index] <= suffixTokenLimit,
|
||||
);
|
||||
@@ -302,7 +308,11 @@ export function pruneSupersededToolResults(entries: SessionEntry[], config: Supe
|
||||
return { prunedCount: toPrune.length, tokensSaved };
|
||||
}
|
||||
|
||||
export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = DEFAULT_PRUNE_CONFIG): PruneResult {
|
||||
export function pruneToolOutputs(
|
||||
entries: SessionEntry[],
|
||||
tokenizer: Tokenizer,
|
||||
config: PruneConfig = DEFAULT_PRUNE_CONFIG,
|
||||
): PruneResult {
|
||||
let accumulatedTokens = 0;
|
||||
let tokensSaved = 0;
|
||||
let prunedCount = 0;
|
||||
@@ -311,7 +321,7 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig =
|
||||
const toolCallsById = collectToolCallsById(entries);
|
||||
const supersededMessages = config.supersedeKey
|
||||
? new Set(
|
||||
collectSupersededResults(entries, toolCallsById, config.supersedeKey, config.protectedTools).map(
|
||||
collectSupersededResults(entries, tokenizer, toolCallsById, config.supersedeKey, config.protectedTools).map(
|
||||
candidate => candidate.message,
|
||||
),
|
||||
)
|
||||
@@ -321,6 +331,7 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig =
|
||||
? new Set(
|
||||
collectUselessResults(
|
||||
entries,
|
||||
tokenizer,
|
||||
toolCallsById,
|
||||
config.protectedTools,
|
||||
supersededMessages ?? new Set(),
|
||||
@@ -331,14 +342,15 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig =
|
||||
const boundaryIndex = resolveBoundaryIndex(entries, config.keepBoundaryId);
|
||||
const cacheWarmSuffixTokens = config.cacheWarmSuffixTokens;
|
||||
// All-message suffix per index, only when the cache guard is armed.
|
||||
const messageSuffix = cacheWarmSuffixTokens === undefined ? undefined : computeMessageSuffixTokens(entries);
|
||||
const messageSuffix =
|
||||
cacheWarmSuffixTokens === undefined ? undefined : computeMessageSuffixTokens(entries, tokenizer);
|
||||
|
||||
for (let i = entries.length - 1; i >= 0; i--) {
|
||||
const entry = entries[i];
|
||||
const message = getToolResultMessage(entry);
|
||||
if (!message) continue;
|
||||
|
||||
const tokens = estimateTokens(message as AgentMessage);
|
||||
const tokens = tokenizer.countMessage(message as AgentMessage);
|
||||
const isProtected = isProtectedToolResult(message, toolCallsById.get(message.toolCallId), config.protectedTools);
|
||||
|
||||
if (message.prunedAt !== undefined) {
|
||||
|
||||
@@ -11,9 +11,8 @@
|
||||
*/
|
||||
|
||||
import type { TextContent, ToolResultMessage } from "@oh-my-pi/pi-ai";
|
||||
import { countTokens } from "../tokenizer";
|
||||
import type { Tokenizer } from "../tokenizer";
|
||||
import type { AgentMessage } from "../types";
|
||||
import { estimateTokens } from "./compaction";
|
||||
import type { CustomMessageEntry, SessionEntry, SessionMessageEntry } from "./entries";
|
||||
import { invalidateMessageCache } from "./message-cache";
|
||||
import {
|
||||
@@ -123,15 +122,15 @@ function toolResultText(message: ToolResultMessage): string {
|
||||
}
|
||||
|
||||
/** Estimate the token contribution of an entry for the protect-recent window. */
|
||||
function entryTokens(entry: SessionEntry): number {
|
||||
function entryTokens(entry: SessionEntry, tokenizer: Tokenizer): number {
|
||||
if (entry.type === "message") {
|
||||
return estimateTokens(entry.message);
|
||||
return tokenizer.countMessage(entry.message);
|
||||
}
|
||||
if (entry.type === "custom_message") {
|
||||
const content = entry.content;
|
||||
if (typeof content === "string") return content.length === 0 ? 0 : countTokens(content);
|
||||
if (typeof content === "string") return content.length === 0 ? 0 : tokenizer.countTokens(content);
|
||||
const fragments = content.filter((block): block is TextContent => block.type === "text").map(block => block.text);
|
||||
return fragments.length === 0 ? 0 : countTokens(fragments);
|
||||
return fragments.length === 0 ? 0 : tokenizer.countTokens(fragments);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -222,6 +221,7 @@ function pushBlockRegions(
|
||||
entry: SessionMessageEntry | CustomMessageEntry,
|
||||
blockIndex: number,
|
||||
text: string,
|
||||
tokenizer: Tokenizer,
|
||||
config: ShakeConfig,
|
||||
label: string,
|
||||
out: ShakeRegion[],
|
||||
@@ -229,7 +229,7 @@ function pushBlockRegions(
|
||||
for (const range of scanTextForBlockRanges(text)) {
|
||||
const slice = text.slice(range.start, range.end);
|
||||
if (slice.length === 0) continue;
|
||||
const tokens = countTokens(slice);
|
||||
const tokens = tokenizer.countTokens(slice);
|
||||
if (tokens < config.fenceMinTokens) continue;
|
||||
out.push({
|
||||
kind: "block",
|
||||
@@ -246,6 +246,7 @@ function pushBlockRegions(
|
||||
|
||||
function collectBlockRegions(
|
||||
entry: SessionMessageEntry | CustomMessageEntry,
|
||||
tokenizer: Tokenizer,
|
||||
config: ShakeConfig,
|
||||
out: ShakeRegion[],
|
||||
): void {
|
||||
@@ -254,34 +255,35 @@ function collectBlockRegions(
|
||||
if (message.role === "assistant") {
|
||||
for (let bi = 0; bi < message.content.length; bi++) {
|
||||
const block = message.content[bi];
|
||||
if (block.type === "text") pushBlockRegions(entry, bi, block.text, config, "assistant", out);
|
||||
if (block.type === "text") pushBlockRegions(entry, bi, block.text, tokenizer, config, "assistant", out);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (message.role === "user" || message.role === "developer") {
|
||||
scanContentBlocks(entry, message.content, config, message.role, out);
|
||||
scanContentBlocks(entry, message.content, tokenizer, config, message.role, out);
|
||||
}
|
||||
return;
|
||||
}
|
||||
// custom_message
|
||||
scanContentBlocks(entry, entry.content, config, entry.customType, out);
|
||||
scanContentBlocks(entry, entry.content, tokenizer, config, entry.customType, out);
|
||||
}
|
||||
|
||||
function scanContentBlocks(
|
||||
entry: SessionMessageEntry | CustomMessageEntry,
|
||||
content: string | Array<{ type: string; text?: string }>,
|
||||
tokenizer: Tokenizer,
|
||||
config: ShakeConfig,
|
||||
label: string,
|
||||
out: ShakeRegion[],
|
||||
): void {
|
||||
if (typeof content === "string") {
|
||||
pushBlockRegions(entry, -1, content, config, label, out);
|
||||
pushBlockRegions(entry, -1, content, tokenizer, config, label, out);
|
||||
return;
|
||||
}
|
||||
for (let bi = 0; bi < content.length; bi++) {
|
||||
const block = content[bi];
|
||||
if (block.type === "text" && typeof block.text === "string") {
|
||||
pushBlockRegions(entry, bi, block.text, config, label, out);
|
||||
pushBlockRegions(entry, bi, block.text, tokenizer, config, label, out);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -300,7 +302,7 @@ function scanContentBlocks(
|
||||
* and regions never span a message boundary. When the combined estimated
|
||||
* savings is below `minSavings`, returns `[]` (no-op).
|
||||
*/
|
||||
export function collectShakeRegions(entries: SessionEntry[], config: ShakeConfig): ShakeRegion[] {
|
||||
export function collectShakeRegions(entries: SessionEntry[], tokenizer: Tokenizer, config: ShakeConfig): ShakeRegion[] {
|
||||
const n = entries.length;
|
||||
if (n === 0) return [];
|
||||
|
||||
@@ -309,7 +311,7 @@ export function collectShakeRegions(entries: SessionEntry[], config: ShakeConfig
|
||||
let acc = 0;
|
||||
for (let i = n - 1; i >= 0; i--) {
|
||||
accumulatedAfter[i] = acc;
|
||||
acc += entryTokens(entries[i]);
|
||||
acc += entryTokens(entries[i], tokenizer);
|
||||
}
|
||||
|
||||
const toolCallsById = collectToolCallsById(entries);
|
||||
@@ -342,7 +344,7 @@ export function collectShakeRegions(entries: SessionEntry[], config: ShakeConfig
|
||||
regions.push({
|
||||
kind: "toolResult",
|
||||
entry: entry as SessionMessageEntry,
|
||||
tokens: estimateTokens(toolResult as AgentMessage),
|
||||
tokens: tokenizer.countMessage(toolResult as AgentMessage),
|
||||
originalText: text,
|
||||
label: toolResult.toolName,
|
||||
});
|
||||
@@ -350,7 +352,7 @@ export function collectShakeRegions(entries: SessionEntry[], config: ShakeConfig
|
||||
}
|
||||
|
||||
if (entry.type === "message" || entry.type === "custom_message") {
|
||||
collectBlockRegions(entry as SessionMessageEntry | CustomMessageEntry, config, regions);
|
||||
collectBlockRegions(entry as SessionMessageEntry | CustomMessageEntry, tokenizer, config, regions);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
/**
|
||||
* Provider-anchored transcript token accounting.
|
||||
*
|
||||
* Local tokenization is the expensive way to answer "how big is this
|
||||
* conversation?" — and usually the wrong one, because the provider already
|
||||
* answered it. Every settled assistant turn carries `usage` covering the exact
|
||||
* prompt it was sent: the system prompt, the tool schemas, and every message up
|
||||
* to and including itself. The only genuinely unaccounted-for text is the tail
|
||||
* appended *after* that turn.
|
||||
*
|
||||
* These helpers locate the newest trustworthy usage report and tokenize only
|
||||
* that tail, so a long session pays counting proportional to one turn instead
|
||||
* of to the whole transcript, every turn.
|
||||
*
|
||||
* Trust rules for an anchor (mirroring the provider contract):
|
||||
* - Assistant role only — nothing else carries `usage`.
|
||||
* - Not `aborted` / `error`: those turns report partial or zero usage.
|
||||
* - `hasContextTokenUsage(usage)`: the report must carry usable context numbers.
|
||||
*/
|
||||
|
||||
import type { AssistantMessage } from "@oh-my-pi/pi-ai";
|
||||
import type { MessageCountOptions, Tokenizer } from "../tokenizer";
|
||||
import type { AgentMessage } from "../types";
|
||||
import { calculateContextTokens, hasContextTokenUsage } from "./compaction";
|
||||
|
||||
/** A provider usage report that accounts for a prefix of the transcript. */
|
||||
export interface TranscriptUsageAnchor {
|
||||
/** Index in the scanned array; messages at or before it are provider-accounted. */
|
||||
index: number;
|
||||
/** The anchoring assistant turn. */
|
||||
message: AssistantMessage;
|
||||
/** Conversation tokens the provider reported for that prompt. */
|
||||
tokens: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Whether this message's provider usage may anchor transcript accounting.
|
||||
*
|
||||
* The single home for the trust rules — every anchor scan MUST route through
|
||||
* it so a stale-usage rule can never drift between the transcript walkers and
|
||||
* the session-entry walkers.
|
||||
*/
|
||||
export function isTranscriptUsageAnchor(message: AgentMessage): message is AssistantMessage {
|
||||
if (message.role !== "assistant") return false;
|
||||
const assistant = message as AssistantMessage;
|
||||
if (assistant.stopReason === "aborted" || assistant.stopReason === "error") return false;
|
||||
return assistant.usage !== undefined && hasContextTokenUsage(assistant.usage);
|
||||
}
|
||||
|
||||
/**
|
||||
* Newest assistant turn in `messages[fromIndex..]` whose usage can anchor the
|
||||
* transcript, or `undefined` when none qualifies (fresh context, or every
|
||||
* recent turn aborted/errored).
|
||||
*
|
||||
* `fromIndex` excludes turns whose usage is stale — anything a compaction
|
||||
* summarized away describes a prompt that is no longer sent.
|
||||
*/
|
||||
export function findTranscriptUsageAnchor(
|
||||
messages: readonly AgentMessage[],
|
||||
fromIndex = 0,
|
||||
): TranscriptUsageAnchor | undefined {
|
||||
for (let index = messages.length - 1; index >= fromIndex; index--) {
|
||||
const message = messages[index];
|
||||
if (!isTranscriptUsageAnchor(message)) continue;
|
||||
return { index, message, tokens: calculateContextTokens(message.usage) };
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/** Options for {@link estimateTranscriptTokens}. */
|
||||
export interface TranscriptTokenOptions {
|
||||
/**
|
||||
* Gates the anchor search only: usage at or before this index is stale (a
|
||||
* compaction rewrote the prompt it describes) and must not anchor. Content
|
||||
* accounting is governed separately by {@link countFromIndex}.
|
||||
*/
|
||||
anchorFromIndex?: number;
|
||||
/**
|
||||
* First message whose content is counted locally when no anchor is found.
|
||||
* Defaults to 0 (count the whole transcript), which is what a floor
|
||||
* estimate wants; pass the compaction boundary to skip summarized-away
|
||||
* messages entirely.
|
||||
*/
|
||||
countFromIndex?: number;
|
||||
/** Forwarded to {@link Tokenizer.countMessage} for every locally counted message. */
|
||||
excludeEncryptedReasoning?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Conversation tokens for `messages`: the provider's own report for everything
|
||||
* it already covers, plus a local count of only the unaccounted-for tail.
|
||||
*
|
||||
* An anchored result already includes the non-message prefix (system prompt +
|
||||
* tool schemas) because the provider charged it; an unanchored result is a
|
||||
* message-only sum. Callers that add non-message tokens on top MUST branch on
|
||||
* {@link findTranscriptUsageAnchor} rather than assuming one shape.
|
||||
*/
|
||||
export function estimateTranscriptTokens(
|
||||
messages: readonly AgentMessage[],
|
||||
tokenizer: Tokenizer,
|
||||
options?: TranscriptTokenOptions,
|
||||
): number {
|
||||
const estimateOptions: MessageCountOptions | undefined =
|
||||
options?.excludeEncryptedReasoning === true ? { excludeEncryptedReasoning: true } : undefined;
|
||||
const anchor = findTranscriptUsageAnchor(messages, options?.anchorFromIndex ?? 0);
|
||||
let total = anchor?.tokens ?? 0;
|
||||
for (let index = anchor ? anchor.index + 1 : (options?.countFromIndex ?? 0); index < messages.length; index++) {
|
||||
total += tokenizer.countMessage(messages[index], estimateOptions);
|
||||
}
|
||||
return total;
|
||||
}
|
||||
+282
-19
@@ -1,27 +1,290 @@
|
||||
import { countTokens as countTokensNat } from "@oh-my-pi/pi-natives";
|
||||
import { bareModelId, parseAnthropicModel, semverGte } from "@oh-my-pi/pi-catalog/identity";
|
||||
import { countTokens as countTokensNat, Encoding } from "@oh-my-pi/pi-natives";
|
||||
import { stringifyJson } from "@oh-my-pi/pi-utils";
|
||||
import * as snapcompact from "@oh-my-pi/snapcompact";
|
||||
import { isEstimateCacheable, messageEstimateVersion } from "./compaction/message-cache";
|
||||
import type { AgentMessage } from "./types";
|
||||
|
||||
const accurate = process.env.PI_TOKENIZER_ACCURATE === "1" && Bun.env.NODE_ENV !== "test";
|
||||
const testEnv = Bun.env.NODE_ENV === "test";
|
||||
const accurate = process.env.PI_TOKENIZER_ACCURATE === "1" && !testEnv;
|
||||
|
||||
function estimateTokens(text: string) {
|
||||
/**
|
||||
* ctok encoding for a Claude model id, or `null` for non-Claude models.
|
||||
*
|
||||
* Family routing mirrors the ctok reconstruction plus live measurement:
|
||||
* Claude 3 through Opus 4.6 (and every non-opus Claude below 5) count with
|
||||
* v3, Opus 4.7–4.9 with v4.7, Opus 5+ with v5, and the non-opus 5-series
|
||||
* (sonnet/fable/mythos) with the sonnet-5 frame variant. Ids
|
||||
* `parseAnthropicModel` cannot classify (e.g. haiku) fall back to v3, which
|
||||
* covers every such model shipped to date.
|
||||
*/
|
||||
export function claudeEncodingForModel(modelId: string): Encoding | null {
|
||||
const bare = bareModelId(modelId);
|
||||
const parsed = parseAnthropicModel(bare);
|
||||
if (parsed) {
|
||||
if (parsed.kind === "opus") {
|
||||
if (semverGte(parsed.version, "5")) return Encoding.ClaudeV5;
|
||||
if (semverGte(parsed.version, "4.7")) return Encoding.ClaudeV47;
|
||||
return Encoding.ClaudeV3;
|
||||
}
|
||||
// Sonnet, Fable, and Mythos: the 4.7 family is opus-only, and their
|
||||
// 5-series message frame differs from opus-5's (measured live).
|
||||
return semverGte(parsed.version, "5") ? Encoding.ClaudeV5Sonnet : Encoding.ClaudeV3;
|
||||
}
|
||||
return /(^|[-/.:])claude([-.:]|$)/i.test(bare) ? Encoding.ClaudeV3 : null;
|
||||
}
|
||||
|
||||
/**
|
||||
* `strict` always pays for an exact native count (Claude ctok when the model
|
||||
* is Claude, o200k_base otherwise). `approximate` and `upperbound` prefer the
|
||||
* same exact count when a Claude encoding is known or `PI_TOKENIZER_ACCURATE=1`
|
||||
* is set, and otherwise fall back to a cheap heuristic: `approximate` a
|
||||
* bytes/4 guess, `upperbound` the raw byte length (never undercounts).
|
||||
*/
|
||||
export type TokenCountMode = "strict" | "approximate" | "upperbound";
|
||||
|
||||
/** Options for {@link Tokenizer.countMessage} / {@link Tokenizer.countMessages}. */
|
||||
export interface MessageCountOptions {
|
||||
/**
|
||||
* Drop opaque provider reasoning payloads (`thinkingSignature`,
|
||||
* `redactedThinking`, native server-tool blocks) from the estimate. Those
|
||||
* are billed by the provider on replay, so the default counts them — but
|
||||
* their *local* byte size can diverge wildly from what the provider
|
||||
* charges, so the compaction floor (which only needs the reliably-countable,
|
||||
* on-wire-compressible content) excludes them to avoid false triggers on
|
||||
* thinking-heavy turns.
|
||||
*/
|
||||
excludeEncryptedReasoning?: boolean;
|
||||
}
|
||||
|
||||
function byteEstimate(text: string): number {
|
||||
return (Buffer.byteLength(text, "utf-8") + 3) >> 2;
|
||||
}
|
||||
|
||||
export function countTokens(text: string | string[]): number {
|
||||
if (accurate) {
|
||||
return countTokensNat(text);
|
||||
} else if (Array.isArray(text)) {
|
||||
return text.reduce((sum, t) => sum + estimateTokens(t), 0);
|
||||
} else {
|
||||
return estimateTokens(text);
|
||||
}
|
||||
}
|
||||
|
||||
export function countTokensConservatively(text: string | string[]): number {
|
||||
if (accurate) {
|
||||
return countTokensNat(text);
|
||||
} else if (Array.isArray(text)) {
|
||||
return text.reduce((sum, value) => sum + Buffer.byteLength(value, "utf-8"), 0);
|
||||
} else {
|
||||
function byteLength(text: string): number {
|
||||
return Buffer.byteLength(text, "utf-8");
|
||||
}
|
||||
|
||||
function sumFragments(text: string | string[], perFragment: (t: string) => number): number {
|
||||
return Array.isArray(text) ? text.reduce((sum, t) => sum + perFragment(t), 0) : perFragment(text);
|
||||
}
|
||||
|
||||
/** Verdict from {@link Tokenizer.checkTokenBudget}. */
|
||||
export interface TokenBudgetCheck {
|
||||
/** Whether the text fits the budget. */
|
||||
fits: boolean;
|
||||
/**
|
||||
* Token count behind the verdict: the exact native count when `exact` is
|
||||
* set, otherwise the cheap byte upper bound (which already fit, so it is
|
||||
* only an over-estimate of a count known to be under budget).
|
||||
*/
|
||||
tokens: number;
|
||||
/** Whether the exact tokenizer had to run because the cheap bound busted. */
|
||||
exact: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Image content has no tokenizer representation; charge a fixed estimate
|
||||
* matching what providers typically bill for inline images.
|
||||
*/
|
||||
const IMAGE_TOKEN_ESTIMATE = 1200;
|
||||
|
||||
/**
|
||||
* Memoized estimates for one message under this tokenizer's encoding, split by
|
||||
* the {@link MessageCountOptions.excludeEncryptedReasoning} option so the two
|
||||
* variants never collide. `version` snapshots {@link messageEstimateVersion} at
|
||||
* write time; an owner mutation (prune/shake/strip-images) bumps the version,
|
||||
* which invalidates the entry in every live Tokenizer at once.
|
||||
*/
|
||||
interface MessageEstimate {
|
||||
version: number;
|
||||
default?: number;
|
||||
floored?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Model-aware local token counter. Immutable: the encoding is fixed at
|
||||
* construction, so a cached count can never straddle two encodings. An `Agent`
|
||||
* owns one for its active model (swapping the instance when the model's
|
||||
* encoding changes) and exposes it as `agent.tokenizer`; one-shot flows
|
||||
* (summarization, snapcompact sizing) construct their own for the model that
|
||||
* will be billed. Claude models get exact native ctok counts; everything else
|
||||
* keeps the fast byte estimate (or o200k when `PI_TOKENIZER_ACCURATE=1`).
|
||||
*/
|
||||
export class Tokenizer {
|
||||
readonly #encoding: Encoding | null;
|
||||
|
||||
/**
|
||||
* Per-message estimate memo. Keyed by message identity, deliberately not a
|
||||
* symbol-tagged property: callers spread messages to derive throwaway
|
||||
* variants for counting (`estimateBranchSummaryTokens` does
|
||||
* `countMessage({ ...message, content: truncated })`), and a property-borne
|
||||
* cache would ride along the spread. Identity keying gives clones a fresh
|
||||
* count.
|
||||
*/
|
||||
#estimates = new WeakMap<AgentMessage, MessageEstimate>();
|
||||
|
||||
constructor(modelId?: string | null) {
|
||||
this.#encoding = modelId ? claudeEncodingForModel(modelId) : null;
|
||||
}
|
||||
|
||||
get encoding(): Encoding | null {
|
||||
return this.#encoding;
|
||||
}
|
||||
|
||||
countTokens(text: string | string[], mode: TokenCountMode = "approximate"): number {
|
||||
if (mode === "strict") return countTokensNat(text, this.#encoding);
|
||||
if (!testEnv && this.#encoding !== null) return countTokensNat(text, this.#encoding);
|
||||
if (accurate) return countTokensNat(text);
|
||||
return sumFragments(text, mode === "upperbound" ? byteLength : byteEstimate);
|
||||
}
|
||||
|
||||
/**
|
||||
* Cheap-first budget probe — the way to ask "does this fit in `budget`
|
||||
* tokens?" without tokenizing the world.
|
||||
*
|
||||
* Byte length is a hard upper bound on token count (every token consumes at
|
||||
* least one input byte), so text whose raw bytes already fit the budget
|
||||
* cannot possibly exceed it — that verdict is returned without tokenizing at
|
||||
* all. Only text that busts the bound is ambiguous, and only that case pays
|
||||
* for the exact count. Since the bound overshoots ~4x on ordinary prose, the
|
||||
* common "comfortably under budget" answer is free.
|
||||
*/
|
||||
checkTokenBudget(text: string | string[], budget: number): TokenBudgetCheck {
|
||||
const bound = sumFragments(text, byteLength);
|
||||
if (bound <= budget) return { fits: true, tokens: bound, exact: false };
|
||||
const tokens = this.countTokens(text, "strict");
|
||||
return { fits: tokens <= budget, tokens, exact: true };
|
||||
}
|
||||
|
||||
/**
|
||||
* Token estimate for one message under this tokenizer's encoding.
|
||||
*
|
||||
* Settled historical messages are counted once and reused until an owner
|
||||
* (prune/shake/strip-images) calls `invalidateMessageCache`; streaming
|
||||
* assistants bypass the memo entirely (see the message-cache settle-gate
|
||||
* invariant). Image blocks charge a fixed per-image estimate.
|
||||
*/
|
||||
countMessage(message: AgentMessage, options?: MessageCountOptions): number {
|
||||
const floored = options?.excludeEncryptedReasoning === true;
|
||||
if (!isEstimateCacheable(message)) return this.#measureMessage(message, floored);
|
||||
const version = messageEstimateVersion(message);
|
||||
let entry = this.#estimates.get(message);
|
||||
if (entry === undefined || entry.version !== version) {
|
||||
entry = { version };
|
||||
this.#estimates.set(message, entry);
|
||||
}
|
||||
const cached = floored ? entry.floored : entry.default;
|
||||
if (cached !== undefined) return cached;
|
||||
const result = this.#measureMessage(message, floored);
|
||||
if (floored) entry.floored = result;
|
||||
else entry.default = result;
|
||||
return result;
|
||||
}
|
||||
|
||||
/** Sum of {@link countMessage} over `messages`. */
|
||||
countMessages(messages: readonly AgentMessage[], options?: MessageCountOptions): number {
|
||||
let total = 0;
|
||||
for (const message of messages) total += this.countMessage(message, options);
|
||||
return total;
|
||||
}
|
||||
|
||||
#measureMessage(message: AgentMessage, excludeEncryptedReasoning: boolean): number {
|
||||
const fragments: string[] = [];
|
||||
let extra = 0;
|
||||
// Declaration-merged app roles (the coding-agent's bashExecution) are
|
||||
// invisible to this package's union, so the discriminant is read as data.
|
||||
const role: string = message.role;
|
||||
if (role === "bashExecution") {
|
||||
if ("command" in message && typeof message.command === "string") fragments.push(message.command);
|
||||
if ("output" in message && typeof message.output === "string") fragments.push(message.output);
|
||||
return fragments.length === 0 ? 0 : this.countTokens(fragments);
|
||||
}
|
||||
|
||||
switch (message.role) {
|
||||
case "user": {
|
||||
const content: string | Array<{ type: string; text?: string }> = message.content;
|
||||
if (typeof content === "string") {
|
||||
fragments.push(content);
|
||||
} else if (Array.isArray(content)) {
|
||||
for (const block of content) {
|
||||
if (block.type === "text" && block.text) {
|
||||
fragments.push(block.text);
|
||||
}
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "assistant": {
|
||||
for (const block of message.content) {
|
||||
if (block.type === "text") {
|
||||
fragments.push(block.text);
|
||||
} else if (block.type === "thinking") {
|
||||
fragments.push(block.thinking);
|
||||
// Providers charge for the opaque signature/reasoning payload that
|
||||
// rides alongside the thinking text (OpenAI Responses encrypted
|
||||
// reasoning items, Anthropic signed thinking blocks, etc.). Without
|
||||
// counting it, this estimator can read ~half of the provider-reported
|
||||
// usage on thinking-heavy turns — see #2275 for the resulting
|
||||
// compaction-trigger / post-check metric divergence. The compaction
|
||||
// floor excludes it (its local byte size diverges from provider billing).
|
||||
if (block.thinkingSignature && !excludeEncryptedReasoning) {
|
||||
fragments.push(block.thinkingSignature);
|
||||
}
|
||||
} else if (block.type === "toolCall") {
|
||||
fragments.push(block.name);
|
||||
fragments.push(stringifyJson(block.arguments) ?? "null");
|
||||
} else if (block.type === "redactedThinking") {
|
||||
// Encrypted reasoning blob the provider still bills for on replay;
|
||||
// excluded from the compaction floor for the same reason as above.
|
||||
if (!excludeEncryptedReasoning) fragments.push(block.data);
|
||||
} else if (block.type === "anthropicServerTool") {
|
||||
// Native Anthropic server-tool call/result replayed verbatim on the
|
||||
// wire (server_tool_use input and opaque result content). The provider
|
||||
// still bills for it on same-provider replay; excluded from the
|
||||
// compaction floor like other encrypted reasoning because its local
|
||||
// byte size diverges from provider billing.
|
||||
if (!excludeEncryptedReasoning) fragments.push(stringifyJson(block.block) ?? "null");
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "hookMessage":
|
||||
case "toolResult": {
|
||||
if (typeof message.content === "string") {
|
||||
fragments.push(message.content);
|
||||
} else {
|
||||
for (const block of message.content) {
|
||||
if (block.type === "text" && block.text) {
|
||||
fragments.push(block.text);
|
||||
} else if (block.type === "image") {
|
||||
extra += IMAGE_TOKEN_ESTIMATE;
|
||||
}
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
case "branchSummary":
|
||||
case "compactionSummary": {
|
||||
fragments.push(message.summary);
|
||||
if (message.role === "compactionSummary") {
|
||||
if (message.blocks) {
|
||||
for (const block of message.blocks) {
|
||||
if (block.type === "text") fragments.push(block.text);
|
||||
else extra += snapcompact.FRAME_TOKEN_ESTIMATE;
|
||||
}
|
||||
} else if (message.images) {
|
||||
// Snapcompact frames render at ≥1568px; providers bill the downscaled cap.
|
||||
extra += message.images.length * snapcompact.FRAME_TOKEN_ESTIMATE;
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (fragments.length === 0) return extra;
|
||||
return extra + this.countTokens(fragments);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -208,7 +208,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions {
|
||||
* @example
|
||||
* ```typescript
|
||||
* transformContext: async (messages) => {
|
||||
* if (estimateTokens(messages) > MAX_TOKENS) {
|
||||
* if (agent.tokenizer.countMessages(messages) > MAX_TOKENS) {
|
||||
* return pruneOldMessages(messages);
|
||||
* }
|
||||
* return messages;
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import {
|
||||
type GenerateBranchSummaryOptions,
|
||||
generateBranchSummary,
|
||||
@@ -183,7 +184,7 @@ describe("branch summarization", () => {
|
||||
];
|
||||
|
||||
// Budget tight enough that the useless blob alone would blow it out.
|
||||
const { messages } = prepareBranchEntries(entries, 100);
|
||||
const { messages } = prepareBranchEntries(entries, new Tokenizer(), 100);
|
||||
|
||||
const userMessages = messages.filter((m): m is Extract<typeof m, { role: "user" }> => m.role === "user");
|
||||
expect(userMessages).toHaveLength(1);
|
||||
@@ -226,7 +227,7 @@ describe("branch summarization", () => {
|
||||
},
|
||||
];
|
||||
|
||||
const { messages } = prepareBranchEntries(entries, 700);
|
||||
const { messages } = prepareBranchEntries(entries, new Tokenizer(), 700);
|
||||
|
||||
expect(messages.some(m => m.role === "toolResult")).toBe(true);
|
||||
});
|
||||
|
||||
@@ -1,17 +1,18 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AgentMessage, Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { SessionMessageEntry } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import {
|
||||
applyShakeRegion,
|
||||
collectShakeRegions,
|
||||
DEFAULT_PRUNE_CONFIG,
|
||||
estimateTokens,
|
||||
invalidateMessageCache,
|
||||
isEstimateCacheable,
|
||||
pruneToolOutputs,
|
||||
} from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import type { AssistantMessage, ToolResultMessage, Usage } from "@oh-my-pi/pi-ai";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
let idCounter = 0;
|
||||
function nextId(): string {
|
||||
return `mc-${idCounter++}`;
|
||||
@@ -83,9 +84,9 @@ describe("estimate cache settle gate", () => {
|
||||
usage: usage(0),
|
||||
stopReason: "stop",
|
||||
};
|
||||
const before = estimateTokens(streaming as AgentMessage);
|
||||
const before = tokenizer.countMessage(streaming as AgentMessage);
|
||||
streaming.content = [{ type: "text", text: "first chunk plus a much longer continuation of streamed text" }];
|
||||
const after = estimateTokens(streaming as AgentMessage);
|
||||
const after = tokenizer.countMessage(streaming as AgentMessage);
|
||||
// Unsettled assistants never read the cache, so the grown content is recounted.
|
||||
expect(after).toBeGreaterThan(before);
|
||||
});
|
||||
@@ -103,12 +104,12 @@ describe("estimate cache option split", () => {
|
||||
};
|
||||
// Prime the default map first, then the floored one; the floored estimate
|
||||
// (which drops the encrypted-reasoning blob) must not read the default entry.
|
||||
const withBlob = estimateTokens(msg as AgentMessage);
|
||||
const floored = estimateTokens(msg as AgentMessage, { excludeEncryptedReasoning: true });
|
||||
const withBlob = tokenizer.countMessage(msg as AgentMessage);
|
||||
const floored = tokenizer.countMessage(msg as AgentMessage, { excludeEncryptedReasoning: true });
|
||||
expect(withBlob).toBeGreaterThan(floored + 500);
|
||||
// Cached reads return the same split values.
|
||||
expect(estimateTokens(msg as AgentMessage)).toBe(withBlob);
|
||||
expect(estimateTokens(msg as AgentMessage, { excludeEncryptedReasoning: true })).toBe(floored);
|
||||
expect(tokenizer.countMessage(msg as AgentMessage)).toBe(withBlob);
|
||||
expect(tokenizer.countMessage(msg as AgentMessage, { excludeEncryptedReasoning: true })).toBe(floored);
|
||||
});
|
||||
|
||||
test("counts native server-tool blocks by default and drops them from the compaction floor", () => {
|
||||
@@ -127,12 +128,12 @@ describe("estimate cache option split", () => {
|
||||
},
|
||||
],
|
||||
};
|
||||
const textOnly = estimateTokens({
|
||||
const textOnly = tokenizer.countMessage({
|
||||
...settledAssistant("x"),
|
||||
content: [{ type: "text", text: "answer" }],
|
||||
} as AgentMessage);
|
||||
const withServerTool = estimateTokens(msg as AgentMessage);
|
||||
const floored = estimateTokens(msg as AgentMessage, { excludeEncryptedReasoning: true });
|
||||
const withServerTool = tokenizer.countMessage(msg as AgentMessage);
|
||||
const floored = tokenizer.countMessage(msg as AgentMessage, { excludeEncryptedReasoning: true });
|
||||
// Default estimate charges for the serialized server-tool payload…
|
||||
expect(withServerTool).toBeGreaterThan(floored + 500);
|
||||
// …while the compaction floor ignores the opaque encrypted blob entirely.
|
||||
@@ -144,24 +145,28 @@ describe("estimate cache invalidation seams", () => {
|
||||
test("pruneToolOutputs drops the cached estimate of a pruned result", () => {
|
||||
const big = toolResult("x".repeat(20_000));
|
||||
const entries = [messageEntry(big as AgentMessage)];
|
||||
const before = estimateTokens(big as AgentMessage);
|
||||
const before = tokenizer.countMessage(big as AgentMessage);
|
||||
expect(before).toBeGreaterThan(1000);
|
||||
|
||||
const result = pruneToolOutputs(entries, { ...DEFAULT_PRUNE_CONFIG, protectTokens: 0, minimumSavings: 0 });
|
||||
const result = pruneToolOutputs(entries, tokenizer, {
|
||||
...DEFAULT_PRUNE_CONFIG,
|
||||
protectTokens: 0,
|
||||
minimumSavings: 0,
|
||||
});
|
||||
expect(result.prunedCount).toBe(1);
|
||||
|
||||
// After the in-place prune the estimate must reflect the short placeholder,
|
||||
// not the stale full-content count.
|
||||
const after = estimateTokens(big as AgentMessage);
|
||||
const after = tokenizer.countMessage(big as AgentMessage);
|
||||
expect(after).toBeLessThan(before);
|
||||
});
|
||||
|
||||
test("applyShakeRegion drops the cached estimate of a shaken result", () => {
|
||||
const big = toolResult(`\`\`\`ts\n${"const value = compute(a, b, c, d, e);\n".repeat(400)}\`\`\``);
|
||||
const entry = messageEntry(big as AgentMessage);
|
||||
const before = estimateTokens(big as AgentMessage);
|
||||
const before = tokenizer.countMessage(big as AgentMessage);
|
||||
|
||||
const regions = collectShakeRegions([entry], {
|
||||
const regions = collectShakeRegions([entry], tokenizer, {
|
||||
protectTokens: 0,
|
||||
minSavings: 0,
|
||||
protectedTools: [],
|
||||
@@ -170,18 +175,23 @@ describe("estimate cache invalidation seams", () => {
|
||||
expect(regions.length).toBeGreaterThan(0);
|
||||
applyShakeRegion(regions[0], "[shaken]");
|
||||
|
||||
const after = estimateTokens(big as AgentMessage);
|
||||
const after = tokenizer.countMessage(big as AgentMessage);
|
||||
expect(after).toBeLessThan(before);
|
||||
});
|
||||
|
||||
test("explicit invalidateMessageCache forces a recount", () => {
|
||||
test("explicit invalidateMessageCache forces a recount in every tokenizer instance", () => {
|
||||
const result = toolResult("original content here");
|
||||
const before = estimateTokens(result as AgentMessage);
|
||||
// Mutate content directly (simulating an owner rewrite) then invalidate.
|
||||
const second = new Tokenizer();
|
||||
const before = tokenizer.countMessage(result as AgentMessage);
|
||||
expect(second.countMessage(result as AgentMessage)).toBe(before);
|
||||
// Mutate content directly (simulating an owner rewrite); without
|
||||
// invalidation both instances still return their stale memo.
|
||||
result.content = [{ type: "text", text: "a much longer replacement body that should count higher than before" }];
|
||||
// Without invalidation the stale cached value would still be returned.
|
||||
expect(estimateTokens(result as AgentMessage)).toBe(before);
|
||||
expect(tokenizer.countMessage(result as AgentMessage)).toBe(before);
|
||||
expect(second.countMessage(result as AgentMessage)).toBe(before);
|
||||
// One version-tag bump invalidates the memo in BOTH instances.
|
||||
invalidateMessageCache(result as AgentMessage);
|
||||
expect(estimateTokens(result as AgentMessage)).toBeGreaterThan(before);
|
||||
expect(tokenizer.countMessage(result as AgentMessage)).toBeGreaterThan(before);
|
||||
expect(second.countMessage(result as AgentMessage)).toBeGreaterThan(before);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import {
|
||||
type CompactionPreparation,
|
||||
compact,
|
||||
@@ -630,7 +631,7 @@ describe("remote compaction input forwarding", () => {
|
||||
{ fetch: fetchMock },
|
||||
);
|
||||
|
||||
const trimmed = trimRemoteCompactionInputToContextWindow(nativeInput, 1_000, "compact");
|
||||
const trimmed = trimRemoteCompactionInputToContextWindow(nativeInput, new Tokenizer(), 1_000, "compact");
|
||||
expect(trimmed.estimatedTokensAfter).toBeLessThanOrEqual(1_000);
|
||||
expect(requestInput?.some(item => item.type === "custom_tool_call")).toBe(true);
|
||||
expect(requestInput?.find(item => item.type === "custom_tool_call_output")?.output).toBe(
|
||||
@@ -646,7 +647,7 @@ describe("remote compaction input forwarding", () => {
|
||||
{ type: "function_call_output", call_id: "call_2", output: "b".repeat(8_000) },
|
||||
];
|
||||
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, 1_000, "compact");
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, new Tokenizer(), 1_000, "compact");
|
||||
|
||||
expect(result.rewrittenOutputs).toBe(2);
|
||||
expect(result.input.slice(0, 2)).toEqual(input.slice(0, 2));
|
||||
@@ -666,7 +667,7 @@ describe("remote compaction input forwarding", () => {
|
||||
{ type: "function_call_output", call_id: "call_2", output: "useful latest result" },
|
||||
];
|
||||
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, 1_000, "compact");
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, new Tokenizer(), 1_000, "compact");
|
||||
|
||||
expect(result.rewrittenOutputs).toBe(0);
|
||||
expect(result.input).toEqual(input);
|
||||
@@ -686,7 +687,7 @@ describe("remote compaction input forwarding", () => {
|
||||
{ type: "function_call_output", call_id: "call_1", output: "useful result" },
|
||||
];
|
||||
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, 15_000, "compact");
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, new Tokenizer(), 15_000, "compact");
|
||||
|
||||
expect(result.rewrittenOutputs).toBe(0);
|
||||
expect(result.input).toEqual(input);
|
||||
@@ -698,7 +699,7 @@ describe("remote compaction input forwarding", () => {
|
||||
const output = Array.from({ length: 1_000 }, (_, index) => index.toString(16).padStart(8, "0")).join("");
|
||||
const input = [{ type: "function_call_output", call_id: "call_1", output }];
|
||||
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, 3_000, "compact");
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, new Tokenizer(), 3_000, "compact");
|
||||
|
||||
expect(result.estimatedTokensBefore).toBeGreaterThan(3_000);
|
||||
expect(result.rewrittenOutputs).toBe(1);
|
||||
@@ -720,7 +721,7 @@ describe("remote compaction input forwarding", () => {
|
||||
attachment,
|
||||
];
|
||||
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, 15_000, "compact");
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, new Tokenizer(), 15_000, "compact");
|
||||
|
||||
expect(result.rewrittenOutputs).toBe(1);
|
||||
expect(result.input[0].output).toBe(CONTEXT_WINDOW_TRUNCATED_OUTPUT_MESSAGE);
|
||||
@@ -740,7 +741,7 @@ describe("remote compaction input forwarding", () => {
|
||||
{ type: "function_call_output", call_id: "call_1", output: "useful result".repeat(2_000) },
|
||||
];
|
||||
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, 15_000, "compact");
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, new Tokenizer(), 15_000, "compact");
|
||||
|
||||
expect(result.estimatedTokensBefore).toBeGreaterThan(15_000);
|
||||
expect(result.rewrittenOutputs).toBe(1);
|
||||
@@ -750,7 +751,7 @@ describe("remote compaction input forwarding", () => {
|
||||
test("returns semantically unchanged input when it already fits", () => {
|
||||
const input = [{ type: "function_call_output", call_id: "call_1", output: "small" }];
|
||||
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, 1_000, "compact");
|
||||
const result = trimRemoteCompactionInputToContextWindow(input, new Tokenizer(), 1_000, "compact");
|
||||
|
||||
expect(result.rewrittenOutputs).toBe(0);
|
||||
expect(result.input).toEqual(input);
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AgentMessage, Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { SessionEntry, SessionMessageEntry, ShakeConfig } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import {
|
||||
AGGRESSIVE_SHAKE_CONFIG,
|
||||
@@ -7,11 +7,12 @@ import {
|
||||
applyShakeRegions,
|
||||
collectShakeRegions,
|
||||
DEFAULT_SHAKE_CONFIG,
|
||||
estimateTokens,
|
||||
RESCUE_SHAKE_CONFIG,
|
||||
} from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import type { AssistantMessage, TextContent, ToolCall, ToolResultMessage } from "@oh-my-pi/pi-ai";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
let idCounter = 0;
|
||||
function nextId(): string {
|
||||
return `entry-${idCounter++}`;
|
||||
@@ -74,7 +75,7 @@ describe("collectShakeRegions — tool results", () => {
|
||||
test("collects unprotected tool results and applyShakeRegion sets prunedAt", () => {
|
||||
const tr = toolResultMessage("bash", "x".repeat(400));
|
||||
const entry = messageEntry(tr);
|
||||
const regions = collectShakeRegions([entry], cfg());
|
||||
const regions = collectShakeRegions([entry], tokenizer, cfg());
|
||||
|
||||
expect(regions).toHaveLength(1);
|
||||
const region = regions[0];
|
||||
@@ -88,13 +89,13 @@ describe("collectShakeRegions — tool results", () => {
|
||||
|
||||
test("never collects protected tools", () => {
|
||||
const entry = messageEntry(toolResultMessage("skill", "y".repeat(800)));
|
||||
const regions = collectShakeRegions([entry], cfg({ protectedTools: ["skill"] }));
|
||||
const regions = collectShakeRegions([entry], tokenizer, cfg({ protectedTools: ["skill"] }));
|
||||
expect(regions).toHaveLength(0);
|
||||
});
|
||||
|
||||
test("never collects already-pruned tool results", () => {
|
||||
const entry = messageEntry(toolResultMessage("bash", "z".repeat(800), { prunedAt: Date.now() }));
|
||||
const regions = collectShakeRegions([entry], cfg());
|
||||
const regions = collectShakeRegions([entry], tokenizer, cfg());
|
||||
expect(regions).toHaveLength(0);
|
||||
});
|
||||
|
||||
@@ -103,9 +104,13 @@ describe("collectShakeRegions — tool results", () => {
|
||||
const older = messageEntry(toolResultMessage("bash", text));
|
||||
const middle = messageEntry(toolResultMessage("bash", text));
|
||||
const recent = messageEntry(toolResultMessage("bash", text));
|
||||
const perEntry = estimateTokens(older.message);
|
||||
const perEntry = tokenizer.countMessage(older.message);
|
||||
// Window covers the most recent ~1.5 entries → middle & recent protected, older eligible.
|
||||
const regions = collectShakeRegions([older, middle, recent], cfg({ protectTokens: Math.floor(perEntry * 1.5) }));
|
||||
const regions = collectShakeRegions(
|
||||
[older, middle, recent],
|
||||
tokenizer,
|
||||
cfg({ protectTokens: Math.floor(perEntry * 1.5) }),
|
||||
);
|
||||
|
||||
expect(regions).toHaveLength(1);
|
||||
expect(regions[0].entry).toBe(older);
|
||||
@@ -113,9 +118,9 @@ describe("collectShakeRegions — tool results", () => {
|
||||
|
||||
test("minSavings gates the whole batch", () => {
|
||||
const entry = messageEntry(toolResultMessage("bash", "q".repeat(800)));
|
||||
const tokens = estimateTokens(entry.message);
|
||||
expect(collectShakeRegions([entry], cfg({ minSavings: tokens * 10 }))).toHaveLength(0);
|
||||
expect(collectShakeRegions([entry], cfg({ minSavings: 0 }))).toHaveLength(1);
|
||||
const tokens = tokenizer.countMessage(entry.message);
|
||||
expect(collectShakeRegions([entry], tokenizer, cfg({ minSavings: tokens * 10 }))).toHaveLength(0);
|
||||
expect(collectShakeRegions([entry], tokenizer, cfg({ minSavings: 0 }))).toHaveLength(1);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -124,7 +129,7 @@ describe("collectShakeRegions — fenced / XML blocks", () => {
|
||||
const fence = fencedBlock(120);
|
||||
const text = `intro line\n${fence}\noutro line`;
|
||||
const entry = messageEntry(assistantMessage([{ type: "text", text }]));
|
||||
const regions = collectShakeRegions([entry], cfg());
|
||||
const regions = collectShakeRegions([entry], tokenizer, cfg());
|
||||
|
||||
expect(regions).toHaveLength(1);
|
||||
const region = regions[0];
|
||||
@@ -140,14 +145,14 @@ describe("collectShakeRegions — fenced / XML blocks", () => {
|
||||
test("ignores fenced blocks below fenceMinTokens", () => {
|
||||
const text = "intro\n```ts\nconst a = 1;\n```\noutro";
|
||||
const entry = messageEntry(assistantMessage([{ type: "text", text }]));
|
||||
expect(collectShakeRegions([entry], cfg({ fenceMinTokens: 400 }))).toHaveLength(0);
|
||||
expect(collectShakeRegions([entry], tokenizer, cfg({ fenceMinTokens: 400 }))).toHaveLength(0);
|
||||
});
|
||||
|
||||
test("detects a top-level XML block", () => {
|
||||
const xml = xmlBlock(120);
|
||||
const text = `before\n${xml}\nafter`;
|
||||
const entry = messageEntry(assistantMessage([{ type: "text", text }]));
|
||||
const regions = collectShakeRegions([entry], cfg());
|
||||
const regions = collectShakeRegions([entry], tokenizer, cfg());
|
||||
|
||||
expect(regions).toHaveLength(1);
|
||||
const region = regions[0];
|
||||
@@ -161,7 +166,7 @@ describe("collectShakeRegions — fenced / XML blocks", () => {
|
||||
const entry = messageEntry(
|
||||
assistantMessage([{ type: "text", text: "tiny" }, toolCall, { type: "text", text: `pre\n${fence}\npost` }]),
|
||||
);
|
||||
const regions = collectShakeRegions([entry], cfg());
|
||||
const regions = collectShakeRegions([entry], tokenizer, cfg());
|
||||
|
||||
expect(regions).toHaveLength(1);
|
||||
const region = regions[0];
|
||||
@@ -172,7 +177,7 @@ describe("collectShakeRegions — fenced / XML blocks", () => {
|
||||
test("does not cross message boundaries — each large block stays in its own entry", () => {
|
||||
const a = messageEntry(assistantMessage([{ type: "text", text: `a\n${fencedBlock(120)}\na` }]));
|
||||
const b = messageEntry(assistantMessage([{ type: "text", text: `b\n${fencedBlock(120, "py")}\nb` }]));
|
||||
const regions = collectShakeRegions([a, b], cfg());
|
||||
const regions = collectShakeRegions([a, b], tokenizer, cfg());
|
||||
|
||||
expect(regions).toHaveLength(2);
|
||||
expect(regions[0].entry).toBe(a);
|
||||
@@ -182,7 +187,7 @@ describe("collectShakeRegions — fenced / XML blocks", () => {
|
||||
test("ignores unterminated fences (conservative)", () => {
|
||||
const text = `intro\n\`\`\`ts\n${"const a = 1;\n".repeat(60)}`; // never closes
|
||||
const entry = messageEntry(assistantMessage([{ type: "text", text }]));
|
||||
expect(collectShakeRegions([entry], cfg())).toHaveLength(0);
|
||||
expect(collectShakeRegions([entry], tokenizer, cfg())).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -192,7 +197,7 @@ describe("applyShakeRegions — multi-region ordering", () => {
|
||||
const second = fencedBlock(80, "py");
|
||||
const text = `head\n${first}\nmiddle\n${second}\ntail`;
|
||||
const entry = messageEntry(assistantMessage([{ type: "text", text }]));
|
||||
const regions = collectShakeRegions([entry], cfg());
|
||||
const regions = collectShakeRegions([entry], tokenizer, cfg());
|
||||
expect(regions).toHaveLength(2);
|
||||
|
||||
applyShakeRegions([
|
||||
@@ -214,7 +219,7 @@ describe("shake config presets", () => {
|
||||
test("manual shake preserves the recent tool-result tail instead of stripping everything", () => {
|
||||
const older = messageEntry(toolResultMessage("bash", "old-result ".repeat(300)));
|
||||
const recent = messageEntry(toolResultMessage("bash", "recent-result ".repeat(3000)));
|
||||
const regions = collectShakeRegions([older, recent], AGGRESSIVE_SHAKE_CONFIG);
|
||||
const regions = collectShakeRegions([older, recent], tokenizer, AGGRESSIVE_SHAKE_CONFIG);
|
||||
|
||||
// The recent result sits inside the preserved tail; the older one is
|
||||
// still shaken aggressively.
|
||||
@@ -229,13 +234,13 @@ describe("shake config presets", () => {
|
||||
|
||||
test("rescue preset overrides the manual tail so it can elide the newest result", () => {
|
||||
const recent = messageEntry(toolResultMessage("bash", "oversized-result ".repeat(2000)));
|
||||
const regions = collectShakeRegions([recent], RESCUE_SHAKE_CONFIG);
|
||||
const regions = collectShakeRegions([recent], tokenizer, RESCUE_SHAKE_CONFIG);
|
||||
expect(regions).toHaveLength(1);
|
||||
expect(regions[0].entry).toBe(recent);
|
||||
});
|
||||
|
||||
test("empty branch yields no regions", () => {
|
||||
expect(collectShakeRegions([] as SessionEntry[], AGGRESSIVE_SHAKE_CONFIG)).toHaveLength(0);
|
||||
expect(collectShakeRegions([] as SessionEntry[], tokenizer, AGGRESSIVE_SHAKE_CONFIG)).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -245,13 +250,13 @@ describe("collectShakeRegions — useless results", () => {
|
||||
const flagged = messageEntry(toolResultMessage("search", text, { useless: true }));
|
||||
const plain = messageEntry(toolResultMessage("search", text));
|
||||
// Window far larger than the whole branch: only the flagged result bypasses it.
|
||||
const regions = collectShakeRegions([flagged, plain], cfg({ protectTokens: 1_000_000 }));
|
||||
const regions = collectShakeRegions([flagged, plain], tokenizer, cfg({ protectTokens: 1_000_000 }));
|
||||
expect(regions).toHaveLength(1);
|
||||
expect(regions[0].entry).toBe(flagged);
|
||||
});
|
||||
|
||||
test("an error result never bypasses the window even when flagged", () => {
|
||||
const entry = messageEntry(toolResultMessage("search", "boom\n".repeat(50), { useless: true, isError: true }));
|
||||
expect(collectShakeRegions([entry], cfg({ protectTokens: 1_000_000 }))).toHaveLength(0);
|
||||
expect(collectShakeRegions([entry], tokenizer, cfg({ protectTokens: 1_000_000 }))).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
import { describe, expect, it } from "bun:test";
|
||||
import type { ImageContent } from "@oh-my-pi/pi-ai";
|
||||
import * as snapcompact from "@oh-my-pi/snapcompact";
|
||||
import { estimateTokens } from "../src/compaction/compaction";
|
||||
import { createCompactionSummaryMessage, defaultConvertToLlm } from "../src/compaction/messages";
|
||||
import { Tokenizer } from "../src/tokenizer";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
describe("compaction summary message with snapcompact frames", () => {
|
||||
const images: ImageContent[] = [
|
||||
@@ -10,7 +12,7 @@ describe("compaction summary message with snapcompact frames", () => {
|
||||
{ type: "image", data: "ZmFrZTI=", mimeType: "image/png" },
|
||||
];
|
||||
|
||||
it("estimateTokens charges per attached frame", () => {
|
||||
it("countMessage charges per attached frame", () => {
|
||||
const bare = createCompactionSummaryMessage("summary text", 1000, new Date().toISOString());
|
||||
const withFrames = createCompactionSummaryMessage(
|
||||
"summary text",
|
||||
@@ -20,7 +22,9 @@ describe("compaction summary message with snapcompact frames", () => {
|
||||
undefined,
|
||||
images,
|
||||
);
|
||||
expect(estimateTokens(withFrames) - estimateTokens(bare)).toBe(2 * snapcompact.FRAME_TOKEN_ESTIMATE);
|
||||
expect(tokenizer.countMessage(withFrames) - tokenizer.countMessage(bare)).toBe(
|
||||
2 * snapcompact.FRAME_TOKEN_ESTIMATE,
|
||||
);
|
||||
});
|
||||
|
||||
it("defaultConvertToLlm appends frames as image blocks after the summary text", () => {
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AgentMessage, Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { SessionEntry, SessionMessageEntry } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import {
|
||||
DEFAULT_PRUNE_CONFIG,
|
||||
@@ -13,6 +13,8 @@ import {
|
||||
import type { ProtectedToolContext } from "@oh-my-pi/pi-agent-core/compaction/tool-protection";
|
||||
import type { AssistantMessage, TextContent, ToolResultMessage } from "@oh-my-pi/pi-ai";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
let idCounter = 0;
|
||||
function nextId(): string {
|
||||
return `entry-${idCounter++}`;
|
||||
@@ -138,7 +140,7 @@ describe("pruneSupersededToolResults — tail case", () => {
|
||||
const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000);
|
||||
const entries: SessionEntry[] = [call1, result1, call2, result2];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ now: T0 + 1_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ now: T0 + 1_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(1);
|
||||
expect(result.tokensSaved).toBeGreaterThan(0);
|
||||
@@ -156,7 +158,7 @@ describe("pruneSupersededToolResults — tail case", () => {
|
||||
const big = textEntry(BIG_TEXT, T0 + 2_000);
|
||||
const entries: SessionEntry[] = [call1, result1, call2, result2, big];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ suffixTokenLimit: 200, now: T0 + 2_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ suffixTokenLimit: 200, now: T0 + 2_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(0);
|
||||
expect(result.tokensSaved).toBe(0);
|
||||
@@ -176,6 +178,7 @@ describe("pruneSupersededToolResults — tail case", () => {
|
||||
// Suffix limit 0 would block every candidate; only the idle gap fires.
|
||||
const result = pruneSupersededToolResults(
|
||||
entries,
|
||||
tokenizer,
|
||||
cfg({ suffixTokenLimit: 0, idleFlushMs: 30 * 60_000, now: T0 + 4_000 + 30 * 60_000 }),
|
||||
);
|
||||
|
||||
@@ -194,6 +197,7 @@ describe("pruneSupersededToolResults — tail case", () => {
|
||||
|
||||
const result = pruneSupersededToolResults(
|
||||
entries,
|
||||
tokenizer,
|
||||
cfg({ suffixTokenLimit: 0, idleFlushMs: 30 * 60_000, now: T0 + 2_000 + 29 * 60_000 }),
|
||||
);
|
||||
|
||||
@@ -210,7 +214,7 @@ describe("pruneSupersededToolResults — selectors", () => {
|
||||
let entries: SessionEntry[] = [callA, resultA, callB, resultB];
|
||||
|
||||
// Different selectors: no candidates.
|
||||
let result = pruneSupersededToolResults(entries, cfg({ now: T0 + 1_000 }));
|
||||
let result = pruneSupersededToolResults(entries, tokenizer, cfg({ now: T0 + 1_000 }));
|
||||
expect(result.prunedCount).toBe(0);
|
||||
expect(resultText(resultA)).toBe(FILE_CONTENT);
|
||||
expect(resultText(resultB)).toBe(FILE_CONTENT);
|
||||
@@ -218,7 +222,7 @@ describe("pruneSupersededToolResults — selectors", () => {
|
||||
// Identical selector strings DO supersede.
|
||||
const [callA2, resultA2] = readPair("src/foo.ts:50-200", FILE_CONTENT, T0 + 2_000);
|
||||
entries = [...entries, callA2, resultA2];
|
||||
result = pruneSupersededToolResults(entries, cfg({ now: T0 + 2_000 }));
|
||||
result = pruneSupersededToolResults(entries, tokenizer, cfg({ now: T0 + 2_000 }));
|
||||
expect(result.prunedCount).toBe(1);
|
||||
expect(resultText(resultA)).toBe(SUPERSEDED_NOTICE);
|
||||
expect(resultText(resultB)).toBe(FILE_CONTENT);
|
||||
@@ -227,7 +231,7 @@ describe("pruneSupersededToolResults — selectors", () => {
|
||||
// A later selector-free read supersedes every selector-carrying read of the base path.
|
||||
const [callFull, resultFull] = readPair("src/foo.ts", FILE_CONTENT, T0 + 3_000);
|
||||
entries = [...entries, callFull, resultFull];
|
||||
result = pruneSupersededToolResults(entries, cfg({ now: T0 + 3_000 }));
|
||||
result = pruneSupersededToolResults(entries, tokenizer, cfg({ now: T0 + 3_000 }));
|
||||
expect(result.prunedCount).toBe(2);
|
||||
expect(resultText(resultB)).toBe(SUPERSEDED_NOTICE);
|
||||
expect(resultText(resultA2)).toBe(SUPERSEDED_NOTICE);
|
||||
@@ -239,7 +243,7 @@ describe("pruneSupersededToolResults — selectors", () => {
|
||||
const [callRange, resultRange] = readPair("src/foo.ts:50-200", FILE_CONTENT, T0 + 1_000);
|
||||
const entries: SessionEntry[] = [callFull, resultFull, callRange, resultRange];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ now: T0 + 1_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ now: T0 + 1_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(0);
|
||||
expect(resultText(resultFull)).toBe(FILE_CONTENT);
|
||||
@@ -254,7 +258,7 @@ describe("pruneSupersededToolResults — protection & latest", () => {
|
||||
const [call3, result3] = readPair("src/foo.ts", FILE_CONTENT, T0 + 2_000);
|
||||
const entries: SessionEntry[] = [call1, result1, call2, result2, call3, result3];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ now: T0 + 2_000 + 60 * 60_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ now: T0 + 2_000 + 60 * 60_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(2);
|
||||
expect(resultText(result1)).toBe(SUPERSEDED_NOTICE);
|
||||
@@ -281,7 +285,11 @@ describe("pruneSupersededToolResults — protection & latest", () => {
|
||||
fooResult2,
|
||||
];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ protectedTools: [protectPlan], now: T0 + 3_000 }));
|
||||
const result = pruneSupersededToolResults(
|
||||
entries,
|
||||
tokenizer,
|
||||
cfg({ protectedTools: [protectPlan], now: T0 + 3_000 }),
|
||||
);
|
||||
|
||||
expect(result.prunedCount).toBe(1);
|
||||
expect(resultText(planResult1)).toBe(FILE_CONTENT);
|
||||
@@ -297,7 +305,7 @@ describe("pruneSupersededToolResults — protection & latest", () => {
|
||||
const entries: SessionEntry[] = [call1, result1, call2, result2];
|
||||
|
||||
// The only newer same-key read is itself pruned -> result1 has no live superseder.
|
||||
const result = pruneSupersededToolResults(entries, cfg({ now: T0 + 2_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ now: T0 + 2_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(0);
|
||||
expect(resultText(result1)).toBe(FILE_CONTENT);
|
||||
@@ -310,7 +318,7 @@ describe("pruneToolOutputs — supersede priority fold", () => {
|
||||
const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000);
|
||||
const entries: SessionEntry[] = [call1, result1, call2, result2];
|
||||
|
||||
const result = pruneToolOutputs(entries, {
|
||||
const result = pruneToolOutputs(entries, tokenizer, {
|
||||
protectTokens: 1_000_000, // everything inside the protect window
|
||||
minimumSavings: 0,
|
||||
protectedTools: [],
|
||||
@@ -335,7 +343,7 @@ describe("pruneToolOutputs — supersede priority fold", () => {
|
||||
|
||||
// Protect window covers everything: nothing pruned, superseded reads included.
|
||||
const protectedFixture = buildEntries();
|
||||
const protectedRun = pruneToolOutputs(protectedFixture.entries, {
|
||||
const protectedRun = pruneToolOutputs(protectedFixture.entries, tokenizer, {
|
||||
protectTokens: 1_000_000,
|
||||
minimumSavings: 0,
|
||||
protectedTools: [],
|
||||
@@ -347,7 +355,7 @@ describe("pruneToolOutputs — supersede priority fold", () => {
|
||||
// Protect window empty: every result past it pruned with the legacy
|
||||
// truncation placeholder — never the supersede placeholder.
|
||||
const unprotectedFixture = buildEntries();
|
||||
const unprotectedRun = pruneToolOutputs(unprotectedFixture.entries, {
|
||||
const unprotectedRun = pruneToolOutputs(unprotectedFixture.entries, tokenizer, {
|
||||
protectTokens: 0,
|
||||
minimumSavings: 0,
|
||||
protectedTools: [],
|
||||
@@ -375,6 +383,7 @@ describe("pruneSupersededToolResults — useless results", () => {
|
||||
// Suffix limit 0 blocks the tail rule; only the idle gap fires.
|
||||
const result = pruneSupersededToolResults(
|
||||
entries,
|
||||
tokenizer,
|
||||
cfg({ pruneUseless: true, suffixTokenLimit: 0, now: T0 + 1_000 + 31 * 60_000 }),
|
||||
);
|
||||
|
||||
@@ -388,7 +397,7 @@ describe("pruneSupersededToolResults — useless results", () => {
|
||||
const [call1, result1] = uselessPair("search", NO_MATCH_TEXT, T0);
|
||||
const entries: SessionEntry[] = [call1, result1];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ pruneUseless: true, now: T0 + 1_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ pruneUseless: true, now: T0 + 1_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(1);
|
||||
expect(resultText(result1)).toBe(USELESS_NOTICE);
|
||||
@@ -401,6 +410,7 @@ describe("pruneSupersededToolResults — useless results", () => {
|
||||
|
||||
const result = pruneSupersededToolResults(
|
||||
entries,
|
||||
tokenizer,
|
||||
cfg({ pruneUseless: true, suffixTokenLimit: 200, now: T0 + 2_000 }),
|
||||
);
|
||||
|
||||
@@ -413,7 +423,7 @@ describe("pruneSupersededToolResults — useless results", () => {
|
||||
const [call1, result1] = uselessPair("search", "No matches found", T0);
|
||||
const entries: SessionEntry[] = [call1, result1];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ pruneUseless: true, now: T0 + 31 * 60_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ pruneUseless: true, now: T0 + 31 * 60_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(0);
|
||||
expect(resultText(result1)).toBe("No matches found");
|
||||
@@ -425,6 +435,7 @@ describe("pruneSupersededToolResults — useless results", () => {
|
||||
|
||||
const result = pruneSupersededToolResults(
|
||||
entries,
|
||||
tokenizer,
|
||||
cfg({ pruneUseless: true, protectedTools: ["search"], now: T0 + 31 * 60_000 }),
|
||||
);
|
||||
|
||||
@@ -436,7 +447,7 @@ describe("pruneSupersededToolResults — useless results", () => {
|
||||
const [call1, result1] = uselessPair("search", NO_MATCH_TEXT, T0);
|
||||
const entries: SessionEntry[] = [call1, result1];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, {
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, {
|
||||
protectedTools: [],
|
||||
pruneUseless: true,
|
||||
now: T0 + 1_000,
|
||||
@@ -452,7 +463,7 @@ describe("pruneSupersededToolResults — useless results", () => {
|
||||
const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000);
|
||||
const entries: SessionEntry[] = [call1, result1, call2, result2];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ pruneUseless: true, now: T0 + 1_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ pruneUseless: true, now: T0 + 1_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(1);
|
||||
expect(resultText(result1)).toBe(SUPERSEDED_NOTICE);
|
||||
@@ -463,7 +474,7 @@ describe("pruneSupersededToolResults — useless results", () => {
|
||||
const [call1, result1] = uselessPair("search", NO_MATCH_TEXT, T0, { isError: true });
|
||||
const entries: SessionEntry[] = [call1, result1];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ pruneUseless: true, now: T0 + 31 * 60_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ pruneUseless: true, now: T0 + 31 * 60_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(0);
|
||||
expect(resultText(result1)).toBe(NO_MATCH_TEXT);
|
||||
@@ -476,7 +487,7 @@ describe("pruneToolOutputs — useless results", () => {
|
||||
const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000);
|
||||
const entries: SessionEntry[] = [call1, result1, call2, result2];
|
||||
|
||||
const result = pruneToolOutputs(entries, {
|
||||
const result = pruneToolOutputs(entries, tokenizer, {
|
||||
protectTokens: 1_000_000, // everything inside the protect window
|
||||
minimumSavings: 0,
|
||||
protectedTools: [],
|
||||
@@ -492,7 +503,7 @@ describe("pruneToolOutputs — useless results", () => {
|
||||
const [call1, result1] = uselessPair("search", NO_MATCH_TEXT, T0);
|
||||
const entries: SessionEntry[] = [call1, result1];
|
||||
|
||||
const result = pruneToolOutputs(entries, {
|
||||
const result = pruneToolOutputs(entries, tokenizer, {
|
||||
protectTokens: 1_000_000,
|
||||
minimumSavings: 0,
|
||||
protectedTools: [],
|
||||
@@ -513,7 +524,7 @@ describe("pruneToolOutputs — small-result floor", () => {
|
||||
const entries: SessionEntry[] = [tinyCall, tinyResult, bigCall, bigResult];
|
||||
|
||||
// Protect window empty and zero savings threshold: only size keeps the tiny one.
|
||||
const result = pruneToolOutputs(entries, { protectTokens: 0, minimumSavings: 0, protectedTools: [] });
|
||||
const result = pruneToolOutputs(entries, tokenizer, { protectTokens: 0, minimumSavings: 0, protectedTools: [] });
|
||||
|
||||
expect(result.prunedCount).toBe(1);
|
||||
expect(resultText(tinyResult)).toBe("ok");
|
||||
@@ -543,14 +554,14 @@ describe("cache-stable boundary — warm prefix protection", () => {
|
||||
|
||||
// Legacy (no cacheWarmSuffixTokens): superseded result1 bypasses the window -> pruned.
|
||||
const legacy = build();
|
||||
const legacyRun = pruneToolOutputs(legacy.entries, base);
|
||||
const legacyRun = pruneToolOutputs(legacy.entries, tokenizer, base);
|
||||
expect(legacyRun.prunedCount).toBe(1);
|
||||
expect(resultText(legacy.result1)).toBe(SUPERSEDED_NOTICE);
|
||||
|
||||
// Guard armed: result1's all-message suffix (BIG_TEXT + call2 + result2) far
|
||||
// exceeds the window, so it is part of the warm cached prefix and is kept.
|
||||
const guarded = build();
|
||||
const guardedRun = pruneToolOutputs(guarded.entries, { ...base, cacheWarmSuffixTokens: 200 });
|
||||
const guardedRun = pruneToolOutputs(guarded.entries, tokenizer, { ...base, cacheWarmSuffixTokens: 200 });
|
||||
expect(guardedRun.prunedCount).toBe(0);
|
||||
expect(resultText(guarded.result1)).toBe(FILE_CONTENT);
|
||||
expect(resultMessage(guarded.result1).prunedAt).toBeUndefined();
|
||||
@@ -562,7 +573,7 @@ describe("cache-stable boundary — warm prefix protection", () => {
|
||||
const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000);
|
||||
const entries: SessionEntry[] = [call1, result1, big, call2, result2];
|
||||
|
||||
const result = pruneToolOutputs(entries, {
|
||||
const result = pruneToolOutputs(entries, tokenizer, {
|
||||
protectTokens: 1_000_000,
|
||||
minimumSavings: 0,
|
||||
protectedTools: [],
|
||||
@@ -583,7 +594,7 @@ describe("cache-stable boundary — warm prefix protection", () => {
|
||||
const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000);
|
||||
const entries: SessionEntry[] = [call1, result1, call2, result2];
|
||||
|
||||
const result = pruneToolOutputs(entries, {
|
||||
const result = pruneToolOutputs(entries, tokenizer, {
|
||||
protectTokens: 1_000_000,
|
||||
minimumSavings: 0,
|
||||
protectedTools: [],
|
||||
@@ -601,7 +612,7 @@ describe("cache-stable boundary — warm prefix protection", () => {
|
||||
const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000);
|
||||
const entries: SessionEntry[] = [call1, result1, call2, result2];
|
||||
|
||||
const result = pruneSupersededToolResults(entries, cfg({ keepBoundaryId: call1.id, now: T0 + 1_000 }));
|
||||
const result = pruneSupersededToolResults(entries, tokenizer, cfg({ keepBoundaryId: call1.id, now: T0 + 1_000 }));
|
||||
|
||||
expect(result.prunedCount).toBe(1);
|
||||
expect(resultText(result1)).toBe(SUPERSEDED_NOTICE);
|
||||
@@ -620,6 +631,7 @@ describe("cache-stable boundary — warm prefix protection", () => {
|
||||
// Cold cache (idle > threshold) with suffixTokenLimit 0: only the idle path can fire.
|
||||
const result = pruneSupersededToolResults(
|
||||
entries,
|
||||
tokenizer,
|
||||
cfg({
|
||||
keepBoundaryId: call2.id,
|
||||
suffixTokenLimit: 0,
|
||||
@@ -642,7 +654,7 @@ describe("cache-stable boundary — warm prefix protection", () => {
|
||||
|
||||
// protectTokens 0 -> the age path would prune both; the window is wide so the
|
||||
// guard does not protect either; only keepBoundaryId shields result1.
|
||||
pruneToolOutputs(entries, {
|
||||
pruneToolOutputs(entries, tokenizer, {
|
||||
protectTokens: 0,
|
||||
minimumSavings: 0,
|
||||
protectedTools: [],
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import { Encoding } from "@oh-my-pi/pi-natives";
|
||||
import { claudeEncodingForModel, Tokenizer } from "../src/tokenizer";
|
||||
|
||||
// Contract: local token counting must pick the ctok family that matches the
|
||||
// model's tokenizer generation (v3 for Claude 3 … Opus 4.6 and every
|
||||
// non-opus < 5, v4.7 for Opus 4.7–4.9, v5 for the 5-series). A wrong family
|
||||
// silently skews every context-budget and compaction decision for that model.
|
||||
describe("claudeEncodingForModel", () => {
|
||||
test("opus routes on the 4.7 and 5.0 version thresholds", () => {
|
||||
expect(claudeEncodingForModel("claude-opus-4-5")).toBe(Encoding.ClaudeV3);
|
||||
expect(claudeEncodingForModel("claude-opus-4-6")).toBe(Encoding.ClaudeV3);
|
||||
expect(claudeEncodingForModel("claude-opus-4-7")).toBe(Encoding.ClaudeV47);
|
||||
expect(claudeEncodingForModel("claude-opus-4-9-20260101")).toBe(Encoding.ClaudeV47);
|
||||
expect(claudeEncodingForModel("claude-opus-5")).toBe(Encoding.ClaudeV5);
|
||||
});
|
||||
|
||||
test("non-opus kinds skip the opus-only 4.7 family and use the sonnet-5 frame", () => {
|
||||
expect(claudeEncodingForModel("claude-sonnet-4-5-20250929")).toBe(Encoding.ClaudeV3);
|
||||
expect(claudeEncodingForModel("claude-sonnet-5")).toBe(Encoding.ClaudeV5Sonnet);
|
||||
expect(claudeEncodingForModel("claude-fable-5")).toBe(Encoding.ClaudeV5Sonnet);
|
||||
});
|
||||
test("unclassifiable claude ids fall back to v3; provider prefixes are stripped", () => {
|
||||
expect(claudeEncodingForModel("claude-3-5-haiku-20241022")).toBe(Encoding.ClaudeV3);
|
||||
expect(claudeEncodingForModel("claude-haiku-4-5")).toBe(Encoding.ClaudeV3);
|
||||
expect(claudeEncodingForModel("anthropic/claude-opus-4-7")).toBe(Encoding.ClaudeV47);
|
||||
});
|
||||
|
||||
test("non-claude models get no ctok encoding", () => {
|
||||
expect(claudeEncodingForModel("gpt-5.4")).toBeNull();
|
||||
expect(claudeEncodingForModel("gemini-3-pro")).toBeNull();
|
||||
expect(claudeEncodingForModel("glm-4.7")).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("Tokenizer", () => {
|
||||
test("defaults to null encoding and byte estimation", () => {
|
||||
const tokenizer = new Tokenizer();
|
||||
expect(tokenizer.encoding).toBeNull();
|
||||
expect(tokenizer.countTokens("hello world")).toBe(3);
|
||||
});
|
||||
|
||||
test("encoding is fixed at construction per model id", () => {
|
||||
expect(new Tokenizer("claude-opus-4-7").encoding).toBe(Encoding.ClaudeV47);
|
||||
expect(new Tokenizer("claude-opus-5").encoding).toBe(Encoding.ClaudeV5);
|
||||
expect(new Tokenizer("gpt-5.4").encoding).toBeNull();
|
||||
expect(new Tokenizer(undefined).encoding).toBeNull();
|
||||
});
|
||||
|
||||
test("separate instances do not interfere with each other", () => {
|
||||
const t1 = new Tokenizer("claude-opus-4-7");
|
||||
const t2 = new Tokenizer("claude-opus-5");
|
||||
const t3 = new Tokenizer("gpt-5.4");
|
||||
|
||||
expect(t1.encoding).toBe(Encoding.ClaudeV47);
|
||||
expect(t2.encoding).toBe(Encoding.ClaudeV5);
|
||||
expect(t3.encoding).toBeNull();
|
||||
|
||||
const t4 = new Tokenizer("claude-sonnet-4-5-20250929");
|
||||
expect(t4.encoding).toBe(Encoding.ClaudeV3);
|
||||
expect(t1.encoding).toBe(Encoding.ClaudeV47);
|
||||
expect(t2.encoding).toBe(Encoding.ClaudeV5);
|
||||
expect(t3.encoding).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("countTokens with modes", () => {
|
||||
test("approximate mode uses fast estimation", () => {
|
||||
const tokenizer = new Tokenizer();
|
||||
expect(tokenizer.countTokens("hello world", "approximate")).toBe(3);
|
||||
});
|
||||
|
||||
test("upperbound mode uses byte length", () => {
|
||||
const tokenizer = new Tokenizer();
|
||||
expect(tokenizer.countTokens("hello world", "upperbound")).toBe(11);
|
||||
});
|
||||
|
||||
test("strict mode uses native counting regardless of encoding", () => {
|
||||
const noEncoding = new Tokenizer();
|
||||
expect(noEncoding.countTokens("hello world", "strict")).toBe(2);
|
||||
|
||||
const claudeEncoding = new Tokenizer("claude-opus-4-7");
|
||||
expect(claudeEncoding.countTokens("hello world", "strict")).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
test("mode is per-call; encoding stays independently model-scoped in strict mode", () => {
|
||||
// approximate/upperbound skip the encoding entirely under NODE_ENV=test
|
||||
// (fast estimate for a snappy suite); strict is testEnv-independent, so
|
||||
// it is the mode that proves per-instance encoding isolation here.
|
||||
const claude = new Tokenizer("claude-opus-4-7");
|
||||
const generic = new Tokenizer("gpt-5.4");
|
||||
expect(claude.countTokens("hello world", "strict")).not.toBe(generic.countTokens("hello world", "strict"));
|
||||
});
|
||||
});
|
||||
@@ -1,9 +1,12 @@
|
||||
import { describe, expect, it } from "bun:test";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { SessionMessageEntry } from "@oh-my-pi/pi-agent-core/compaction/entries";
|
||||
import { DEFAULT_PRUNE_CONFIG, pruneToolOutputs } from "@oh-my-pi/pi-agent-core/compaction/pruning";
|
||||
import { AGGRESSIVE_SHAKE_CONFIG, collectShakeRegions } from "@oh-my-pi/pi-agent-core/compaction/shake";
|
||||
import type { AssistantMessage, TextContent, ToolResultMessage, Usage } from "@oh-my-pi/pi-ai";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
function usage(): Usage {
|
||||
return {
|
||||
input: 0,
|
||||
@@ -61,7 +64,11 @@ describe("conditional tool-result protection", () => {
|
||||
fileResult,
|
||||
];
|
||||
|
||||
const result = pruneToolOutputs(entries, { ...DEFAULT_PRUNE_CONFIG, protectTokens: 0, minimumSavings: 0 });
|
||||
const result = pruneToolOutputs(entries, tokenizer, {
|
||||
...DEFAULT_PRUNE_CONFIG,
|
||||
protectTokens: 0,
|
||||
minimumSavings: 0,
|
||||
});
|
||||
|
||||
expect(result.prunedCount).toBe(1);
|
||||
expect((skillResult.message as ToolResultMessage).prunedAt).toBeUndefined();
|
||||
@@ -86,7 +93,7 @@ describe("conditional tool-result protection", () => {
|
||||
|
||||
// protectTokens: 0 isolates the matcher behavior from the aggressive
|
||||
// preset's recent-tail window (covered by shake.test.ts).
|
||||
const regions = collectShakeRegions(entries, { ...AGGRESSIVE_SHAKE_CONFIG, protectTokens: 0 });
|
||||
const regions = collectShakeRegions(entries, tokenizer, { ...AGGRESSIVE_SHAKE_CONFIG, protectTokens: 0 });
|
||||
|
||||
expect(regions).toHaveLength(1);
|
||||
expect(regions[0]?.kind).toBe("toolResult");
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
import { describe, expect, test } from "bun:test";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import {
|
||||
estimateTranscriptTokens,
|
||||
findTranscriptUsageAnchor,
|
||||
isTranscriptUsageAnchor,
|
||||
} from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import type { AgentMessage } from "@oh-my-pi/pi-agent-core/types";
|
||||
import type { AssistantMessage, Usage } from "@oh-my-pi/pi-ai";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
function usage(overrides: Partial<Usage> = {}): Usage {
|
||||
return {
|
||||
input: 1_000,
|
||||
output: 200,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 1_200,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function assistant(text: string, overrides: Partial<AssistantMessage> = {}): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text }],
|
||||
api: "mock",
|
||||
provider: "mock",
|
||||
model: "mock-model",
|
||||
usage: usage(),
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function user(text: string): AgentMessage {
|
||||
return { role: "user", content: text, timestamp: Date.now() };
|
||||
}
|
||||
|
||||
// Contract: transcript sizing charges the provider's own report for the prefix
|
||||
// it already measured and tokenizes only the tail appended after it. Getting
|
||||
// this wrong either re-tokenizes the whole history every turn (the cost this
|
||||
// exists to avoid) or double-counts the anchored prefix into a false
|
||||
// compaction trigger.
|
||||
describe("findTranscriptUsageAnchor", () => {
|
||||
test("anchors on the newest settled assistant with usable usage", () => {
|
||||
const older = assistant("older");
|
||||
const newer = assistant("newer");
|
||||
const anchor = findTranscriptUsageAnchor([user("a"), older, user("b"), newer, user("c")]);
|
||||
expect(anchor?.message).toBe(newer);
|
||||
expect(anchor?.index).toBe(3);
|
||||
expect(anchor?.tokens).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
test("skips aborted and errored turns and falls back to an older usable one", () => {
|
||||
const good = assistant("good");
|
||||
const messages = [
|
||||
good,
|
||||
assistant("aborted", { stopReason: "aborted" }),
|
||||
assistant("errored", { stopReason: "error" }),
|
||||
];
|
||||
expect(findTranscriptUsageAnchor(messages)?.message).toBe(good);
|
||||
});
|
||||
|
||||
test("rejects usage carrying no usable context numbers", () => {
|
||||
const empty = assistant("empty", {
|
||||
usage: usage({ input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 }),
|
||||
});
|
||||
expect(isTranscriptUsageAnchor(empty)).toBe(false);
|
||||
expect(findTranscriptUsageAnchor([empty])).toBeUndefined();
|
||||
});
|
||||
|
||||
test("fromIndex hides stale pre-compaction usage from the scan", () => {
|
||||
const stale = assistant("summarized away");
|
||||
const messages = [stale, user("kept")];
|
||||
expect(findTranscriptUsageAnchor(messages)?.message).toBe(stale);
|
||||
expect(findTranscriptUsageAnchor(messages, 1)).toBeUndefined();
|
||||
});
|
||||
|
||||
test("non-assistant roles never anchor", () => {
|
||||
expect(isTranscriptUsageAnchor(user("hi"))).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("estimateTranscriptTokens", () => {
|
||||
test("charges provider usage for the prefix and tokenizes only the tail", () => {
|
||||
const anchorMessage = assistant("anchored turn");
|
||||
const tail = user("x".repeat(4_000));
|
||||
const anchored = estimateTranscriptTokens([user("ancient"), anchorMessage, tail], tokenizer);
|
||||
|
||||
// Provider report for the prefix, plus a local count of the tail only —
|
||||
// the pre-anchor "ancient" message is already inside the report.
|
||||
const anchorTokens = findTranscriptUsageAnchor([anchorMessage])?.tokens ?? 0;
|
||||
expect(anchored).toBeGreaterThan(anchorTokens);
|
||||
expect(anchored).toBe(anchorTokens + tokenizer.countMessages([tail]));
|
||||
});
|
||||
|
||||
test("the local floor ignores provider usage so it cannot be deflated", () => {
|
||||
const messages = [user("a".repeat(400)), assistant("reply"), user("b".repeat(400))];
|
||||
const anchored = estimateTranscriptTokens(messages, tokenizer);
|
||||
const local = tokenizer.countMessages(messages);
|
||||
// The anchored arm inherits the provider's 1.2k-token report; the local
|
||||
// floor counts only message bytes, so the two must not coincide.
|
||||
expect(anchored).not.toBe(local);
|
||||
expect(local).toBeLessThan(anchored);
|
||||
});
|
||||
|
||||
test("unanchored transcripts fall back to counting every message from countFromIndex", () => {
|
||||
const messages = [user("a".repeat(400)), user("b".repeat(400))];
|
||||
const all = estimateTranscriptTokens(messages, tokenizer);
|
||||
const skipFirst = estimateTranscriptTokens(messages, tokenizer, { countFromIndex: 1 });
|
||||
expect(all).toBeGreaterThan(skipFirst);
|
||||
expect(skipFirst).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
test("a hidden anchor still counts the whole transcript locally", () => {
|
||||
// anchorFromIndex gates only the anchor scan; content accounting stays
|
||||
// governed by countFromIndex so a stale-usage guard cannot silently drop
|
||||
// real messages from the estimate.
|
||||
const messages = [assistant("stale usage"), user("y".repeat(400))];
|
||||
const hidden = estimateTranscriptTokens(messages, tokenizer, { anchorFromIndex: messages.length });
|
||||
expect(hidden).toBe(tokenizer.countMessages(messages));
|
||||
});
|
||||
});
|
||||
@@ -9,6 +9,8 @@
|
||||
### Changed
|
||||
|
||||
- `/settings` rows can now carry a risk note: a warning glyph on the row plus a warning-colored line above the description. `External Thinking` (`externalThinking`, `--external-thinking`) is the first user — providers have flagged the request shape it produces as abuse, up to account-level enforcement, so both the settings entry and `--help` now say so.
|
||||
- Token counting is now scoped to the model being billed rather than to a process-global tokenizer: session maintenance, stats, advisors, `/context`, snapcompact inline imaging, and `compress` each count through the owning agent's `Tokenizer` (`agent.tokenizer`). Message counting is `Tokenizer.countMessage`/`countMessages` (replacing the free `estimateTokens(message, tokenizer)` helper; the legacy shim keeps a compat `estimateTokens` export for legacy pi extensions). `estimateToolSchemaTokens`, `estimateSkillsTokens`, `computeNonMessageTokens`, and `computeNonMessageBreakdown` take an explicit tokenizer; `scripts/measure-prompt-tokens.ts` accepts an optional model id (argv) so its numbers match what that model is charged.
|
||||
- The advisor runtime's `maintainContext` hook now receives the pending update as a message instead of a pre-computed token count — sizing it needs the advisor model's tokenizer, which the host owns.
|
||||
|
||||
## [17.3.8] - 2026-08-19
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
* Before each model call and during compaction accounting, the agent walks the
|
||||
* full live `AgentMessage[]` history through:
|
||||
* 1. `convertToLlm(messages)` — role-specific conversion into provider `Message[]`.
|
||||
* 2. `estimateTokens(message)` — cl100k-style token counting for prune/shake/floors.
|
||||
* 2. `Tokenizer.countMessage(message)` — cl100k-style token counting for prune/shake/floors.
|
||||
*
|
||||
* In a long session those historical objects are settled, yet before the memo
|
||||
* both paths recompute from scratch on every pass. This bench measures cold
|
||||
@@ -24,11 +24,12 @@
|
||||
* Env: `LLM_ASSEMBLY_N` overrides the history length (default 5000);
|
||||
* `PI_TOKENIZER_ACCURATE=1` uses the native cl100k tokenizer.
|
||||
*/
|
||||
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import { estimateTokens } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import { type AgentMessage, Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { AssistantMessage, ToolResultMessage, Usage } from "@oh-my-pi/pi-ai";
|
||||
import { convertToLlm } from "../src/session/messages";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
const N = Number(Bun.env.LLM_ASSEMBLY_N ?? 5000);
|
||||
const WARMUP = 5;
|
||||
const SAMPLES = 25;
|
||||
@@ -134,12 +135,6 @@ function sample<T>(makeWorkload: () => T, run: (workload: T) => void, batch = 1)
|
||||
return stats(samples);
|
||||
}
|
||||
|
||||
function estimateAll(messages: AgentMessage[]): number {
|
||||
let total = 0;
|
||||
for (const m of messages) total += estimateTokens(m);
|
||||
return total;
|
||||
}
|
||||
|
||||
console.log(`\nBenchmark: llm-assembly (N=${N}, warmup=${WARMUP}, samples=${SAMPLES})\n`);
|
||||
|
||||
// ─── convertToLlm ─────────────────────────────────────────────────────────────
|
||||
@@ -187,21 +182,21 @@ const convertGrow = sample(
|
||||
},
|
||||
);
|
||||
|
||||
// ─── estimateTokens ───────────────────────────────────────────────────────────
|
||||
// ─── Tokenizer.countMessage ───────────────────────────────────────────────────
|
||||
// Cold: fresh-identity history per sample → every estimate is a cache miss.
|
||||
const estimateFirst = sample(
|
||||
() => buildHistory(N),
|
||||
history => {
|
||||
estimateAll(history);
|
||||
tokenizer.countMessages(history);
|
||||
},
|
||||
);
|
||||
// Warm: one history, primed once, re-counted every sample from the cache.
|
||||
const warmEstimate = buildHistory(N);
|
||||
estimateAll(warmEstimate);
|
||||
tokenizer.countMessages(warmEstimate);
|
||||
const estimateSecond = sample(
|
||||
() => warmEstimate,
|
||||
history => {
|
||||
estimateAll(history);
|
||||
tokenizer.countMessages(history);
|
||||
},
|
||||
);
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { countTokens } from "@oh-my-pi/pi-agent-core";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import { estimateToolSchemaTokens } from "@oh-my-pi/pi-coding-agent/modes/utils/context-usage";
|
||||
import { buildSystemPrompt } from "@oh-my-pi/pi-coding-agent/system-prompt";
|
||||
@@ -14,6 +14,11 @@ function est(s: string): number {
|
||||
await Settings.init({ inMemory: true, cwd: process.cwd() });
|
||||
const settings = Settings.isolated({});
|
||||
|
||||
// Optional model id (argv[2]) scopes the counter to that model's tokenizer, so
|
||||
// the numbers match what the agent will actually be charged. Without it the
|
||||
// counts are the fast byte estimate the runtime uses for non-Claude models.
|
||||
const tokenizer = new Tokenizer(process.argv[2]);
|
||||
|
||||
const session: ToolSession = {
|
||||
cwd: process.cwd(),
|
||||
hasUI: false,
|
||||
@@ -29,14 +34,14 @@ console.log(`active tools (${tools.length}): ${tools.map(t => t.name).join(", ")
|
||||
|
||||
const rows: Array<{ name: string; descBytes: number; tok: number; schemaTok: number }> = [];
|
||||
for (const t of tools) {
|
||||
const tok = estimateToolSchemaTokens([t as never]);
|
||||
const tok = estimateToolSchemaTokens([t as never], tokenizer);
|
||||
const descBytes = bytes(t.description ?? "");
|
||||
const descTok = est(t.description ?? "");
|
||||
rows.push({ name: t.name, descBytes, tok, schemaTok: tok - descTok });
|
||||
}
|
||||
rows.sort((a, b) => b.tok - a.tok);
|
||||
|
||||
const totalTok = estimateToolSchemaTokens(tools as never);
|
||||
const totalTok = estimateToolSchemaTokens(tools as never, tokenizer);
|
||||
console.log("per-tool tokens (sorted): name | total tok | desc bytes | ~schema tok");
|
||||
for (const r of rows) {
|
||||
console.log(
|
||||
@@ -59,5 +64,5 @@ const parts = built.systemPrompt;
|
||||
const part0 = parts[0] ?? "";
|
||||
const rest = parts.slice(1).join("\n");
|
||||
console.log(`system prompt parts: ${parts.length}`);
|
||||
console.log(`SYSTEM PROMPT tokens (part0, no skills): ${countTokens(part0)} (bytes=${bytes(part0)})`);
|
||||
console.log(`SYSTEM CONTEXT tokens (parts[1..]): ${countTokens(rest)} (bytes=${bytes(rest)})`);
|
||||
console.log(`SYSTEM PROMPT tokens (part0, no skills): ${tokenizer.countTokens(part0)} (bytes=${bytes(part0)})`);
|
||||
console.log(`SYSTEM CONTEXT tokens (parts[1..]): ${tokenizer.countTokens(rest)} (bytes=${bytes(rest)})`);
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import { estimateTokens } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import type { AssistantMessage, ImageContent, TextContent } from "@oh-my-pi/pi-ai";
|
||||
import * as AIError from "@oh-my-pi/pi-ai/error";
|
||||
import { raceWithSignal } from "@oh-my-pi/pi-ai/utils/abort";
|
||||
@@ -48,9 +47,12 @@ export interface AdvisorRuntimeHost {
|
||||
* when the advisor must clear its own context before sending the current
|
||||
* incremental update. The cursor stays at the current primary position: this
|
||||
* recovery path must never replay the full primary transcript.
|
||||
*
|
||||
* Takes the pending update as a message rather than a token count: sizing it
|
||||
* needs the advisor model's tokenizer, which the host owns.
|
||||
* Optional: hosts that omit it get no proactive maintenance.
|
||||
*/
|
||||
maintainContext?(incomingTokens: number, signal: AbortSignal): Promise<boolean>;
|
||||
maintainContext?(incoming: AgentMessage, signal: AbortSignal): Promise<boolean>;
|
||||
/**
|
||||
* Called immediately before each `agent.prompt(batch)` cycle. Lets the host
|
||||
* clear per-update advisor state and apply the in-progress delivery policy.
|
||||
@@ -962,10 +964,12 @@ export class AdvisorRuntime {
|
||||
for (let round = 0; round < MAX_COALESCE_ROUNDS; round++) {
|
||||
if (this.#sessionTransitionPaused) break;
|
||||
if (this.host.maintainContext) {
|
||||
const incomingTokens = estimateTokens({ role: "user", content: batchText, timestamp: Date.now() });
|
||||
let shouldResetContext = false;
|
||||
try {
|
||||
shouldResetContext = await this.host.maintainContext(incomingTokens, signal);
|
||||
shouldResetContext = await this.host.maintainContext(
|
||||
{ role: "user", content: batchText, timestamp: Date.now() },
|
||||
signal,
|
||||
);
|
||||
} catch (err) {
|
||||
logger.debug("advisor context maintenance failed", { err: String(err) });
|
||||
}
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
* // …drive a session, then read protocol.latest / protocol.approved
|
||||
*/
|
||||
import { type } from "@oh-my-pi/omptype";
|
||||
import { countTokens } from "@oh-my-pi/pi-agent-core";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { ToolDefinition } from "../extensibility/extensions";
|
||||
import approveDescription from "../prompts/tools/approve.md" with { type: "text" };
|
||||
import rewriteDescription from "../prompts/tools/rewrite.md" with { type: "text" };
|
||||
@@ -66,6 +66,7 @@ function words(text: string): number {
|
||||
|
||||
/** Draft ledger shared by the protocol tools and the command loop. */
|
||||
export class CompressProtocol {
|
||||
readonly #tokenizer: Tokenizer;
|
||||
readonly #sourceWords: number;
|
||||
readonly #sourceTokens: number;
|
||||
readonly #drafts: CompressDraft[] = [];
|
||||
@@ -73,9 +74,16 @@ export class CompressProtocol {
|
||||
#approved = false;
|
||||
#verdict: string | undefined;
|
||||
|
||||
constructor(source: string) {
|
||||
/**
|
||||
* `modelId` scopes the token counter to the compressing model. Metrics are
|
||||
* source-vs-draft ratios measured with one counter, so they stay coherent
|
||||
* even when the model is unknown at construction time (the session that
|
||||
* resolves it is built from this protocol).
|
||||
*/
|
||||
constructor(source: string, modelId?: string | null) {
|
||||
this.#tokenizer = new Tokenizer(modelId);
|
||||
this.#sourceWords = words(source);
|
||||
this.#sourceTokens = countTokens(source);
|
||||
this.#sourceTokens = this.#tokenizer.countTokens(source);
|
||||
}
|
||||
|
||||
/** Newest submitted draft, or undefined before the first `rewrite`. */
|
||||
@@ -110,7 +118,7 @@ export class CompressProtocol {
|
||||
|
||||
/** Size of `draft` against the source. */
|
||||
metrics(draft: CompressDraft): CompressMetrics {
|
||||
const draftTokens = countTokens(draft.text);
|
||||
const draftTokens = this.#tokenizer.countTokens(draft.text);
|
||||
return {
|
||||
sourceWords: this.#sourceWords,
|
||||
draftWords: words(draft.text),
|
||||
|
||||
@@ -15,7 +15,13 @@
|
||||
import { Database } from "bun:sqlite";
|
||||
import * as fs from "node:fs";
|
||||
import * as path from "node:path";
|
||||
import type { AgentToolResult, AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
|
||||
import {
|
||||
type AgentMessage,
|
||||
type AgentToolResult,
|
||||
type AgentToolUpdateCallback,
|
||||
type MessageCountOptions,
|
||||
Tokenizer,
|
||||
} from "@oh-my-pi/pi-agent-core";
|
||||
import { type AuthCredential, SqliteAuthCredentialStore, type TSchema } from "@oh-my-pi/pi-ai";
|
||||
import { piEscapeRegexLiteral, piJoinPath } from "@oh-my-pi/pi-ai/providers/cursor-pi-args";
|
||||
import { getKeybindings, type Keybinding, Text } from "@oh-my-pi/pi-tui";
|
||||
@@ -1449,11 +1455,24 @@ export function getPackageDir(): string {
|
||||
|
||||
// Legacy pi's `@earendil-works/pi-coding-agent` re-exported `estimateTokens`,
|
||||
// `compact`, and `serializeConversation` from its package root (via
|
||||
// `./core/compaction/index.ts`). In omp they live in
|
||||
// `@oh-my-pi/pi-agent-core/compaction`, and the coding-agent barrel below does
|
||||
// not forward them, so legacy extensions importing them fail Bun's static
|
||||
// export check during validation (issues #6583, #7174, #7403).
|
||||
export { compact, estimateTokens, serializeConversation } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
// `./core/compaction/index.ts`). In omp `compact` and `serializeConversation`
|
||||
// live in `@oh-my-pi/pi-agent-core/compaction`, and the coding-agent barrel
|
||||
// below does not forward them, so legacy extensions importing them fail Bun's
|
||||
// static export check during validation (issues #6583, #7174, #7403).
|
||||
export { compact, serializeConversation } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
|
||||
const legacyTokenizer = new Tokenizer();
|
||||
|
||||
/**
|
||||
* Legacy `estimateTokens(message, tokenizer?, options?)` export. The core API
|
||||
* became `Tokenizer.countMessage`, but legacy pi extensions still import this
|
||||
* free function by name (issues #6583, #7174, #7403), so the export surface
|
||||
* must survive; a shared model-agnostic Tokenizer backs the tokenizer-less
|
||||
* legacy call shape.
|
||||
*/
|
||||
export function estimateTokens(message: AgentMessage, tokenizer?: Tokenizer, options?: MessageCountOptions): number {
|
||||
return (tokenizer ?? legacyTokenizer).countMessage(message, options);
|
||||
}
|
||||
|
||||
// Same barrel gap for two more legacy package-root exports: pi re-exported the
|
||||
// `CONFIG_DIR_NAME` constant and the CLI parser `parseArgs`. In omp
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { countTokens } from "@oh-my-pi/pi-agent-core";
|
||||
import type { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { CompactionSettings } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import { effectiveReserveTokens, estimateTokens, resolveThresholdTokens } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import { effectiveReserveTokens, resolveThresholdTokens } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import type { Tool as AiTool, Model } from "@oh-my-pi/pi-ai";
|
||||
import { toolWireSchema } from "@oh-my-pi/pi-ai/utils/schema";
|
||||
import { formatNumber } from "@oh-my-pi/pi-utils";
|
||||
@@ -72,18 +72,19 @@ function renderedSkills(
|
||||
return skills.filter(skill => skill.hide !== true);
|
||||
}
|
||||
|
||||
export function estimateSkillsTokens(skills: readonly Skill[]): number {
|
||||
export function estimateSkillsTokens(skills: readonly Skill[], tokenizer: Tokenizer): number {
|
||||
const fragments: string[] = [];
|
||||
for (const skill of skills) {
|
||||
// "- name: description\n" wire framing tokenizes ~identically to the
|
||||
// concatenated form, so encode each piece separately and sum.
|
||||
fragments.push(skill.name, skill.description);
|
||||
}
|
||||
return countTokens(fragments);
|
||||
return tokenizer.countTokens(fragments);
|
||||
}
|
||||
|
||||
export function estimateToolSchemaTokens(
|
||||
tools: ReadonlyArray<Pick<Tool, "name" | "description" | "parameters">>,
|
||||
tokenizer: Tokenizer,
|
||||
): number {
|
||||
const fragments: string[] = [];
|
||||
for (const tool of tools) {
|
||||
@@ -99,7 +100,7 @@ export function estimateToolSchemaTokens(
|
||||
// Schema may contain functions or cycles; ignore.
|
||||
}
|
||||
}
|
||||
return countTokens(fragments);
|
||||
return tokenizer.countTokens(fragments);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -127,6 +128,9 @@ interface NonMessageTokenCache {
|
||||
systemPromptRef: readonly string[];
|
||||
toolsRef: ReadonlyArray<Pick<Tool, "name" | "description" | "parameters">>;
|
||||
skillsRef: readonly Skill[];
|
||||
// The Agent swaps its Tokenizer instance when the model's encoding changes,
|
||||
// so instance identity doubles as the encoding key.
|
||||
tokenizerRef: Tokenizer;
|
||||
tokens: number | undefined;
|
||||
breakdown:
|
||||
| {
|
||||
@@ -144,7 +148,7 @@ interface CachedNonMessageTokenSource extends NonMessageTokenSource {
|
||||
[NON_MESSAGE_TOKEN_CACHE]?: NonMessageTokenCache;
|
||||
}
|
||||
|
||||
function nonMessageTokenCacheEntry(session: NonMessageTokenSource): NonMessageTokenCache {
|
||||
function nonMessageTokenCacheEntry(session: NonMessageTokenSource, tokenizer: Tokenizer): NonMessageTokenCache {
|
||||
const cachedSession: CachedNonMessageTokenSource = session;
|
||||
const systemPromptRef = session.systemPrompt ?? EMPTY_STRING_PARTS;
|
||||
const toolsRef = session.agent?.state?.tools ?? EMPTY_TOOLS;
|
||||
@@ -154,21 +158,22 @@ function nonMessageTokenCacheEntry(session: NonMessageTokenSource): NonMessageTo
|
||||
entry &&
|
||||
entry.systemPromptRef === systemPromptRef &&
|
||||
entry.toolsRef === toolsRef &&
|
||||
entry.skillsRef === skillsRef
|
||||
entry.skillsRef === skillsRef &&
|
||||
entry.tokenizerRef === tokenizer
|
||||
) {
|
||||
return entry;
|
||||
}
|
||||
entry = { systemPromptRef, toolsRef, skillsRef, tokens: undefined, breakdown: undefined };
|
||||
entry = { systemPromptRef, toolsRef, skillsRef, tokenizerRef: tokenizer, tokens: undefined, breakdown: undefined };
|
||||
cachedSession[NON_MESSAGE_TOKEN_CACHE] = entry;
|
||||
return entry;
|
||||
}
|
||||
|
||||
export function computeNonMessageTokens(session: NonMessageTokenSource): number {
|
||||
const entry = nonMessageTokenCacheEntry(session);
|
||||
export function computeNonMessageTokens(session: NonMessageTokenSource, tokenizer: Tokenizer): number {
|
||||
const entry = nonMessageTokenCacheEntry(session, tokenizer);
|
||||
if (entry.tokens !== undefined) return entry.tokens;
|
||||
const systemPromptParts = session.systemPrompt ?? EMPTY_STRING_PARTS;
|
||||
const tools = session.agent?.state?.tools ?? EMPTY_TOOLS;
|
||||
const tokens = countTokens(systemPromptParts) + estimateToolSchemaTokens(tools);
|
||||
const tokens = tokenizer.countTokens(systemPromptParts) + estimateToolSchemaTokens(tools, tokenizer);
|
||||
entry.tokens = tokens;
|
||||
return tokens;
|
||||
}
|
||||
@@ -179,20 +184,23 @@ export function computeNonMessageTokens(session: NonMessageTokenSource): number
|
||||
* the status-line fast path intentionally uses the equivalent collapsed total
|
||||
* in `computeNonMessageTokens`.
|
||||
*/
|
||||
export function computeNonMessageBreakdown(session: NonMessageTokenSource): {
|
||||
export function computeNonMessageBreakdown(
|
||||
session: NonMessageTokenSource,
|
||||
tokenizer: Tokenizer,
|
||||
): {
|
||||
skillsTokens: number;
|
||||
toolsTokens: number;
|
||||
systemContextTokens: number;
|
||||
systemPromptTokens: number;
|
||||
} {
|
||||
const entry = nonMessageTokenCacheEntry(session);
|
||||
const entry = nonMessageTokenCacheEntry(session, tokenizer);
|
||||
if (entry.breakdown) return entry.breakdown;
|
||||
const tools = session.agent?.state?.tools ?? EMPTY_TOOLS;
|
||||
const skillsTokens = estimateSkillsTokens(renderedSkills(session.skills ?? EMPTY_SKILLS, tools));
|
||||
const toolsTokens = estimateToolSchemaTokens(tools);
|
||||
const skillsTokens = estimateSkillsTokens(renderedSkills(session.skills ?? EMPTY_SKILLS, tools), tokenizer);
|
||||
const toolsTokens = estimateToolSchemaTokens(tools, tokenizer);
|
||||
const systemPromptParts = session.systemPrompt ?? EMPTY_STRING_PARTS;
|
||||
const systemContextTokens = countTokens(systemPromptParts.slice(1));
|
||||
const systemPromptTokens = Math.max(0, countTokens(systemPromptParts[0] ?? "") - skillsTokens);
|
||||
const systemContextTokens = tokenizer.countTokens(systemPromptParts.slice(1));
|
||||
const systemPromptTokens = Math.max(0, tokenizer.countTokens(systemPromptParts[0] ?? "") - skillsTokens);
|
||||
const breakdown = { skillsTokens, toolsTokens, systemContextTokens, systemPromptTokens };
|
||||
entry.breakdown = breakdown;
|
||||
return breakdown;
|
||||
@@ -207,6 +215,7 @@ export function computeContextBreakdown(
|
||||
options?: { snapcompactSavings?: boolean },
|
||||
): ContextBreakdown {
|
||||
const model = session.model;
|
||||
const tokenizer = session.agent.tokenizer;
|
||||
const contextWindow = model?.contextWindow ?? 0;
|
||||
|
||||
const breakdown = typeof session.getContextBreakdown === "function" ? session.getContextBreakdown() : undefined;
|
||||
@@ -226,13 +235,10 @@ export function computeContextBreakdown(
|
||||
systemPromptTokens = breakdown.systemPromptTokens;
|
||||
usedTokens = breakdown.usedTokens;
|
||||
} else {
|
||||
const convo = session.messages;
|
||||
if (convo) {
|
||||
for (const message of convo) {
|
||||
messagesTokens += estimateTokens(message);
|
||||
}
|
||||
}
|
||||
const nonMessage = computeNonMessageBreakdown(session);
|
||||
// Category split needs a messages-only number, so this walk stays local:
|
||||
// an anchored total folds the system prompt and tool schemas into it.
|
||||
messagesTokens = tokenizer.countMessages(session.messages ?? []);
|
||||
const nonMessage = computeNonMessageBreakdown(session, tokenizer);
|
||||
skillsTokens = nonMessage.skillsTokens;
|
||||
toolsTokens = nonMessage.toolsTokens;
|
||||
systemContextTokens = nonMessage.systemContextTokens;
|
||||
|
||||
@@ -49,7 +49,6 @@ import {
|
||||
type CompactionResult,
|
||||
calculatePromptTokens,
|
||||
collectEntriesForBranchSummary,
|
||||
estimateTokens,
|
||||
generateBranchSummary,
|
||||
type ShakeConfig,
|
||||
} from "@oh-my-pi/pi-agent-core/compaction";
|
||||
@@ -2368,7 +2367,8 @@ export class AgentSession {
|
||||
if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) {
|
||||
assistantMsg.contextSnapshot = {
|
||||
promptTokens: calculatePromptTokens(assistantMsg.usage),
|
||||
nonMessageTokens: this.#stats.pendingNonMessageTokens ?? computeNonMessageTokens(this),
|
||||
nonMessageTokens:
|
||||
this.#stats.pendingNonMessageTokens ?? computeNonMessageTokens(this, this.agent.tokenizer),
|
||||
compactionEpoch: this.#stats.compactionEpoch,
|
||||
};
|
||||
}
|
||||
@@ -5725,14 +5725,14 @@ export class AgentSession {
|
||||
}
|
||||
|
||||
const agentPromptOptions = options?.toolChoice ? { toolChoice: options.toolChoice } : undefined;
|
||||
const nonMessageTokens = computeNonMessageTokens(this);
|
||||
const nonMessageTokens = computeNonMessageTokens(this, this.agent.tokenizer);
|
||||
const contextWindow = this.model?.contextWindow ?? 0;
|
||||
const breakdown = this.getContextBreakdown({ contextWindow, pendingMessages: messages });
|
||||
const promptTokens =
|
||||
breakdown?.usedTokens ??
|
||||
nonMessageTokens +
|
||||
this.messages.reduce((sum, msg) => sum + estimateTokens(msg), 0) +
|
||||
messages.reduce((sum, msg) => sum + estimateTokens(msg), 0);
|
||||
this.agent.tokenizer.countMessages(this.messages) +
|
||||
this.agent.tokenizer.countMessages(messages);
|
||||
this.#stats.setPendingSnapshot({
|
||||
promptTokens,
|
||||
nonMessageTokens,
|
||||
|
||||
@@ -5,18 +5,17 @@ import {
|
||||
type AgentToolContext,
|
||||
AppendOnlyContextManager,
|
||||
type CompactionSummaryMessage,
|
||||
countTokens,
|
||||
resolveTelemetry,
|
||||
type StreamFn,
|
||||
ThinkingLevel,
|
||||
type Tokenizer,
|
||||
} from "@oh-my-pi/pi-agent-core";
|
||||
import {
|
||||
type CompactionResult,
|
||||
calculateContextTokens,
|
||||
compact,
|
||||
compactionContextTokens,
|
||||
createCompactionSummaryMessage,
|
||||
estimateTokens,
|
||||
estimateTranscriptTokens,
|
||||
NativeCompactionError,
|
||||
prepareCompaction,
|
||||
type SessionMessageEntry,
|
||||
@@ -892,8 +891,7 @@ export class SessionAdvisors {
|
||||
const runtime = new AdvisorRuntime(advisorAgentFacade, {
|
||||
snapshotMessages: () => this.#host.agent.state.messages,
|
||||
enqueueAdvice: (note, severity) => this.#routeAdvice(advisorRef, note, severity),
|
||||
maintainContext: (incomingTokens, signal) =>
|
||||
this.#maintainAdvisorContext(advisorRef, incomingTokens, signal),
|
||||
maintainContext: (incoming, signal) => this.#maintainAdvisorContext(advisorRef, incoming, signal),
|
||||
obfuscator: this.#host.obfuscator,
|
||||
getModelIdentity: () => formatModelString(advisorRef.agent.state.model),
|
||||
beginAdvisorUpdate: inProgress => {
|
||||
@@ -1318,11 +1316,12 @@ export class SessionAdvisors {
|
||||
|
||||
async #maintainAdvisorContext(
|
||||
advisor: ActiveAdvisor,
|
||||
incomingTokens: number,
|
||||
incoming: AgentMessage,
|
||||
signal: AbortSignal,
|
||||
): Promise<boolean> {
|
||||
await this.#maybeRestoreAdvisorRetryFallbackPrimary(advisor, signal);
|
||||
const agent = advisor.agent;
|
||||
const incomingTokens = agent.tokenizer.countMessage(incoming);
|
||||
|
||||
const compactionSettings = this.#host.settings.getGroup("compaction");
|
||||
if (compactionSettings.strategy === "off") return false;
|
||||
@@ -1333,20 +1332,16 @@ export class SessionAdvisors {
|
||||
if (contextWindow <= 0) return false;
|
||||
|
||||
const messages = agent.state.messages;
|
||||
const estimateOptions = { excludeEncryptedReasoning: true } as const;
|
||||
let storedConversationTokens = 0;
|
||||
for (const message of messages) {
|
||||
storedConversationTokens += estimateTokens(message, estimateOptions);
|
||||
}
|
||||
const storedConversationTokens = agent.tokenizer.countMessages(messages, { excludeEncryptedReasoning: true });
|
||||
// Provider usage (including cache reads and generated output) is the
|
||||
// trustworthy anchor for accumulated context. Add only the trailing incoming
|
||||
// delta to that arm. Floor it by a full local estimate — fixed advisor system
|
||||
// prompt, tool schemas, stored messages, and incoming delta — so provider
|
||||
// under-reporting or payload transforms cannot suppress maintenance.
|
||||
const providerContextTokens = this.#estimateAdvisorContextTokens(messages) + incomingTokens;
|
||||
const providerContextTokens = this.#estimateAdvisorContextTokens(messages, agent.tokenizer) + incomingTokens;
|
||||
const localContextTokens =
|
||||
countTokens(agent.state.systemPrompt) +
|
||||
estimateToolSchemaTokens(agent.state.tools) +
|
||||
agent.tokenizer.countTokens(agent.state.systemPrompt) +
|
||||
estimateToolSchemaTokens(agent.state.tools, agent.tokenizer) +
|
||||
storedConversationTokens +
|
||||
incomingTokens;
|
||||
const contextTokens = compactionContextTokens(providerContextTokens, localContextTokens);
|
||||
@@ -1406,7 +1401,7 @@ export class SessionAdvisors {
|
||||
this.#host.sessionId(),
|
||||
advisor.slug,
|
||||
);
|
||||
const preparation = prepareCompaction(pathEntries, compactionSettings, advisorModel);
|
||||
const preparation = prepareCompaction(pathEntries, compactionSettings, advisorModel, agent.tokenizer);
|
||||
if (!preparation) {
|
||||
// Cannot prepare compaction, fallback to re-prime
|
||||
return true;
|
||||
@@ -1758,7 +1753,7 @@ export class SessionAdvisors {
|
||||
#computeAdvisorStat(advisor: ActiveAdvisor): PerAdvisorStat {
|
||||
const model = advisor.agent.state.model;
|
||||
const messages = advisor.agent.state.messages;
|
||||
const contextTokens = this.#estimateAdvisorContextTokens(messages);
|
||||
const contextTokens = this.#estimateAdvisorContextTokens(messages, advisor.agent.tokenizer);
|
||||
let input = 0;
|
||||
let output = 0;
|
||||
let reasoning = 0;
|
||||
@@ -1848,7 +1843,7 @@ export class SessionAdvisors {
|
||||
* retained pre-compaction messages is stale and must not immediately retrigger
|
||||
* maintenance on the newly compacted context.
|
||||
*/
|
||||
#estimateAdvisorContextTokens(messages: AgentMessage[]): number {
|
||||
#estimateAdvisorContextTokens(messages: AgentMessage[], tokenizer: Tokenizer): number {
|
||||
let usageAnchorStartIndex = 0;
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
const message = messages[i];
|
||||
@@ -1860,33 +1855,10 @@ export class SessionAdvisors {
|
||||
usageAnchorStartIndex = advisorSummary.advisorUsageAnchorStartIndex ?? messages.length;
|
||||
break;
|
||||
}
|
||||
|
||||
let lastUsageIndex: number | undefined;
|
||||
let lastUsage: AssistantMessage["usage"] | undefined;
|
||||
for (let i = messages.length - 1; i >= usageAnchorStartIndex; i--) {
|
||||
const message = messages[i];
|
||||
if (message.role !== "assistant") continue;
|
||||
const assistant = message as AssistantMessage;
|
||||
if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error" && assistant.usage) {
|
||||
lastUsage = assistant.usage;
|
||||
lastUsageIndex = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const estimateOptions = { excludeEncryptedReasoning: true } as const;
|
||||
if (!lastUsage || lastUsageIndex === undefined) {
|
||||
let estimated = 0;
|
||||
for (const message of messages) {
|
||||
estimated += estimateTokens(message, estimateOptions);
|
||||
}
|
||||
return estimated;
|
||||
}
|
||||
let trailingTokens = 0;
|
||||
for (let i = lastUsageIndex + 1; i < messages.length; i++) {
|
||||
trailingTokens += estimateTokens(messages[i], estimateOptions);
|
||||
}
|
||||
return calculateContextTokens(lastUsage) + trailingTokens;
|
||||
return estimateTranscriptTokens(messages, tokenizer, {
|
||||
anchorFromIndex: usageAnchorStartIndex,
|
||||
excludeEncryptedReasoning: true,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -5,7 +5,6 @@ import {
|
||||
type Agent,
|
||||
type AgentMessage,
|
||||
type AgentTurnEndContext,
|
||||
countTokens,
|
||||
resolveTelemetry,
|
||||
type StreamFn,
|
||||
type ThinkingLevel,
|
||||
@@ -25,8 +24,7 @@ import {
|
||||
createCompactionSummaryMessage,
|
||||
DEFAULT_SHAKE_CONFIG,
|
||||
effectiveReserveTokens,
|
||||
estimateTokens,
|
||||
hasContextTokenUsage,
|
||||
isTranscriptUsageAnchor,
|
||||
NativeCompactionError,
|
||||
prepareCompaction,
|
||||
RESCUE_SHAKE_CONFIG,
|
||||
@@ -292,6 +290,10 @@ export class SessionMaintenance {
|
||||
return this.#host.model();
|
||||
}
|
||||
|
||||
get #tokenizer() {
|
||||
return this.#host.agent.tokenizer;
|
||||
}
|
||||
|
||||
get #goalModeState(): GoalModeState | undefined {
|
||||
return this.#host.goalModeState();
|
||||
}
|
||||
@@ -338,6 +340,7 @@ export class SessionMaintenance {
|
||||
const keepBoundaryId = getLatestCompactionEntry(branchEntries)?.firstKeptEntryId;
|
||||
const result = pruneToolOutputs(
|
||||
branchEntries,
|
||||
this.#tokenizer,
|
||||
this.#withPlanProtection({
|
||||
...DEFAULT_PRUNE_CONFIG,
|
||||
pruneUseless: this.#host.settings.getGroup("compaction").dropUseless,
|
||||
@@ -380,6 +383,7 @@ export class SessionMaintenance {
|
||||
const keepBoundaryId = getLatestCompactionEntry(branchEntries)?.firstKeptEntryId;
|
||||
const result = pruneSupersededToolResults(
|
||||
branchEntries,
|
||||
this.#tokenizer,
|
||||
this.#withPlanProtection({
|
||||
supersedeKey: supersedeReads ? readToolSupersedeKey : undefined,
|
||||
pruneUseless: dropUseless,
|
||||
@@ -482,7 +486,7 @@ export class SessionMaintenance {
|
||||
// the active model cannot replay still hides its prefix from the prompt.
|
||||
keepBoundaryId: latestCompaction?.firstKeptEntryId,
|
||||
});
|
||||
const regions = collectShakeRegions(branchEntries, config);
|
||||
const regions = collectShakeRegions(branchEntries, this.#tokenizer, config);
|
||||
if (regions.length === 0) {
|
||||
return { mode, toolResultsDropped: 0, blocksDropped: 0, tokensFreed: 0 };
|
||||
}
|
||||
@@ -495,18 +499,10 @@ export class SessionMaintenance {
|
||||
let anchorIndex = -1;
|
||||
for (let index = branchEntries.length - 1; index > compactionIndex; index--) {
|
||||
const entry = branchEntries[index];
|
||||
if (entry.type !== "message" || entry.message.role !== "assistant") continue;
|
||||
const assistant = entry.message;
|
||||
if (
|
||||
assistant.stopReason !== "aborted" &&
|
||||
assistant.stopReason !== "error" &&
|
||||
assistant.usage &&
|
||||
hasContextTokenUsage(assistant.usage)
|
||||
) {
|
||||
if (entry.type !== "message" || !isTranscriptUsageAnchor(entry.message)) continue;
|
||||
anchorIndex = index;
|
||||
break;
|
||||
}
|
||||
}
|
||||
const entryIndexes = new Map(branchEntries.map((entry, index) => [entry, index]));
|
||||
|
||||
let toolResultsDropped = 0;
|
||||
@@ -519,7 +515,7 @@ export class SessionMaintenance {
|
||||
else blocksDropped++;
|
||||
originalTokens += region.tokens;
|
||||
const replacement = replacements[index];
|
||||
const replacementTokenCount = replacement.length > 0 ? countTokens(replacement) : 0;
|
||||
const replacementTokenCount = replacement.length > 0 ? this.#tokenizer.countTokens(replacement) : 0;
|
||||
replacementTokens += replacementTokenCount;
|
||||
const entryIndex = entryIndexes.get(region.entry) ?? -1;
|
||||
if (
|
||||
@@ -638,7 +634,7 @@ export class SessionMaintenance {
|
||||
compactionCandidates = this.#getCompactionModelCandidates(availableModels);
|
||||
}
|
||||
const pathEntries = this.#host.sessionManager.getBranch();
|
||||
const preparation = prepareCompaction(pathEntries, effectiveSettings, this.#model);
|
||||
const preparation = prepareCompaction(pathEntries, effectiveSettings, this.#model, this.#tokenizer);
|
||||
if (!preparation) {
|
||||
// Check why we can't compact
|
||||
const lastEntry = pathEntries[pathEntries.length - 1];
|
||||
@@ -999,15 +995,17 @@ export class SessionMaintenance {
|
||||
* suppress it.
|
||||
*/
|
||||
#estimateStoredContextTokens(pendingMessages: AgentMessage[] = []): number {
|
||||
// Exclude encrypted reasoning (thinkingSignature / redactedThinking): its
|
||||
// local byte size diverges from what the provider bills, so counting it here
|
||||
// would let a thinking-heavy turn falsely trip the floor. The provider usage
|
||||
// (the other arm of compactionContextTokens) already accounts for it.
|
||||
// Local counting is the whole point of this arm: provider usage is
|
||||
// exactly what it must not trust. Exclude encrypted reasoning
|
||||
// (thinkingSignature / redactedThinking) too — its local byte size
|
||||
// diverges from what the provider bills, so counting it would let a
|
||||
// thinking-heavy turn falsely trip the floor. The provider usage (the
|
||||
// other arm of compactionContextTokens) already accounts for it.
|
||||
const opts = { excludeEncryptedReasoning: true } as const;
|
||||
return (
|
||||
computeNonMessageTokens(this.#host.nonMessageTokenSource()) +
|
||||
this.#host.messages().reduce((sum, msg) => sum + estimateTokens(msg, opts), 0) +
|
||||
pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg, opts), 0)
|
||||
computeNonMessageTokens(this.#host.nonMessageTokenSource(), this.#tokenizer) +
|
||||
this.#tokenizer.countMessages(this.#host.messages(), opts) +
|
||||
this.#tokenizer.countMessages(pendingMessages, opts)
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1032,7 +1030,8 @@ export class SessionMaintenance {
|
||||
if (!shouldCompact(contextTokens, contextWindow, compactionSettings)) return;
|
||||
if (
|
||||
pendingMidTurnDeadEnd &&
|
||||
prepareCompaction(this.#host.sessionManager.getBranch(), compactionSettings, model) === undefined
|
||||
prepareCompaction(this.#host.sessionManager.getBranch(), compactionSettings, model, this.#tokenizer) ===
|
||||
undefined
|
||||
) {
|
||||
// The prior tool loop already attempted the rescue and warned for this
|
||||
// persisted oversized turn. Only a later persisted cut point makes a
|
||||
@@ -1131,7 +1130,8 @@ export class SessionMaintenance {
|
||||
// soon as one appears.
|
||||
if (
|
||||
!model ||
|
||||
prepareCompaction(this.#host.sessionManager.getBranch(), compactionSettings, model) === undefined
|
||||
prepareCompaction(this.#host.sessionManager.getBranch(), compactionSettings, model, this.#tokenizer) ===
|
||||
undefined
|
||||
) {
|
||||
return;
|
||||
}
|
||||
@@ -1717,16 +1717,14 @@ export class SessionMaintenance {
|
||||
const ctxWindow = this.#model?.contextWindow ?? 0;
|
||||
if (ctxWindow <= 0) return Math.min(snapcompact.MAX_FRAMES_DEFAULT, snapcompact.maxFramesForDataBudget());
|
||||
const reserve = effectiveReserveTokens(ctxWindow, settings);
|
||||
let baseTokens = computeNonMessageTokens(this.#host.nonMessageTokenSource());
|
||||
for (const message of preparation.recentMessages) {
|
||||
baseTokens += estimateTokens(message);
|
||||
}
|
||||
let baseTokens = computeNonMessageTokens(this.#host.nonMessageTokenSource(), this.#tokenizer);
|
||||
baseTokens += this.#tokenizer.countMessages(preparation.recentMessages);
|
||||
const totalBudget = ctxWindow - reserve;
|
||||
// Skip iff there is no headroom whatsoever; a text-only archive costs
|
||||
// far less than the cap reserve below, so any positive residual is
|
||||
// worth attempting and the projection guard catches actual overflow.
|
||||
if (baseTokens >= totalBudget) return 0;
|
||||
// Cap reserve mirrors what `estimateTokens(summaryMessage)` will charge
|
||||
// Cap reserve mirrors what `countMessage(summaryMessage)` will charge
|
||||
// when frames > 0: `countTokens(summaryTemplate ‖ textHead ‖ textTail)`
|
||||
// plus `numFrames × FRAME_TOKEN_ESTIMATE`. Resolve the shape this
|
||||
// snapcompact pass will actually use (matches the `shape` argument
|
||||
@@ -1788,10 +1786,10 @@ export class SessionMaintenance {
|
||||
undefined,
|
||||
blocks,
|
||||
);
|
||||
let tokens = computeNonMessageTokens(this.#host.nonMessageTokenSource()) + estimateTokens(summaryMessage);
|
||||
for (const message of preparation.recentMessages) {
|
||||
tokens += estimateTokens(message);
|
||||
}
|
||||
let tokens =
|
||||
computeNonMessageTokens(this.#host.nonMessageTokenSource(), this.#tokenizer) +
|
||||
this.#tokenizer.countMessage(summaryMessage);
|
||||
tokens += this.#tokenizer.countMessages(preparation.recentMessages);
|
||||
return tokens;
|
||||
}
|
||||
|
||||
@@ -1860,9 +1858,9 @@ export class SessionMaintenance {
|
||||
if (contextWindow <= 0) return true;
|
||||
const activeExcludedMessage =
|
||||
excludedMessage && this.#host.messages().includes(excludedMessage) ? excludedMessage : undefined;
|
||||
const providerExcludedTokens = activeExcludedMessage ? estimateTokens(activeExcludedMessage) : 0;
|
||||
const providerExcludedTokens = activeExcludedMessage ? this.#tokenizer.countMessage(activeExcludedMessage) : 0;
|
||||
const storedExcludedTokens = activeExcludedMessage
|
||||
? estimateTokens(activeExcludedMessage, { excludeEncryptedReasoning: true })
|
||||
? this.#tokenizer.countMessage(activeExcludedMessage, { excludeEncryptedReasoning: true })
|
||||
: 0;
|
||||
const compactionSettings = this.#host.settings.getGroup("compaction");
|
||||
const residualTokens = compactionContextTokens(
|
||||
@@ -1998,7 +1996,7 @@ export class SessionMaintenance {
|
||||
if (ctxWindow <= 0) return Math.min(snapcompact.MAX_FRAMES_DEFAULT, snapcompact.maxFramesForDataBudget());
|
||||
const thresholdTokens = resolveThresholdTokens(ctxWindow, settings);
|
||||
const recoveryBandTokens = Math.floor(thresholdTokens * COMPACTION_RECOVERY_BAND);
|
||||
const baseTokens = computeNonMessageTokens(this.#host.nonMessageTokenSource());
|
||||
const baseTokens = computeNonMessageTokens(this.#host.nonMessageTokenSource(), this.#tokenizer);
|
||||
const shape = snapcompact.resolveShape(this.#model, this.#host.settings.get("snapcompact.shape"));
|
||||
const edgeCap = snapcompact.geometry(shape).capacity;
|
||||
const textEdgeTokens = Math.ceil((2 * edgeCap * 1.15) / 4);
|
||||
@@ -2069,7 +2067,7 @@ export class SessionMaintenance {
|
||||
}
|
||||
if (!inKeptRegion) continue;
|
||||
const message = (entry as { message?: AgentMessage }).message;
|
||||
if (message) keptTailTokens += estimateTokens(message);
|
||||
if (message) keptTailTokens += this.#tokenizer.countMessage(message);
|
||||
}
|
||||
const archive = snapcompact.getPreservedArchive(staleEntry.preserveData);
|
||||
if (!archive || archive.frames.length <= 1) return undefined;
|
||||
@@ -2359,7 +2357,12 @@ export class SessionMaintenance {
|
||||
const pathEntries = this.#host.sessionManager.getBranch();
|
||||
|
||||
let pathEntriesForCompaction = pathEntries;
|
||||
let preparation = prepareCompaction(pathEntriesForCompaction, compactionSettings, this.#model);
|
||||
let preparation = prepareCompaction(
|
||||
pathEntriesForCompaction,
|
||||
compactionSettings,
|
||||
this.#model,
|
||||
this.#tokenizer,
|
||||
);
|
||||
if (!preparation) {
|
||||
// prepareCompaction found nothing to summarize because the kept region
|
||||
// is a single oversized recent turn — findCutPoint never cuts inside a
|
||||
@@ -2409,7 +2412,12 @@ export class SessionMaintenance {
|
||||
// branch has been rewritten either way.
|
||||
rescueRewroteHistory = true;
|
||||
pathEntriesForCompaction = this.#host.sessionManager.getBranch();
|
||||
preparation = prepareCompaction(pathEntriesForCompaction, compactionSettings, this.#model);
|
||||
preparation = prepareCompaction(
|
||||
pathEntriesForCompaction,
|
||||
compactionSettings,
|
||||
this.#model,
|
||||
this.#tokenizer,
|
||||
);
|
||||
return preparation !== undefined;
|
||||
},
|
||||
});
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
import type { Agent, AgentMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import {
|
||||
calculatePromptTokens,
|
||||
estimateTokens,
|
||||
hasContextTokenUsage,
|
||||
findTranscriptUsageAnchor,
|
||||
isTranscriptUsageAnchor,
|
||||
type SessionMessageEntry,
|
||||
} from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import type { AssistantMessage, Model, ProviderResponseMetadata, Usage } from "@oh-my-pi/pi-ai";
|
||||
@@ -56,6 +56,30 @@ export class SessionStatsTracker {
|
||||
this.#host = host;
|
||||
}
|
||||
|
||||
get #tokenizer() {
|
||||
return this.#host.agent.tokenizer;
|
||||
}
|
||||
|
||||
/**
|
||||
* Anchored used-token arithmetic shared by every anchored branch: provider
|
||||
* base + non-message growth since the anchor + local tail + pending.
|
||||
*/
|
||||
#anchoredUsedTokens(
|
||||
base: number,
|
||||
anchorNonMessageTokens: number,
|
||||
currentNonMessageTokens: number,
|
||||
tailFromIndex: number,
|
||||
activeMessages: readonly AgentMessage[],
|
||||
pendingTokens: number,
|
||||
): number {
|
||||
return (
|
||||
base +
|
||||
Math.max(0, currentNonMessageTokens - anchorNonMessageTokens) +
|
||||
this.#tokenizer.countMessages(activeMessages.slice(tailFromIndex)) +
|
||||
pendingTokens
|
||||
);
|
||||
}
|
||||
|
||||
/** Returns aggregate message, token, and cost statistics for the session. */
|
||||
getSessionStats(): SessionStats {
|
||||
const state = this.#host.agent.state;
|
||||
@@ -128,32 +152,26 @@ export class SessionStatsTracker {
|
||||
const contextWindow = Number.isFinite(rawContextWindow) && rawContextWindow > 0 ? rawContextWindow : 0;
|
||||
const { skillsTokens, toolsTokens, systemContextTokens, systemPromptTokens } = computeNonMessageBreakdown(
|
||||
this.#host.session,
|
||||
this.#tokenizer,
|
||||
);
|
||||
const categoryNonMessageTokens = skillsTokens + toolsTokens + systemContextTokens + systemPromptTokens;
|
||||
const currentNonMessageTokens = computeNonMessageTokens(this.#host.session);
|
||||
const currentNonMessageTokens = computeNonMessageTokens(this.#host.session, this.#tokenizer);
|
||||
const branchEntries = this.#host.sessionManager.getBranch();
|
||||
const latestCompaction = getLatestCompactionEntry(branchEntries);
|
||||
const compactionIndex = latestCompaction ? branchEntries.lastIndexOf(latestCompaction) : -1;
|
||||
let usedTokens = 0;
|
||||
let anchored = false;
|
||||
const pendingMessages = options?.pendingMessages ?? [];
|
||||
const pendingTokens = this.#tokenizer.countMessages(pendingMessages);
|
||||
const pending = this.#pendingContextSnapshot;
|
||||
|
||||
let anchorEntry: SessionMessageEntry | undefined;
|
||||
for (let index = branchEntries.length - 1; index > compactionIndex; index--) {
|
||||
const entry = branchEntries[index];
|
||||
if (entry.type !== "message" || entry.message.role !== "assistant") continue;
|
||||
const assistant = entry.message;
|
||||
if (
|
||||
assistant.stopReason !== "aborted" &&
|
||||
assistant.stopReason !== "error" &&
|
||||
assistant.usage &&
|
||||
hasContextTokenUsage(assistant.usage)
|
||||
) {
|
||||
if (entry.type !== "message" || !isTranscriptUsageAnchor(entry.message)) continue;
|
||||
anchorEntry = entry;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
const activeMessages = this.#host.agent.state.messages;
|
||||
let anchorIndex = -1;
|
||||
@@ -175,67 +193,49 @@ export class SessionStatsTracker {
|
||||
anchorIndex !== -1 &&
|
||||
(!pending || (anchorIndex >= pending.cutoffCount && anchorEpoch >= pending.epoch));
|
||||
if (useAnchor && anchorAssistant) {
|
||||
const promptTokens = correctedPromptTokens(anchorAssistant);
|
||||
const nonMessageTokens =
|
||||
anchorAssistant.contextSnapshot?.nonMessageTokens ?? computeNonMessageTokens(this.#host.session);
|
||||
anchorAssistant.contextSnapshot?.nonMessageTokens ??
|
||||
computeNonMessageTokens(this.#host.session, this.#tokenizer);
|
||||
anchored = true;
|
||||
let tailTokens = 0;
|
||||
for (let index = anchorIndex + 1; index < activeMessages.length; index++) {
|
||||
tailTokens += estimateTokens(activeMessages[index]);
|
||||
}
|
||||
usedTokens =
|
||||
promptTokens +
|
||||
Math.max(0, currentNonMessageTokens - nonMessageTokens) +
|
||||
tailTokens +
|
||||
pendingMessages.reduce((sum, message) => sum + estimateTokens(message), 0);
|
||||
usedTokens = this.#anchoredUsedTokens(
|
||||
correctedPromptTokens(anchorAssistant),
|
||||
nonMessageTokens,
|
||||
currentNonMessageTokens,
|
||||
anchorIndex + 1,
|
||||
activeMessages,
|
||||
pendingTokens,
|
||||
);
|
||||
} else if (pending) {
|
||||
anchored = true;
|
||||
let tailTokens = 0;
|
||||
for (let index = pending.cutoffCount; index < activeMessages.length; index++) {
|
||||
tailTokens += estimateTokens(activeMessages[index]);
|
||||
}
|
||||
usedTokens =
|
||||
pending.promptTokens +
|
||||
Math.max(0, currentNonMessageTokens - pending.nonMessageTokens) +
|
||||
tailTokens +
|
||||
pendingMessages.reduce((sum, message) => sum + estimateTokens(message), 0);
|
||||
usedTokens = this.#anchoredUsedTokens(
|
||||
pending.promptTokens,
|
||||
pending.nonMessageTokens,
|
||||
currentNonMessageTokens,
|
||||
pending.cutoffCount,
|
||||
activeMessages,
|
||||
pendingTokens,
|
||||
);
|
||||
}
|
||||
|
||||
if (!anchored && !pending && branchEntries.length === 0) {
|
||||
for (let index = activeMessages.length - 1; index >= 0; index--) {
|
||||
const message = activeMessages[index];
|
||||
if (
|
||||
message.role !== "assistant" ||
|
||||
message.stopReason === "aborted" ||
|
||||
message.stopReason === "error" ||
|
||||
!message.usage ||
|
||||
!hasContextTokenUsage(message.usage)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
const promptTokens = correctedPromptTokens(message);
|
||||
const liveAnchor = findTranscriptUsageAnchor(activeMessages);
|
||||
if (liveAnchor) {
|
||||
const nonMessageTokens =
|
||||
message.contextSnapshot?.nonMessageTokens ?? computeNonMessageTokens(this.#host.session);
|
||||
let tailTokens = 0;
|
||||
for (let tailIndex = index + 1; tailIndex < activeMessages.length; tailIndex++) {
|
||||
tailTokens += estimateTokens(activeMessages[tailIndex]);
|
||||
}
|
||||
usedTokens =
|
||||
promptTokens +
|
||||
Math.max(0, currentNonMessageTokens - nonMessageTokens) +
|
||||
tailTokens +
|
||||
pendingMessages.reduce((sum, pendingMessage) => sum + estimateTokens(pendingMessage), 0);
|
||||
liveAnchor.message.contextSnapshot?.nonMessageTokens ??
|
||||
computeNonMessageTokens(this.#host.session, this.#tokenizer);
|
||||
usedTokens = this.#anchoredUsedTokens(
|
||||
correctedPromptTokens(liveAnchor.message),
|
||||
nonMessageTokens,
|
||||
currentNonMessageTokens,
|
||||
liveAnchor.index + 1,
|
||||
activeMessages,
|
||||
pendingTokens,
|
||||
);
|
||||
anchored = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!anchored) {
|
||||
let messagesTokens = 0;
|
||||
for (const message of activeMessages) messagesTokens += estimateTokens(message);
|
||||
usedTokens =
|
||||
currentNonMessageTokens +
|
||||
messagesTokens +
|
||||
pendingMessages.reduce((sum, message) => sum + estimateTokens(message), 0);
|
||||
usedTokens = currentNonMessageTokens + this.#tokenizer.countMessages(activeMessages) + pendingTokens;
|
||||
}
|
||||
return {
|
||||
contextWindow,
|
||||
@@ -296,21 +296,13 @@ export class SessionStatsTracker {
|
||||
const compactionIndex = latestCompaction ? branchEntries.lastIndexOf(latestCompaction) : -1;
|
||||
for (let index = branchEntries.length - 1; index > compactionIndex; index--) {
|
||||
const entry = branchEntries[index];
|
||||
if (entry.type !== "message" || entry.message.role !== "assistant") continue;
|
||||
if (entry.type !== "message" || !isTranscriptUsageAnchor(entry.message)) continue;
|
||||
const assistant = entry.message;
|
||||
if (
|
||||
assistant.stopReason === "aborted" ||
|
||||
assistant.stopReason === "error" ||
|
||||
!assistant.usage ||
|
||||
!hasContextTokenUsage(assistant.usage)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (!assistant.contextSnapshot) {
|
||||
assistant.contextSnapshot = {
|
||||
promptTokens: calculatePromptTokens(assistant.usage),
|
||||
nonMessageTokens: computeNonMessageTokens(this.#host.session),
|
||||
nonMessageTokens: computeNonMessageTokens(this.#host.session, this.#tokenizer),
|
||||
compactionEpoch: this.#compactionEpoch,
|
||||
};
|
||||
}
|
||||
@@ -331,10 +323,10 @@ export class SessionStatsTracker {
|
||||
rebaseAfterCompaction(): void {
|
||||
this.#compactionEpoch++;
|
||||
if (!this.#pendingContextSnapshot) return;
|
||||
const nonMessageTokens = computeNonMessageTokens(this.#host.session);
|
||||
const nonMessageTokens = computeNonMessageTokens(this.#host.session, this.#tokenizer);
|
||||
const messages = this.#host.agent.state.messages;
|
||||
this.setPendingSnapshot({
|
||||
promptTokens: nonMessageTokens + messages.reduce((sum, message) => sum + estimateTokens(message), 0),
|
||||
promptTokens: nonMessageTokens + this.#tokenizer.countMessages(messages),
|
||||
nonMessageTokens,
|
||||
cutoffCount: messages.length,
|
||||
});
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
* estimate (`estimateInlineSavings`) so the two can never disagree.
|
||||
*/
|
||||
|
||||
import { countTokens } from "@oh-my-pi/pi-agent-core";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { Context, ImageContent, Model, TextContent, ToolResultMessage, UserMessage } from "@oh-my-pi/pi-ai";
|
||||
import * as snapcompact from "@oh-my-pi/snapcompact";
|
||||
import contextFramesNote from "../prompts/system/snapcompact-context-frames-note.md" with { type: "text" };
|
||||
@@ -282,6 +282,7 @@ export function estimateInlineSavings(input: {
|
||||
}
|
||||
|
||||
const shape = snapcompact.resolveShape(model, options.shape);
|
||||
const tokenizer = new Tokenizer(model.id);
|
||||
let existingImages = 0;
|
||||
for (const message of input.messages) {
|
||||
if (!Array.isArray(message.content)) continue;
|
||||
@@ -304,7 +305,7 @@ export function estimateInlineSavings(input: {
|
||||
.filter(block => block.type === "text" && typeof block.text === "string")
|
||||
.map(block => block.text as string)
|
||||
.join("\n");
|
||||
const textTokens = text.length > 0 ? countTokens(text) : 0;
|
||||
const textTokens = text.length > 0 ? tokenizer.countTokens(text) : 0;
|
||||
candidates.push({
|
||||
id: message.toolCallId,
|
||||
textTokens,
|
||||
@@ -320,7 +321,7 @@ export function estimateInlineSavings(input: {
|
||||
systemPromptTarget = selectSystemPromptImageTarget(input.systemPrompt, options.renderSystemPrompt);
|
||||
if (systemPromptTarget) {
|
||||
systemPromptCandidate = {
|
||||
textTokens: countTokens(systemPromptTarget.text),
|
||||
textTokens: tokenizer.countTokens(systemPromptTarget.text),
|
||||
frames: snapcompact.frames(systemPromptTarget.text, { shape }),
|
||||
};
|
||||
}
|
||||
@@ -420,6 +421,7 @@ export class SnapcompactInlineTransformer {
|
||||
if (!model.input.includes("image")) return context;
|
||||
|
||||
const shape = snapcompact.resolveShape(model, this.options.shape);
|
||||
const tokenizer = new Tokenizer(model.id);
|
||||
const budget = snapcompact.providerImageBudget(model.provider) - countContextImages(context);
|
||||
if (budget <= 0) return context;
|
||||
|
||||
@@ -444,7 +446,7 @@ export class SnapcompactInlineTransformer {
|
||||
.filter(isTextContent)
|
||||
.map(block => block.text)
|
||||
.join("\n");
|
||||
const textTokens = text.length > 0 ? countTokens(text) : 0;
|
||||
const textTokens = text.length > 0 ? tokenizer.countTokens(text) : 0;
|
||||
candidates.push({
|
||||
id: message.toolCallId,
|
||||
textTokens,
|
||||
@@ -461,7 +463,7 @@ export class SnapcompactInlineTransformer {
|
||||
systemPromptTarget = selectSystemPromptImageTarget(context.systemPrompt, this.options.renderSystemPrompt);
|
||||
if (systemPromptTarget) {
|
||||
systemPromptCandidate = {
|
||||
textTokens: countTokens(systemPromptTarget.text),
|
||||
textTokens: tokenizer.countTokens(systemPromptTarget.text),
|
||||
frames: snapcompact.frames(systemPromptTarget.text, { shape }),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ import { describe, expect, it, spyOn } from "bun:test";
|
||||
import * as fs from "node:fs/promises";
|
||||
import * as os from "node:os";
|
||||
import * as path from "node:path";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type {
|
||||
ResetCreditAccountStatus,
|
||||
ResetCreditRedeemOutcome,
|
||||
@@ -1149,7 +1150,7 @@ describe("wave 5 — adapters and polish", () => {
|
||||
contextWindow: 200_000,
|
||||
};
|
||||
(session as unknown as Record<string, unknown>).skills = [];
|
||||
(session as unknown as Record<string, unknown>).agent = { state: { tools: [] } };
|
||||
(session as unknown as Record<string, unknown>).agent = { state: { tools: [] }, tokenizer: new Tokenizer() };
|
||||
(session as unknown as Record<string, unknown>).systemPrompt = ["You are a helpful assistant."];
|
||||
session.messages = [
|
||||
{ role: "user", content: "Hello, how are you?" },
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "bun:test";
|
||||
import { Agent, type AgentMessage, type CompactionSummaryMessage, countTokens } from "@oh-my-pi/pi-agent-core";
|
||||
import { Agent, type AgentMessage, type CompactionSummaryMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import * as compactionModule from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import { calculateContextTokens, estimateTokens, resolveThresholdTokens } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import { calculateContextTokens, resolveThresholdTokens } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import type { AssistantMessage } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel, type MockModel, registerMockApi } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
|
||||
@@ -198,7 +198,8 @@ describe("AgentSession advisor context maintenance", () => {
|
||||
const update = advisorCall.context.messages.find(message => message.role === "user");
|
||||
if (!update) throw new Error("Expected the advisor's incremental update");
|
||||
const threshold = resolveThresholdTokens(CONTEXT_WINDOW, settings.getGroup("compaction"));
|
||||
const providerAndUpdateTokens = calculateContextTokens(anchor.usage) + estimateTokens(update as AgentMessage);
|
||||
const providerAndUpdateTokens =
|
||||
calculateContextTokens(anchor.usage) + advisor.tokenizer.countMessage(update as AgentMessage);
|
||||
expect(calculateContextTokens(anchor.usage)).toBe(CACHE_READ_TOKENS + INPUT_TOKENS + OUTPUT_TOKENS);
|
||||
expect(providerAndUpdateTokens).toBeGreaterThan(threshold);
|
||||
|
||||
@@ -249,8 +250,10 @@ describe("AgentSession advisor context maintenance", () => {
|
||||
const { advisor, advisorMock, settings } = createHarness();
|
||||
const seed: AgentMessage = { role: "user", content: "small stored advisor message", timestamp: 1 };
|
||||
advisor.state.messages.push(seed);
|
||||
const storedTokens = estimateTokens(seed, { excludeEncryptedReasoning: true });
|
||||
const fixedPrefixTokens = countTokens(advisor.state.systemPrompt) + estimateToolSchemaTokens(advisor.state.tools);
|
||||
const storedTokens = advisor.tokenizer.countMessage(seed, { excludeEncryptedReasoning: true });
|
||||
const fixedPrefixTokens =
|
||||
advisor.tokenizer.countTokens(advisor.state.systemPrompt) +
|
||||
estimateToolSchemaTokens(advisor.state.tools, advisor.tokenizer);
|
||||
const threshold = storedTokens + Math.floor(fixedPrefixTokens / 2);
|
||||
settings.set("compaction.thresholdTokens", threshold);
|
||||
|
||||
@@ -259,7 +262,7 @@ describe("AgentSession advisor context maintenance", () => {
|
||||
const advisorCall = advisorMock.calls[0];
|
||||
const update = advisorCall.context.messages.find(message => message.role === "user");
|
||||
if (!update) throw new Error("Expected the advisor's incremental update");
|
||||
const messagesOnlyTokens = storedTokens + estimateTokens(update as AgentMessage);
|
||||
const messagesOnlyTokens = storedTokens + advisor.tokenizer.countMessage(update as AgentMessage);
|
||||
expect(messagesOnlyTokens).toBeLessThan(threshold);
|
||||
expect(messagesOnlyTokens + fixedPrefixTokens).toBeGreaterThan(threshold);
|
||||
expect(JSON.stringify(advisor.state.messages)).not.toContain("small stored advisor message");
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import { describe, expect, it, vi } from "bun:test";
|
||||
import { type } from "@oh-my-pi/omptype";
|
||||
import type { AgentMessage, AgentTelemetryConfig } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AgentMessage, type AgentTelemetryConfig, Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { AssistantMessage } from "@oh-my-pi/pi-ai";
|
||||
import * as AIError from "@oh-my-pi/pi-ai/error";
|
||||
import { kCursorExecResolved } from "@oh-my-pi/pi-ai/utils/block-symbols";
|
||||
@@ -2436,8 +2436,11 @@ describe("advisor", () => {
|
||||
const host: AdvisorRuntimeHost = {
|
||||
snapshotMessages: () => messages,
|
||||
enqueueAdvice: () => {},
|
||||
maintainContext: async tokens => {
|
||||
expect(tokens).toBeGreaterThan(0);
|
||||
maintainContext: async incoming => {
|
||||
// The host receives the pending update itself and sizes it with its
|
||||
// own model's tokenizer, so it must arrive as a non-empty message.
|
||||
expect(incoming.role).toBe("user");
|
||||
expect(promptText([incoming]).length).toBeGreaterThan(0);
|
||||
return shouldResetContext;
|
||||
},
|
||||
};
|
||||
@@ -2674,8 +2677,8 @@ describe("advisor", () => {
|
||||
{
|
||||
snapshotMessages: () => messages,
|
||||
enqueueAdvice: () => {},
|
||||
maintainContext: async incomingTokens => {
|
||||
maintenanceTokens.push(incomingTokens);
|
||||
maintainContext: async incoming => {
|
||||
maintenanceTokens.push(new Tokenizer().countMessage(incoming));
|
||||
if (maintenanceTokens.length === 4) fourthMaintenance.resolve();
|
||||
return false;
|
||||
},
|
||||
@@ -5126,7 +5129,7 @@ describe("advisor", () => {
|
||||
const host: AdvisorRuntimeHost = {
|
||||
snapshotMessages: () => [],
|
||||
enqueueAdvice: () => {},
|
||||
maintainContext: async (_incomingTokens, signal) => {
|
||||
maintainContext: async (_incoming, signal) => {
|
||||
maintenanceSignals.push(signal);
|
||||
return false;
|
||||
},
|
||||
|
||||
@@ -20,7 +20,7 @@
|
||||
|
||||
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test";
|
||||
import { Agent } from "@oh-my-pi/pi-agent-core";
|
||||
import { effectiveReserveTokens, estimateTokens, prepareCompaction } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import { effectiveReserveTokens, prepareCompaction } from "@oh-my-pi/pi-agent-core/compaction";
|
||||
import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
@@ -170,10 +170,8 @@ describe("AgentSession snapcompact frame-budget sizing", () => {
|
||||
// numFrames × FRAME_TOKEN_ESTIMATE + non-message + kept-recent.
|
||||
const preparation = prepareCompaction(branchEntries, settings);
|
||||
if (!preparation) throw new Error("Expected non-empty preparation");
|
||||
let baseTokens = computeNonMessageTokens(session);
|
||||
for (const message of preparation.recentMessages) {
|
||||
baseTokens += estimateTokens(message);
|
||||
}
|
||||
let baseTokens = computeNonMessageTokens(session, session.agent.tokenizer);
|
||||
baseTokens += session.agent.tokenizer.countMessages(preparation.recentMessages);
|
||||
const shape = snapcompact.resolveShape(model);
|
||||
const edgeCap = snapcompact.geometry(shape).capacity;
|
||||
// Worst-case `textHead + textTail` tokenized at the cl100k 4-chars/token
|
||||
|
||||
@@ -1,13 +1,12 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test";
|
||||
import * as path from "node:path";
|
||||
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AgentMessage, Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import {
|
||||
type CompactionSettings,
|
||||
calculateContextTokens,
|
||||
compact,
|
||||
compactionContextTokens,
|
||||
DEFAULT_COMPACTION_SETTINGS,
|
||||
estimateTokens,
|
||||
findCutPoint,
|
||||
getLastAssistantUsage,
|
||||
hasContextTokenUsage,
|
||||
@@ -32,6 +31,8 @@ import { migrateSessionEntries } from "@oh-my-pi/pi-coding-agent/session/session
|
||||
import { mockFetch } from "./helpers/fetch-mock";
|
||||
import { e2eApiKey } from "./utilities";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
// ============================================================================
|
||||
// Test fixtures
|
||||
// ============================================================================
|
||||
@@ -372,7 +373,7 @@ describe("compactionContextTokens", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe("estimateTokens excludeEncryptedReasoning (compaction floor)", () => {
|
||||
describe("Tokenizer.countMessage excludeEncryptedReasoning (compaction floor)", () => {
|
||||
it("drops encrypted reasoning from the floor estimate but counts it by default", () => {
|
||||
const blob = "blob ".repeat(8_000); // large opaque encrypted-reasoning payload
|
||||
const msg: AssistantMessage = {
|
||||
@@ -388,8 +389,8 @@ describe("estimateTokens excludeEncryptedReasoning (compaction floor)", () => {
|
||||
provider: "openai",
|
||||
model: "gpt-5.5",
|
||||
};
|
||||
const withBlob = estimateTokens(msg);
|
||||
const flooredEstimate = estimateTokens(msg, { excludeEncryptedReasoning: true });
|
||||
const withBlob = tokenizer.countMessage(msg);
|
||||
const flooredEstimate = tokenizer.countMessage(msg, { excludeEncryptedReasoning: true });
|
||||
// Default counts the blob (providers bill it on replay); the floor excludes it,
|
||||
// so a thinking-heavy turn can't falsely trip compaction on local byte size.
|
||||
expect(withBlob).toBeGreaterThan(flooredEstimate + 1_000);
|
||||
@@ -408,7 +409,7 @@ describe("estimateTokens excludeEncryptedReasoning (compaction floor)", () => {
|
||||
// Even with the floor option, tool-result content is fully counted — that is
|
||||
// exactly what a before_provider_request compressor (e.g. Headroom) shrinks,
|
||||
// so the floor must still see its real size.
|
||||
expect(estimateTokens(toolMsg, { excludeEncryptedReasoning: true })).toBeGreaterThan(1_000);
|
||||
expect(tokenizer.countMessage(toolMsg, { excludeEncryptedReasoning: true })).toBeGreaterThan(1_000);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1134,7 +1135,7 @@ describe("findCutPoint", () => {
|
||||
|
||||
// 20 entries, last assistant has 10000 tokens
|
||||
// keepRecentTokens = 2500: keep entries where diff < 2500
|
||||
const result = findCutPoint(entries, 0, entries.length, 2500);
|
||||
const result = findCutPoint(entries, tokenizer, 0, entries.length, 2500);
|
||||
|
||||
// Should cut at a valid cut point (user or assistant message)
|
||||
expect(entries[result.firstKeptEntryIndex].type).toBe("message");
|
||||
@@ -1144,7 +1145,7 @@ describe("findCutPoint", () => {
|
||||
|
||||
it("should return startIndex if no valid cut points in range", () => {
|
||||
const entries: SessionEntry[] = [createMessageEntry(createAssistantMessage("a"))];
|
||||
const result = findCutPoint(entries, 0, entries.length, 1000);
|
||||
const result = findCutPoint(entries, tokenizer, 0, entries.length, 1000);
|
||||
expect(result.firstKeptEntryIndex).toBe(0);
|
||||
});
|
||||
|
||||
@@ -1156,7 +1157,7 @@ describe("findCutPoint", () => {
|
||||
createMessageEntry(createAssistantMessage("b", createMockUsage(0, 50, 1000, 0))),
|
||||
];
|
||||
|
||||
const result = findCutPoint(entries, 0, entries.length, 50000);
|
||||
const result = findCutPoint(entries, tokenizer, 0, entries.length, 50000);
|
||||
expect(result.firstKeptEntryIndex).toBe(0);
|
||||
});
|
||||
|
||||
@@ -1172,7 +1173,7 @@ describe("findCutPoint", () => {
|
||||
];
|
||||
|
||||
// With keepRecentTokens = 3000, should cut somewhere in Turn 2
|
||||
const result = findCutPoint(entries, 0, entries.length, 3000);
|
||||
const result = findCutPoint(entries, tokenizer, 0, entries.length, 3000);
|
||||
|
||||
// If cut at assistant message (not user), should indicate split turn
|
||||
const cutEntry = entries[result.firstKeptEntryIndex] as SessionMessageEntry;
|
||||
@@ -1385,7 +1386,7 @@ describe("buildSessionContext", () => {
|
||||
describe("Large session fixture", () => {
|
||||
it("should find cut point in large session", async () => {
|
||||
const entries = await loadLargeSessionEntries();
|
||||
const result = findCutPoint(entries, 0, entries.length, DEFAULT_COMPACTION_SETTINGS.keepRecentTokens);
|
||||
const result = findCutPoint(entries, tokenizer, 0, entries.length, DEFAULT_COMPACTION_SETTINGS.keepRecentTokens);
|
||||
|
||||
// Cut point should be at a message entry (user or assistant)
|
||||
expect(entries[result.firstKeptEntryIndex].type).toBe("message");
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
import { afterAll, beforeAll, describe, expect, it, vi } from "bun:test";
|
||||
import * as path from "node:path";
|
||||
import { Agent, type AgentMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import { estimateTokens } from "@oh-my-pi/pi-agent-core/compaction/compaction";
|
||||
import type { AssistantMessage, Message, Model } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
@@ -301,7 +300,7 @@ describe("Context usage consolidation", () => {
|
||||
const breakdown = session.getContextBreakdown();
|
||||
expect(breakdown?.anchored).toBe(true);
|
||||
|
||||
const customEstimate = estimateTokens(customMsg);
|
||||
const customEstimate = agent.tokenizer.countMessage(customMsg);
|
||||
expect(breakdown?.usedTokens).toBe(150 + customEstimate);
|
||||
|
||||
await tempDir.remove();
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { describe, expect, it } from "bun:test";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import {
|
||||
compact,
|
||||
estimateTokens,
|
||||
@@ -8,15 +9,20 @@ import {
|
||||
// Issue #6583: pi extensions import `estimateTokens` from
|
||||
// `@earendil-works/pi-coding-agent`, which aliases to this shim. Legacy pi
|
||||
// re-exported it from the coding-agent package root (via
|
||||
// `./core/compaction/index.ts`); in omp it lives in
|
||||
// `@oh-my-pi/pi-agent-core/compaction` and the coding-agent barrel does not
|
||||
// forward it, so `export * from "../index"` left the symbol off the shim
|
||||
// surface and a named import threw Bun's static "Export named X not found"
|
||||
// during plugin validation (e.g. `omp plugin install pi-blackhole`). This pins
|
||||
// the re-export through the public package specifier.
|
||||
// `./core/compaction/index.ts`); the core API has since become
|
||||
// `Tokenizer.countMessage`, so the shim now defines a compat wrapper that keeps
|
||||
// the legacy export surface — a named import must not throw Bun's static
|
||||
// "Export named X not found" during plugin validation (e.g.
|
||||
// `omp plugin install pi-blackhole`). This pins the export through the public
|
||||
// package specifier.
|
||||
describe("legacy shim compaction helpers", () => {
|
||||
it("re-exports estimateTokens as a callable token estimator", () => {
|
||||
it("exports estimateTokens as a callable token estimator", () => {
|
||||
expect(typeof estimateTokens).toBe("function");
|
||||
const tokens = estimateTokens({ role: "user", content: "hello world", timestamp: Date.now() }, new Tokenizer());
|
||||
expect(tokens).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it("counts tokens without a tokenizer argument (the legacy pi call shape)", () => {
|
||||
const tokens = estimateTokens({ role: "user", content: "hello world", timestamp: Date.now() });
|
||||
expect(tokens).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
*/
|
||||
import { describe, expect, it } from "bun:test";
|
||||
import { type } from "@oh-my-pi/omptype";
|
||||
import { Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import { arkToWireSchema } from "@oh-my-pi/pi-ai/utils/schema";
|
||||
import {
|
||||
type ContextBreakdown,
|
||||
@@ -16,18 +17,22 @@ import {
|
||||
renderContextUsage,
|
||||
} from "@oh-my-pi/pi-coding-agent/modes/utils/context-usage";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
describe("estimateToolSchemaTokens", () => {
|
||||
it("counts arktype tool schemas by their wire JSON Schema, not arktype internals", () => {
|
||||
const parameters = type({
|
||||
"query /** search query */": "string",
|
||||
"limit?": "number",
|
||||
});
|
||||
const arktypeEstimate = estimateToolSchemaTokens([
|
||||
{ name: "web_search", description: "Searches the web.", parameters } as never,
|
||||
]);
|
||||
const wireEstimate = estimateToolSchemaTokens([
|
||||
{ name: "web_search", description: "Searches the web.", parameters: arkToWireSchema(parameters) } as never,
|
||||
]);
|
||||
const arktypeEstimate = estimateToolSchemaTokens(
|
||||
[{ name: "web_search", description: "Searches the web.", parameters } as never],
|
||||
tokenizer,
|
||||
);
|
||||
const wireEstimate = estimateToolSchemaTokens(
|
||||
[{ name: "web_search", description: "Searches the web.", parameters: arkToWireSchema(parameters) } as never],
|
||||
tokenizer,
|
||||
);
|
||||
expect(arktypeEstimate).toBe(wireEstimate);
|
||||
});
|
||||
});
|
||||
@@ -103,38 +108,40 @@ describe("computeNonMessageTokens / computeNonMessageBreakdown memoization", ()
|
||||
|
||||
it("recomputes when the system prompt reference changes and caches otherwise", () => {
|
||||
const session = makeSession(["system prompt alpha"]);
|
||||
const first = computeNonMessageTokens(session as never);
|
||||
const first = computeNonMessageTokens(session as never, tokenizer);
|
||||
// Same inputs (identical refs) → cached, identical value.
|
||||
expect(computeNonMessageTokens(session as never)).toBe(first);
|
||||
expect(computeNonMessageTokens(session as never, tokenizer)).toBe(first);
|
||||
// Replace the system prompt reference (mirrors setSystemPrompt).
|
||||
session.systemPrompt = ["system prompt beta with more tokens than alpha"];
|
||||
const afterChange = computeNonMessageTokens(session as never);
|
||||
const afterChange = computeNonMessageTokens(session as never, tokenizer);
|
||||
expect(afterChange).toBeGreaterThan(first);
|
||||
// Cached on the new inputs.
|
||||
expect(computeNonMessageTokens(session as never)).toBe(afterChange);
|
||||
expect(computeNonMessageTokens(session as never, tokenizer)).toBe(afterChange);
|
||||
});
|
||||
|
||||
it("recomputes the breakdown when the tools reference changes", () => {
|
||||
const session = makeSession(["base"], []);
|
||||
const before = computeNonMessageBreakdown(session as never);
|
||||
const before = computeNonMessageBreakdown(session as never, tokenizer);
|
||||
expect(before.toolsTokens).toBe(0);
|
||||
// New tools array reference (mirrors setTools).
|
||||
session.agent.state.tools = [{ name: "search", description: "search the web", parameters: {} }];
|
||||
const after = computeNonMessageBreakdown(session as never);
|
||||
const after = computeNonMessageBreakdown(session as never, tokenizer);
|
||||
expect(after.toolsTokens).toBeGreaterThan(0);
|
||||
// Cached on the new tools.
|
||||
expect(computeNonMessageBreakdown(session as never).toolsTokens).toBe(after.toolsTokens);
|
||||
expect(computeNonMessageBreakdown(session as never, tokenizer).toolsTokens).toBe(after.toolsTokens);
|
||||
});
|
||||
|
||||
it("shares one cache entry so tokens and breakdown invalidate together", () => {
|
||||
const session = makeSession(["shared prompt"]);
|
||||
const tokens = computeNonMessageTokens(session as never);
|
||||
const breakdown = computeNonMessageBreakdown(session as never);
|
||||
const tokens = computeNonMessageTokens(session as never, tokenizer);
|
||||
const breakdown = computeNonMessageBreakdown(session as never, tokenizer);
|
||||
// Changing the system prompt ref must invalidate BOTH fields, not just
|
||||
// the one most recently touched.
|
||||
session.systemPrompt = ["shared prompt but longer now to shift the count"];
|
||||
expect(computeNonMessageTokens(session as never)).not.toBe(tokens);
|
||||
expect(computeNonMessageBreakdown(session as never).systemPromptTokens).not.toBe(breakdown.systemPromptTokens);
|
||||
expect(computeNonMessageTokens(session as never, tokenizer)).not.toBe(tokens);
|
||||
expect(computeNonMessageBreakdown(session as never, tokenizer).systemPromptTokens).not.toBe(
|
||||
breakdown.systemPromptTokens,
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -157,16 +164,16 @@ describe("computeNonMessageBreakdown skills filtering", () => {
|
||||
}
|
||||
|
||||
it("excludes hidden skills and does not clamp System prompt to 0", () => {
|
||||
const b = computeNonMessageBreakdown(session([readTool], [hidden, visible]));
|
||||
const b = computeNonMessageBreakdown(session([readTool], [hidden, visible]), tokenizer);
|
||||
// Only the visible skill is counted, not the large hidden one.
|
||||
expect(b.skillsTokens).toBe(computeNonMessageBreakdown(session([readTool], [visible])).skillsTokens);
|
||||
expect(b.skillsTokens).toBe(computeNonMessageBreakdown(session([readTool], [visible]), tokenizer).skillsTokens);
|
||||
expect(b.skillsTokens).toBeLessThan(100);
|
||||
expect(b.systemPromptTokens).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it("counts zero Skills tokens when the read tool is unavailable", () => {
|
||||
const b = computeNonMessageBreakdown(session([], [hidden, visible]));
|
||||
const b = computeNonMessageBreakdown(session([], [hidden, visible]), tokenizer);
|
||||
expect(b.skillsTokens).toBe(0);
|
||||
expect(b.systemPromptTokens).toBe(computeNonMessageBreakdown(session([], [])).systemPromptTokens);
|
||||
expect(b.systemPromptTokens).toBe(computeNonMessageBreakdown(session([], []), tokenizer).systemPromptTokens);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import { describe, expect, it } from "bun:test";
|
||||
import type { AgentToolCall } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AgentToolCall, Tokenizer } from "@oh-my-pi/pi-agent-core";
|
||||
import type { SessionMessageEntry } from "@oh-my-pi/pi-agent-core/compaction/entries";
|
||||
import { DEFAULT_PRUNE_CONFIG, pruneToolOutputs } from "@oh-my-pi/pi-agent-core/compaction/pruning";
|
||||
import { AGGRESSIVE_SHAKE_CONFIG, collectShakeRegions } from "@oh-my-pi/pi-agent-core/compaction/shake";
|
||||
@@ -7,6 +7,8 @@ import type { ProtectedToolContext } from "@oh-my-pi/pi-agent-core/compaction/to
|
||||
import type { AssistantMessage, TextContent, ToolResultMessage, Usage } from "@oh-my-pi/pi-ai";
|
||||
import { createPlanReadMatcher } from "@oh-my-pi/pi-coding-agent/plan-mode/plan-protection";
|
||||
|
||||
const tokenizer = new Tokenizer();
|
||||
|
||||
function context(opts: { toolName?: string; callName?: string | undefined; path?: string }): ProtectedToolContext {
|
||||
const toolResult = {
|
||||
role: "toolResult",
|
||||
@@ -131,7 +133,7 @@ describe("plan-read protection in compaction", () => {
|
||||
fileResult,
|
||||
];
|
||||
|
||||
const result = pruneToolOutputs(entries, {
|
||||
const result = pruneToolOutputs(entries, tokenizer, {
|
||||
...DEFAULT_PRUNE_CONFIG,
|
||||
protectTokens: 0,
|
||||
minimumSavings: 0,
|
||||
@@ -153,7 +155,7 @@ describe("plan-read protection in compaction", () => {
|
||||
fileResult,
|
||||
];
|
||||
|
||||
const regions = collectShakeRegions(entries, {
|
||||
const regions = collectShakeRegions(entries, tokenizer, {
|
||||
...AGGRESSIVE_SHAKE_CONFIG,
|
||||
protectTokens: 0,
|
||||
protectedTools: [...AGGRESSIVE_SHAKE_CONFIG.protectedTools, matcher],
|
||||
|
||||
@@ -2,6 +2,10 @@
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- Added `ClaudeV3`/`ClaudeV47`/`ClaudeV5` encodings to `countTokens`: a Rust rewrite of [ctok](https://github.com/sanderland/ctok) by Sander Land (MIT), reconstructing Anthropic's `count_tokens` offline. Counts are exact on ctok's ~3.4M-response measurement corpora; the port is validated against 493 Python-ctok reference fixtures covering all three families. The pipeline is byte-level throughout — markers occupy one byte, normalization borrows text no rule touches, ASCII and ideographs skip the Unicode tables, and pieces are matched with one Aho-Corasick transition per byte instead of a per-position vocabulary descent — which counts English prose at 64 MiB/s, markdown at 73 MiB/s, source code at 35 MiB/s and CJK at 49 MiB/s per core: 1.5× (CJK, already cheap per byte) to 5.5× (prose, markdown, digits) a straightforward character-level implementation of the same model, which is held to byte-for-byte identical counts across 2.4M randomized differential comparisons.
|
||||
|
||||
## [17.3.8] - 2026-08-19
|
||||
|
||||
### Changed
|
||||
|
||||
Vendored
+13
-4
@@ -634,9 +634,10 @@ export declare function cosineSimilarityPairs(vectors: Float64Array, count: numb
|
||||
* pool is available). Always returns a single token total — use this for any
|
||||
* aggregate budget question without paying a per-element napi crossing.
|
||||
*
|
||||
* Uses ordinary encoding (no special-token handling), which is the right
|
||||
* choice for measuring user/model content rather than wire-protocol tokens.
|
||||
* Defaults to `o200k_base`; pass `Cl100kBase` for older `OpenAI` models.
|
||||
* Measures user/model content, not wire-protocol tokens: BPE encodings use
|
||||
* ordinary encoding (no special-token handling) and the Claude encodings
|
||||
* count message content without the fixed per-message frame. Defaults to
|
||||
* `o200k_base`; pass a `Claude*` encoding for exact Claude counts.
|
||||
*/
|
||||
export declare function countTokens(input: string | Array<string>, encoding?: Encoding | undefined | null): number
|
||||
|
||||
@@ -853,7 +854,15 @@ export declare enum Encoding {
|
||||
/** GPT-4o / o1 / GPT-5 (default). */
|
||||
O200kBase = 'O200kBase',
|
||||
/** GPT-3.5 / GPT-4 / older. */
|
||||
Cl100kBase = 'Cl100kBase'
|
||||
Cl100kBase = 'Cl100kBase',
|
||||
/** Claude 3 … Opus 4.6 (ctok v3 reconstruction). */
|
||||
ClaudeV3 = 'ClaudeV3',
|
||||
/** Claude Opus 4.7–4.9 (ctok v4.7 reconstruction). */
|
||||
ClaudeV47 = 'ClaudeV47',
|
||||
/** Claude Opus 5+ (ctok v5 reconstruction). */
|
||||
ClaudeV5 = 'ClaudeV5',
|
||||
/** Claude Sonnet/Fable 5+ (live-measured non-opus v5 frame). */
|
||||
ClaudeV5Sonnet = 'ClaudeV5Sonnet'
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -101,6 +101,10 @@ export const Ellipsis = {
|
||||
export const Encoding = {
|
||||
O200kBase: "O200kBase",
|
||||
Cl100kBase: "Cl100kBase",
|
||||
ClaudeV3: "ClaudeV3",
|
||||
ClaudeV47: "ClaudeV47",
|
||||
ClaudeV5: "ClaudeV5",
|
||||
ClaudeV5Sonnet: "ClaudeV5Sonnet",
|
||||
};
|
||||
export const FileType = {
|
||||
File: 1,
|
||||
|
||||
@@ -37,6 +37,7 @@
|
||||
"test": "bun test --parallel",
|
||||
"fix": "biome check --write --unsafe .",
|
||||
"fmt": "biome format --write .",
|
||||
"gen:ctok": "bun scripts/gen-ctok-vocab.ts",
|
||||
"gen:native": "bun scripts/embed-native.ts",
|
||||
"gen:native:reset": "bun scripts/embed-native.ts --reset",
|
||||
"gen:npm": "bun scripts/gen-npm-packages.ts",
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
/**
|
||||
* Regenerates the compact ctok vocabulary data embedded by the Rust ctok port
|
||||
* (crates/pi-natives/src/ctok/data/ctok_*.bin).
|
||||
*
|
||||
* Source of truth is the measured vocabulary of sanderland/ctok (MIT), pinned
|
||||
* to a release revision. Compaction drops the per-piece witness metadata,
|
||||
* parses the public `⟨bow⟩the⟨eow⟩` notation into the single-glyph internal
|
||||
* form, adds the glued contraction spellings, and emits the front-coded
|
||||
* binary format below — cutting ~4.7 MB of upstream JSON to ~350 KB of
|
||||
* embedded data. If the pin moves, also regenerate
|
||||
* crates/pi-natives/src/ctok/testdata/fixtures.json against the same ctok
|
||||
* release (`uv run --with ctok …`; see the fixture doc in ctok/mod.rs).
|
||||
*
|
||||
* Format (little-endian; parsed by `VocabCore::parse` in ctok/engine.rs):
|
||||
*
|
||||
* magic b"CTOK"
|
||||
* version u8 = 2
|
||||
* flags u8 (bit 0: fold_quotes)
|
||||
* message_overhead u8
|
||||
* allcaps_min u8 (0 = disabled)
|
||||
* byte_token_count u16
|
||||
* piece_count u32
|
||||
* byte tokens count × { len u8, bytes }
|
||||
* pieces count × { shared_prefix varint, suffix_len varint, suffix bytes }
|
||||
*
|
||||
* Pieces are written in the compact alphabet — marker glyphs as the single
|
||||
* bytes above, everything else UTF-8 — and sorted by those bytes;
|
||||
* `shared_prefix` is the byte length shared with the previous piece (front
|
||||
* coding). Varints are LEB128.
|
||||
*/
|
||||
|
||||
import * as path from "node:path";
|
||||
|
||||
/** Pinned upstream: sanderland/ctok v1.0.0. */
|
||||
const CTOK_REV = "df3b59b5e645289a5eadc8e24036b99d39c333c4";
|
||||
const UPSTREAM = `https://raw.githubusercontent.com/sanderland/ctok/${CTOK_REV}/ctok/data`;
|
||||
|
||||
const DATA_DIR = path.join(import.meta.dir, "../../../crates/pi-natives/src/ctok/data");
|
||||
|
||||
/** Marker glyphs of ctok's internal marked form, keyed by public atom. */
|
||||
const ATOMS: Record<string, string> = {
|
||||
"⟨bow⟩": "\ufdd0",
|
||||
"⟨eow⟩": "\ufdd1",
|
||||
"⟨pad⟩": "\ufdd2",
|
||||
"⟨shift⟩": "\ufdd3",
|
||||
"⟨caps⟩": "\ufdd4",
|
||||
};
|
||||
|
||||
/**
|
||||
* Marker glyph → the single byte the Rust encoder writes for it (mirrors
|
||||
* `MARKERS` in ctok/constants.rs). The tokenizer strips C0 controls from input
|
||||
* before anything else, so these bytes can never collide with text, and one
|
||||
* byte per marker instead of three shrinks both the marked stream and the
|
||||
* matching automaton by about a third.
|
||||
*/
|
||||
const MARKER_BYTES: Record<string, number> = {
|
||||
"\ufdd0": 0x01,
|
||||
"\ufdd1": 0x02,
|
||||
"\ufdd2": 0x03,
|
||||
"\ufdd3": 0x04,
|
||||
"\ufdd4": 0x05,
|
||||
};
|
||||
|
||||
const encoder = new TextEncoder();
|
||||
|
||||
function encodeCompact(piece: string): Uint8Array<ArrayBuffer> {
|
||||
const out: number[] = [];
|
||||
for (const ch of piece) {
|
||||
const marker = MARKER_BYTES[ch];
|
||||
if (marker === undefined) out.push(...encoder.encode(ch));
|
||||
else out.push(marker);
|
||||
}
|
||||
return new Uint8Array(out);
|
||||
}
|
||||
|
||||
const EOW = "\ufdd1";
|
||||
const BYTE_ATOM = /^⟨0x([0-9A-Fa-f]{2})⟩/;
|
||||
|
||||
/**
|
||||
* Parse one public-notation vocabulary key into the internal marked string:
|
||||
* named atoms become single glyphs, `⟨0xNN⟩` escape runs decode back to their
|
||||
* characters, anything else is literal.
|
||||
*/
|
||||
function parseMarked(publicKey: string): string {
|
||||
let out = "";
|
||||
let bytes: number[] = [];
|
||||
const flush = () => {
|
||||
if (bytes.length === 0) return;
|
||||
out += new TextDecoder("utf-8", { fatal: true }).decode(new Uint8Array(bytes));
|
||||
bytes = [];
|
||||
};
|
||||
let rest = publicKey;
|
||||
outer: while (rest.length > 0) {
|
||||
if (rest.startsWith("⟨")) {
|
||||
for (const atom in ATOMS) {
|
||||
if (rest.startsWith(atom)) {
|
||||
flush();
|
||||
out += ATOMS[atom];
|
||||
rest = rest.slice(atom.length);
|
||||
continue outer;
|
||||
}
|
||||
}
|
||||
const byteAtom = BYTE_ATOM.exec(rest);
|
||||
if (byteAtom) {
|
||||
bytes.push(Number.parseInt(byteAtom[1], 16));
|
||||
rest = rest.slice(byteAtom[0].length);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
flush();
|
||||
const ch = String.fromCodePoint(rest.codePointAt(0) as number);
|
||||
out += ch;
|
||||
rest = rest.slice(ch.length);
|
||||
}
|
||||
flush();
|
||||
return out;
|
||||
}
|
||||
|
||||
function pushVarint(out: number[], value: number): void {
|
||||
let v = value;
|
||||
while (v >= 0x80) {
|
||||
out.push((v & 0x7f) | 0x80);
|
||||
v >>>= 7;
|
||||
}
|
||||
out.push(v);
|
||||
}
|
||||
|
||||
interface UpstreamDoc {
|
||||
meta: { message_overhead: number; fold_quotes: boolean; allcaps_min: number | null };
|
||||
tokens: Record<string, Record<string, unknown>>;
|
||||
}
|
||||
|
||||
async function generate(src: string, dst: string): Promise<void> {
|
||||
const url = `${UPSTREAM}/${src}`;
|
||||
const response = await fetch(url);
|
||||
if (!response.ok) throw new Error(`fetch ${url}: ${response.status} ${response.statusText}`);
|
||||
const doc = (await response.json()) as UpstreamDoc;
|
||||
|
||||
const pieces = new Set<string>();
|
||||
const byteTokens: string[] = [];
|
||||
for (const group in doc.tokens) {
|
||||
const entries = doc.tokens[group];
|
||||
if (group === "bytes_fallback") {
|
||||
byteTokens.push(...Object.keys(entries));
|
||||
continue;
|
||||
}
|
||||
for (const key in entries) {
|
||||
const parsed = parseMarked(key);
|
||||
if (group === "contractions") {
|
||||
// The file stores `'t`; the encoder writes `'t⟨eow⟩` (ctok's
|
||||
// glued_contraction). Both spellings join the tiling vocabulary.
|
||||
pieces.add(parsed + EOW);
|
||||
}
|
||||
pieces.add(parsed);
|
||||
}
|
||||
}
|
||||
|
||||
const sorted = [...pieces].map(encodeCompact).sort(Buffer.compare);
|
||||
|
||||
const out: number[] = [0x43, 0x54, 0x4f, 0x4b, 2]; // "CTOK", version
|
||||
out.push(doc.meta.fold_quotes ? 1 : 0);
|
||||
out.push(doc.meta.message_overhead);
|
||||
out.push(doc.meta.allcaps_min ?? 0);
|
||||
byteTokens.sort();
|
||||
out.push(byteTokens.length & 0xff, byteTokens.length >>> 8);
|
||||
const n = sorted.length;
|
||||
out.push(n & 0xff, (n >>> 8) & 0xff, (n >>> 16) & 0xff, (n >>> 24) & 0xff);
|
||||
for (const hex of byteTokens) {
|
||||
out.push(hex.length / 2);
|
||||
for (let i = 0; i < hex.length; i += 2) out.push(Number.parseInt(hex.slice(i, i + 2), 16));
|
||||
}
|
||||
let prev = new Uint8Array(0);
|
||||
for (const piece of sorted) {
|
||||
let shared = 0;
|
||||
const max = Math.min(prev.length, piece.length);
|
||||
while (shared < max && prev[shared] === piece[shared]) shared++;
|
||||
pushVarint(out, shared);
|
||||
pushVarint(out, piece.length - shared);
|
||||
for (let i = shared; i < piece.length; i++) out.push(piece[i]);
|
||||
prev = piece;
|
||||
}
|
||||
|
||||
const dstPath = path.join(DATA_DIR, dst);
|
||||
await Bun.write(dstPath, new Uint8Array(out));
|
||||
console.log(`${dst}: ${n} pieces, ${byteTokens.length} byte tokens, ${out.length} bytes`);
|
||||
}
|
||||
|
||||
await generate("pieces_v3.json", "ctok_v3.bin");
|
||||
await generate("pieces_v4_7.json", "ctok_v4_7.bin");
|
||||
Reference in New Issue
Block a user