From c4cef25f76f126b2aac29200ac757ff36a31e28c Mon Sep 17 00:00:00 2001 From: Rafael Silva Date: Thu, 26 Feb 2026 13:39:01 -0300 Subject: [PATCH] FIX: Chunk large inputs to prevent tiktoken-rs stack overflow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tiktoken-rs v0.9.1 uses fancy-regex for BPE tokenization, which can stack overflow on large inputs due to catastrophic backtracking (https://github.com/openai/tiktoken/issues/245). The Rust code unwraps the error, panics across the FFI boundary, and aborts the entire Ruby process — this is not rescuable from Ruby. A secondary issue is quadratic BPE merge time for long non-whitespace runs (https://github.com/openai/tiktoken/issues/195), which causes the process to hang indefinitely. Python tiktoken fixed the backtracking in v0.8.0 with possessive quantifiers (https://github.com/openai/tiktoken/pull/258), but tiktoken-rs 0.9.1 and tiktoken_ruby 0.0.15.1 have not ported the fix, and no newer versions are available. This adds a safe_encode method that chunks text larger than 50K characters at whitespace boundaries before passing to tiktoken. BPE token boundaries never span whitespace, so chunked encoding produces identical results for normal text. For pathological inputs (e.g. 500K repeated characters with no whitespace), it completes in seconds instead of crashing or hanging. --- CHANGELOG.md | 4 + .../tokenizer/open_ai_tokenizer.rb | 38 +++++- lib/discourse_ai/tokenizers/version.rb | 2 +- .../tokenizers/safe_encode_spec.rb | 113 ++++++++++++++++++ 4 files changed, 153 insertions(+), 4 deletions(-) create mode 100644 spec/discourse_ai/tokenizers/safe_encode_spec.rb diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ca03b0..650987a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## [Unreleased] +## [0.4.1] - 2026-02-26 + +- Fix tiktoken-rs stack overflow crash by chunking large inputs at whitespace boundaries before encoding + ## [0.4.0] - 2026-01-06 - Add Ruby 4.0 compatibility diff --git a/lib/discourse_ai/tokenizer/open_ai_tokenizer.rb b/lib/discourse_ai/tokenizer/open_ai_tokenizer.rb index f43789c..4bfb8ce 100644 --- a/lib/discourse_ai/tokenizer/open_ai_tokenizer.rb +++ b/lib/discourse_ai/tokenizer/open_ai_tokenizer.rb @@ -4,17 +4,22 @@ module DiscourseAi module Tokenizer # Wrapper for OpenAI tokenizer library for compatibility with Discourse AI API class OpenAiTokenizer < BasicTokenizer + # tiktoken-rs uses fancy-regex which can stack overflow on large inputs + # due to catastrophic backtracking (github.com/openai/tiktoken/issues/245). + # Chunking at whitespace boundaries prevents this while preserving accuracy. + SAFE_CHUNK_SIZE = 50_000 + class << self def tokenizer @tokenizer ||= Tiktoken.get_encoding("o200k_base") end def tokenize(text) - tokenizer.encode(text) + safe_encode(text) end def encode(text) - tokenizer.encode(text) + safe_encode(text) end def decode(token_ids) @@ -72,7 +77,34 @@ def below_limit?(text, limit, strict: false) # than can take more than 1 token per char return true if !strict && text.size < limit / 2 - tokenizer.encode(text).length < limit + safe_encode(text).length < limit + end + + private + + def safe_encode(text) + if !text.is_a?(String) || text.size <= SAFE_CHUNK_SIZE + return tokenizer.encode(text) + end + + tokens = [] + offset = 0 + while offset < text.size + chunk_end = offset + SAFE_CHUNK_SIZE + + if chunk_end < text.size + # Split at a whitespace boundary to preserve tokenization accuracy + break_point = text.rindex(/\s/, chunk_end) + chunk_end = break_point if break_point && break_point > offset + else + chunk_end = text.size + end + + tokens.concat(tokenizer.encode(text[offset...chunk_end])) + offset = chunk_end + end + + tokens end end end diff --git a/lib/discourse_ai/tokenizers/version.rb b/lib/discourse_ai/tokenizers/version.rb index 872d415..9b6d014 100644 --- a/lib/discourse_ai/tokenizers/version.rb +++ b/lib/discourse_ai/tokenizers/version.rb @@ -2,6 +2,6 @@ module DiscourseAi module Tokenizers - VERSION = "0.4" + VERSION = "0.4.1" end end diff --git a/spec/discourse_ai/tokenizers/safe_encode_spec.rb b/spec/discourse_ai/tokenizers/safe_encode_spec.rb new file mode 100644 index 0000000..3b48d24 --- /dev/null +++ b/spec/discourse_ai/tokenizers/safe_encode_spec.rb @@ -0,0 +1,113 @@ +# frozen_string_literal: true + +require "spec_helper" + +RSpec.describe DiscourseAi::Tokenizer::OpenAiTokenizer do + let(:chunk_size) { described_class::SAFE_CHUNK_SIZE } + + shared_examples "safe encoding" do |tokenizer_class| + describe "text under chunk size" do + it "encodes normally without chunking" do + text = "Hello world, this is a normal sentence." + raw_tokenizer = tokenizer_class.tokenizer + + expect(tokenizer_class.encode(text)).to eq(raw_tokenizer.encode(text)) + end + end + + describe "large text with whitespace" do + it "produces identical tokens to direct encoding" do + text = "The quick brown fox jumps over the lazy dog. " * 2000 + expect(text.size).to be > chunk_size + + direct_tokens = tokenizer_class.tokenizer.encode(text) + chunked_tokens = tokenizer_class.encode(text) + + expect(chunked_tokens).to eq(direct_tokens) + end + + it "works with tokenize" do + text = "word " * 20_000 + expect(text.size).to be > chunk_size + + tokens = tokenizer_class.tokenize(text) + expect(tokens).to be_an(Array) + expect(tokens.length).to be > 0 + end + + it "works with size" do + text = "word " * 20_000 + expect(text.size).to be > chunk_size + + direct_size = tokenizer_class.tokenizer.encode(text).length + expect(tokenizer_class.size(text)).to eq(direct_size) + end + end + + describe "large text without whitespace" do + it "encodes repeated characters without crashing" do + text = "M" * (chunk_size + 1000) + + tokens = tokenizer_class.encode(text) + expect(tokens).to be_an(Array) + expect(tokens.length).to be > 0 + end + + it "encodes long non-whitespace runs without crashing" do + text = ("a".."z").to_a.join * 3000 + expect(text.size).to be > chunk_size + + tokens = tokenizer_class.encode(text) + expect(tokens).to be_an(Array) + expect(tokens.length).to be > 0 + end + end + + describe "chunking splits at whitespace boundaries" do + it "splits at the last whitespace before chunk boundary" do + before_boundary = "a" * (chunk_size - 10) + after_boundary = "b" * 100 + text = "#{before_boundary} #{after_boundary}" + + direct_tokens = + tokenizer_class.tokenizer.encode(before_boundary) + + tokenizer_class.tokenizer.encode(" #{after_boundary}") + + expect(tokenizer_class.encode(text)).to eq(direct_tokens) + end + end + + describe "below_limit? with large text" do + it "returns correct result for text over chunk size" do + text = "word " * 20_000 + expect(text.size).to be > chunk_size + + actual_token_count = tokenizer_class.tokenizer.encode(text).length + expect( + tokenizer_class.below_limit?(text, actual_token_count + 1) + ).to be true + expect(tokenizer_class.below_limit?(text, 1)).to be false + end + end + + describe "truncate with large text" do + it "truncates large text correctly" do + text = "word " * 20_000 + expect(text.size).to be > chunk_size + + result = tokenizer_class.truncate(text, 100, strict: true) + expect(tokenizer_class.size(result)).to be <= 100 + expect(result.length).to be < text.length + end + end + end + + describe DiscourseAi::Tokenizer::OpenAiTokenizer do + include_examples "safe encoding", DiscourseAi::Tokenizer::OpenAiTokenizer + end + + describe DiscourseAi::Tokenizer::OpenAiCl100kTokenizer do + include_examples "safe encoding", + DiscourseAi::Tokenizer::OpenAiCl100kTokenizer + end +end