FIX: Chunk large inputs to prevent tiktoken-rs stack overflow

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

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

lib/discourse_ai/tokenizers/version.rb

@@ -2,6 +2,6 @@
 
 module DiscourseAi
   module Tokenizers
-    VERSION = "0.4"
+    VERSION = "0.4.1"
   end
 end

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