fix(tokenizer): drain consumed tokens in DecodeStream to bound memory

[slin1237]

Description

Problem

DecodeStream::step appends every generated token to its internal all_token_ids buffer but never drains it. The two offsets (prefix_offset, read_offset) advance, yet the consumed tokens below prefix_offset are kept for the entire lifetime of the stream. For a long-lived streaming response this is unbounded memory growth per stream — the buffer grows by one entry per generated token even though only the small prefix window is ever read by subsequent decode() calls.

This mirrors a class of issue already solved elsewhere in the crate: the Decoder::decode_step default and Sequence both drain consumed tokens to keep memory bounded; DecodeStream was the outlier.

Solution

After each successful step (the branch that emits a chunk and advances the offsets), drain the tokens below the prefix window and rebase both offsets:

self.prefix_offset = self.read_offset;
self.read_offset = self.all_token_ids.len();

if self.prefix_offset > 0 {
    let drained = self.prefix_offset;
    self.all_token_ids.drain(..drained);
    self.prefix_offset -= drained;
    self.read_offset -= drained;
}

The active decode windows (prefix_offset..read_offset and read_offset..) are preserved exactly, so incremental output and flush() behavior are unchanged. Only tokens that can no longer affect any future decode are removed. The public tokens() accessor now returns the retained sliding window rather than full history (documented), matching Sequence::token_ids().

Changes

• crates/tokenizer/src/stream.rs: drain consumed tokens below prefix_offset and rebase offsets after a successful step; update tokens() doc to reflect the sliding-window semantics. Made INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET pub(crate) so the test can assert against the real window size.
• crates/tokenizer/src/tests.rs: added test_decode_stream_history_bounded (1000 steps, asserts the retained buffer stays within the detokenization window and that accumulated output equals a full decode); updated two existing tokens() assertions for the new draining behavior.
• crates/tokenizer/README.md: note that consumed tokens are drained so retained history stays bounded.

Test Plan

Scenario: stream 1000 tokens through a single DecodeStream and assert (a) the retained token buffer never exceeds the detokenization window (INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET + 1) regardless of total tokens, and (b) accumulated incremental output still equals a full decode of all tokens. Correctness is additionally guarded end-to-end by the existing real-tokenizer integration tests test_decode_stream and test_long_sequence_incremental_decode_with_prefill, which compare accumulated step() output against the original prompt.

Gate (sccache disabled, scoped to the changed crate llm-tokenizer):

$ cargo +nightly fmt --all
(clean, exit 0)

$ RUSTC_WRAPPER="" cargo clippy -p llm-tokenizer --all-targets -- -D warnings
    Checking llm-tokenizer v1.3.2 (crates/tokenizer)
    Finished `dev` profile [unoptimized + debuginfo] target(s)
clippy exit: 0

$ RUSTC_WRAPPER="" cargo test -p llm-tokenizer
test tests::test_decode_stream_history_bounded ... ok
test test_decode_stream ... ok
test test_long_sequence_incremental_decode_with_prefill ... ok
test result: ok. 132 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out   (lib)
... (all binaries)
aggregate: 199 passed; 0 failed; 13 ignored

From the codebase audit: crates/tokenizer/src/stream.rs:48 — DecodeStream never drains all_token_ids: unbounded memory per stream.

Checklist

• cargo +nightly fmt passes
• cargo clippy --all-targets --all-features -- -D warnings passes
• (Optional) Documentation updated
• (Optional) Please join us on Slack #sig-smg to discuss, review, and merge PRs

Summary by CodeRabbit

• Bug Fixes

  • Improved token stream management to drain consumed tokens and keep retained history bounded, preventing unbounded memory growth during streaming operations.

• Documentation

  • Clarified token stream behavior documentation to reflect the bounded token window retention model.

• Tests

  • Strengthened validation tests to verify bounded token history behavior and added regression test coverage.

[coderabbitai]

Warning

Review limit reached

@slin1237, we couldn't start this review because you've reached your PR review rate limit.

More reviews will be available in 1 hour, 19 minutes, and 18 seconds. Learn how PR review limits work.

Your organization has used up its prepaid credits, and credit purchases are no longer available. Enable the review add-on in the billing tab to keep reviews running — you're only billed for reviews past your plan's rate limits ($0.25/file).

⌛ How to resolve this issue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available.

Please see our Fair Usage Limits Policy for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: e6e10931-fcca-4635-ba13-6f93cc577799

📥 Commits

Reviewing files that changed from the base of the PR and between ad879fc and 7dad7d3.

📒 Files selected for processing (3)

• crates/tokenizer/README.md
• crates/tokenizer/src/stream.rs
• crates/tokenizer/src/tests.rs

📝 Walkthrough

Walkthrough

DecodeStream implementation bounds retained token history by draining consumed tokens below a prefix window after each decoding step. The window size constant is exposed, offsets are rebased to keep memory bounded, and tests validate the behavior through updated assertions and new regression coverage.

Changes

Bounded token history in DecodeStream

Layer / File(s)	Summary
Implementation and documentation crates/tokenizer/src/stream.rs, crates/tokenizer/README.md	INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET is exposed as pub(crate). DecodeStream::step now drains token ids below the retained prefix window and rebases prefix_offset/read_offset to keep history bounded. DecodeStream::tokens() and README documentation updated to describe the bounded sliding window behavior.
Test validation of bounded history crates/tokenizer/src/tests.rs	Tests import the prefix window constant and update test_decode_stream_basic and test_decode_stream_multiple_steps to verify that consumed tokens are drained after steps. New regression test test_decode_stream_history_bounded exercises 1000 decoding steps, asserts token history stays within window bounds, and validates streamed output matches a full decode.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Suggested labels

tokenizer, tests

Suggested reviewers

• CatherineSue

Poem

🐰 A stream that flows with bounded grace,
Each token finds its rightful place,
Old history drained without a trace,
The window slides at bounded pace,
A thousand steps keep memory spaced!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'fix(tokenizer): drain consumed tokens in DecodeStream to bound memory' directly and clearly describes the main change: fixing a memory-growth bug by draining consumed tokens in DecodeStream.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/tokenizer-decodestream-unbounded

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request bounds the memory usage of DecodeStream by draining consumed tokens below the prefix window after each successful step and rebasing the internal offsets. It also updates the documentation and existing tests to reflect this sliding-window behavior, and adds a new integration test to verify that the retained history remains bounded over many steps while still producing correct decoded output. There are no review comments, and I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[claude]

Clean, correct fix. The drain logic is safe — tokens before prefix_offset are provably never accessed again by step() or flush(). Offset rebasing preserves all invariants. No production callers of tokens() are affected by the semantic change. Tests cover both boundedness and output correctness over 1000 steps.

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 14 days. It will be automatically closed if no further activity occurs within 16 days. Leave a comment if you feel this pull request should remain open. Thank you!

[github-actions]

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!