feat(mesh): expose model-bound tokenizer capability - #1227
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe tokenizer contract now supports bounded structured encoding of byte runs and opaque control descriptors. Native serving plugins use a model-bound tokenizer capability through ABI V2, with host-side validation, callback handling, binding metadata, and explicit error statuses. Standalone N-gram speculation is also enabled for non-staged serving when native MTP is disabled. ChangesTokenizer capability and ABI integration
Standalone N-gram speculation
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Plugin
participant HostTokenizerCapability
participant TokenizerCapability
participant TokenizerSource
Plugin->>HostTokenizerCapability: invoke encode callback
HostTokenizerCapability->>TokenizerCapability: validate pieces and limits
TokenizerCapability->>TokenizerSource: encode structured input
TokenizerSource-->>TokenizerCapability: return token IDs
TokenizerCapability-->>HostTokenizerCapability: return response
HostTokenizerCapability-->>Plugin: write tokens and return status
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
crates/skippy-tokenizer/src/lib.rs (1)
15-15: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winTie the piece bound to the ABI constant.
MAX_TOKENIZE_PIECESis 4,096 here andMAX_TOKENIZER_INPUT_PIECESis 4,096 incrates/mesh-native-serving-plugin-api/src/lib.rs. The host checks the ABI constant and the server checks this constant. If one value changes, the two layers report different statuses for the same request. Add a compile-time assertion incrates/mesh-native-serving-plugin-host/src/lib.rsso the values cannot drift.♻️ Suggested assertion in the host crate
const _: () = assert!( abi::MAX_TOKENIZER_INPUT_PIECES == skippy_tokenizer::MAX_TOKENIZE_PIECES, "ABI and tokenizer piece bounds must match", );🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-tokenizer/src/lib.rs` at line 15, Add a compile-time assertion in the host crate’s lib module comparing abi::MAX_TOKENIZER_INPUT_PIECES with skippy_tokenizer::MAX_TOKENIZE_PIECES, and fail compilation with a clear message when they differ. Keep both existing constants unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/mesh-native-serving-plugin-host/src/lib.rs`:
- Around line 266-340: The tokenizer capability code in
crates/mesh-native-serving-plugin-host/src/lib.rs#L266-L340 must move into a new
tokenizer_capability module, including HostTokenizerCapability,
ActivationInventory, encode_tokenizer, write_encode_output, encode_status, and
the inventory test; update module visibility/imports while preserving behavior.
The binding helpers in crates/skippy-server/src/tokenizer.rs#L441-L474 must move
into a tokenizer/binding module, including inventory_from_stage,
tokenizer_binding_digest, update_string, and their related tests; update
references and keep each new module under 1,000 lines.
- Around line 219-263: Update TokenizerCapability::new so the Arc<Self> is
created before initializing abi.inventory and abi.context, then use Arc::get_mut
to assign both pointers from the finalized Arc allocation. Preserve the existing
inventory_view and tokenizer references, ensuring encode_tokenizer and
activation reads access memory rooted in the returned Arc.
In `@crates/skippy-server/src/tokenizer.rs`:
- Around line 157-184: Update encode to coalesce each consecutive run of
InputPiece::Bytes before UTF-8/NUL validation and tokenization, preserving the
unsupported-input error for non-byte pieces. Tokenize each merged run once so
split UTF-8 characters and token IDs match a single-piece input, and avoid
calling tokenize with a zero remaining limit by returning once max_tokens is
reached or otherwise passing the actual positive bound. Add a test comparing
token IDs for a byte run split across two pieces with the equivalent
single-piece input.
---
Nitpick comments:
In `@crates/skippy-tokenizer/src/lib.rs`:
- Line 15: Add a compile-time assertion in the host crate’s lib module comparing
abi::MAX_TOKENIZER_INPUT_PIECES with skippy_tokenizer::MAX_TOKENIZE_PIECES, and
fail compilation with a clear message when they differ. Keep both existing
constants unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 7a4f83c6-0b7a-4ca5-869e-0981b8950d5e
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (10)
crates/mesh-native-serving-plugin-api/README.mdcrates/mesh-native-serving-plugin-api/src/lib.rscrates/mesh-native-serving-plugin-host/Cargo.tomlcrates/mesh-native-serving-plugin-host/README.mdcrates/mesh-native-serving-plugin-host/src/lib.rscrates/mesh-native-serving-plugin-host/src/test_support.rscrates/skippy-server/README.mdcrates/skippy-server/src/tokenizer.rscrates/skippy-tokenizer/README.mdcrates/skippy-tokenizer/src/lib.rs
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (6)
crates/skippy-server/src/tokenizer/binding.rs (2)
21-34: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPreserve the failure reason instead of collapsing every path to
None.Five distinct root causes return
Nonehere: no GGUF file on disk, an unscannable GGUF, a token-ID conversion overflow, atoken_pieceserror, and a vocabulary/piece length mismatch. The host convertsNoneinto a single message, "bound model does not expose a tokenizer inventory" (crates/mesh-native-serving-plugin-host/src/tokenizer_capability.rsLines 84-88), and then fails plugin activation.Line 31 also discards the
TokenizerCapabilityErrorfromsource.token_pieces. A transientRuntimeUnavailablebecomes indistinguishable from a permanently unsupported model.Return
Result<TokenizerInventory, TokenizerCapabilityError>(or log the reason at each early return) so an operator can tell these cases apart from the activation failure alone.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-server/src/tokenizer/binding.rs` around lines 21 - 34, The inventory_from_stage function currently collapses missing files, scan failures, ID conversion overflow, token_pieces errors, and length mismatches into None. Change its return type and each early-return path to preserve distinct TokenizerCapabilityError reasons, including propagating the error from source.token_pieces, so the host can distinguish transient failures from unsupported inventories.
1-8: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the relocated test here and add digest coverage.
This new module has no test module. The
source_gguf_pathtest stayed incrates/skippy-server/src/tokenizer.rs(Lines 882-896), andtokenizer_binding_digesthas no direct test.Add a
#[cfg(test)]module here that covers digest determinism for identical inputs and digest change when a token ID, a piece kind, or piece bytes change. Move thesource_gguf_pathfallback test into this module.As per coding guidelines: "When modifying a source file over 1,000 lines, extract any separable responsibility into a semantically named module, keep the new file under 1,000 lines, and move relevant tests with the extracted behavior when practical."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-server/src/tokenizer/binding.rs` around lines 1 - 8, Add a #[cfg(test)] module in binding.rs covering tokenizer_binding_digest determinism for identical inputs and digest changes when token IDs, piece kinds, or piece bytes differ. Move the source_gguf_path fallback test from tokenizer.rs into this module, preserving its existing behavior and using the binding module’s relevant symbols.Source: Coding guidelines
crates/mesh-native-serving-plugin-host/src/lib.rs (1)
129-129: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument that this field is load-bearing, and drop the underscore prefix.
The
_prefix marks a field as unused. This field is not unused:context.tokenizer_capabilityat Line 104 points into thisArc<HostTokenizerCapability>allocation, and the plugin keeps that pointer for its whole active lifetime. If a later change removes the field as dead state, the plugin dereferences freed memory.Rename the field to
tokenizer_capabilityand add a comment on theActivePluginfield declaration that states the raw ABI pointers handed to the plugin remain valid only while this field is alive.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/mesh-native-serving-plugin-host/src/lib.rs` at line 129, Rename the ActivePlugin field from _tokenizer_capability to tokenizer_capability and update its initializer and any references accordingly. Add a field-level comment explaining that the raw ABI pointers provided to the plugin remain valid only while this Arc<HostTokenizerCapability> is alive.crates/skippy-server/src/tokenizer.rs (1)
606-631: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCollapse the two loops and consider making the stub content-sensitive.
The check at Lines 611-618 already rejects every
InputPiece::Control, so thelet ... elseat Lines 620-622 is unreachable. A singlematchover the pieces removes the dead branch.The stub also ignores piece content and returns a fixed token list. A stub that derives tokens from the merged bytes would let a test assert that a byte run split across two pieces produces the same token IDs as the single-piece input, which guards the segmentation behavior of
encode.♻️ Proposed refactor
fn encode( &self, pieces: &[InputPiece], max_tokens: usize, ) -> Result<Vec<i32>, TokenizerCapabilityError> { - if pieces - .iter() - .any(|piece| matches!(piece, InputPiece::Control { .. })) - { - return Err(TokenizerCapabilityError::UnsupportedInput { - reason: "recording source does not support controls".to_owned(), - }); - } - for piece in pieces { - let InputPiece::Bytes(bytes) = piece else { - unreachable!(); - }; - std::str::from_utf8(bytes).map_err(|_| { - TokenizerCapabilityError::UnsupportedInput { - reason: "recording source only accepts UTF-8".to_owned(), - } - })?; - } - self.tokenize("", false, max_tokens) + let mut merged = Vec::new(); + for piece in pieces { + match piece { + InputPiece::Bytes(bytes) => merged.extend_from_slice(bytes), + InputPiece::Control { .. } => { + return Err(TokenizerCapabilityError::UnsupportedInput { + reason: "recording source does not support controls".to_owned(), + }); + } + } + } + std::str::from_utf8(&merged).map_err(|_| { + TokenizerCapabilityError::UnsupportedInput { + reason: "recording source only accepts UTF-8".to_owned(), + } + })?; + self.tokenize("", false, max_tokens) }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/skippy-server/src/tokenizer.rs` around lines 606 - 631, Update encode to validate pieces in one match-based pass, rejecting controls and invalid UTF-8 without retaining the unreachable InputPiece::Bytes else branch. Merge the validated byte contents in order and derive the returned token IDs from that merged input instead of always calling tokenize with an empty string, while preserving max_tokens handling and existing error types.crates/mesh-native-serving-plugin-host/src/tokenizer_capability.rs (2)
69-80: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winCorrect the SAFETY comment to state the real aliasing invariant.
Two details in the comment do not match the code.
The value is an
Arc<Self>, not a boxed value.More importantly,
abi.inventorydoes not only referenceself.inventory.ActivationInventory::from_inventorybuilds eachByteSliceat Line 52 from the&abi::TokenizerInventoryargument, soself.inventory.entries[i].bytespoints into heap buffers owned byself.tokenizer. The soundness condition is therefore thatself.tokenizerand its inventory are never mutated, reallocated, or dropped while the plugin holdsabi.inventory.State that condition explicitly. Note also that
tokenizeris declared beforeinventory, so it drops first; that is safe only because droppingVec<TokenizerInventoryEntry>never dereferences aByteSlice.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/mesh-native-serving-plugin-host/src/tokenizer_capability.rs` around lines 69 - 80, Update the SAFETY comment above the Send/Sync implementations for HostTokenizerCapability to refer to the Arc<Self> lifetime rather than a boxed value, and explicitly state that tokenizer and its inventory remain immutable, non-reallocated, and alive while the plugin holds abi.inventory because its ByteSlice values reference tokenizer-owned buffers. Note that tokenizer drops before inventory and this is safe because dropping Vec<TokenizerInventoryEntry> does not dereference ByteSlice.
255-284: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for the validation and status-mapping paths.
This module is the FFI trust boundary, and only the
ActivationInventoryhappy path has a test. Add tests for these cases:
ActivationInventory::from_inventoryrejects a wrongschema_version, an empty token list, a list overMAX_TOKENIZER_INVENTORY_ENTRIES, non-increasing IDs, and an empty control descriptor.write_encode_outputreturnsOUTPUT_TOO_SMALLand still writes the required length, so a plugin can size its buffer and retry.encode_statusmaps eachTokenizerErrorvariant to the intendedTokenizerEncodeStatus.HostTokenizerCapability::newexposes identity bytes that match the bound tokenizer identity, which also guards the dangling-slice defect flagged on Lines 118-127.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/mesh-native-serving-plugin-host/src/tokenizer_capability.rs` around lines 255 - 284, Add focused tests in the existing tests module covering the validation failures in ActivationInventory::from_inventory: incorrect schema_version, empty tokens, more than MAX_TOKENIZER_INVENTORY_ENTRIES, non-increasing token IDs, and an empty control descriptor. Add a write_encode_output test that verifies OUTPUT_TOO_SMALL and the required length is populated for retry. Test every TokenizerError variant through encode_status against its expected TokenizerEncodeStatus, and verify HostTokenizerCapability::new exposes identity bytes equal to the bound tokenizer identity, retaining the returned capability while inspecting the identity slice.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/mesh-native-serving-plugin-host/src/tokenizer_capability.rs`:
- Around line 140-152: In the tokenizer encoding function, after validating that
output_length is non-null, initialize *output_length to zero before the
remaining INVALID_ARGUMENT checks. Remove the redundant zero writes from the
LIMIT_EXCEEDED paths, and update the documented ABI contract in the tokenizer
plugin API README to state that non-OK results report an output length of zero.
- Around line 118-127: Update the tokenizer capability constructor around
TokenizerCapability creation so abi.model_id, abi.source_model_sha256, and
abi.tokenizer_id are built from the tokenizer parameter before it is moved into
the Arc, ensuring their buffers remain valid for the capability lifetime. Keep
inner.inventory_view and abi.context unchanged, and add a test that reads all
three ByteSlice fields from the constructed capability and compares them with
the expected identity strings.
---
Nitpick comments:
In `@crates/mesh-native-serving-plugin-host/src/lib.rs`:
- Line 129: Rename the ActivePlugin field from _tokenizer_capability to
tokenizer_capability and update its initializer and any references accordingly.
Add a field-level comment explaining that the raw ABI pointers provided to the
plugin remain valid only while this Arc<HostTokenizerCapability> is alive.
In `@crates/mesh-native-serving-plugin-host/src/tokenizer_capability.rs`:
- Around line 69-80: Update the SAFETY comment above the Send/Sync
implementations for HostTokenizerCapability to refer to the Arc<Self> lifetime
rather than a boxed value, and explicitly state that tokenizer and its inventory
remain immutable, non-reallocated, and alive while the plugin holds
abi.inventory because its ByteSlice values reference tokenizer-owned buffers.
Note that tokenizer drops before inventory and this is safe because dropping
Vec<TokenizerInventoryEntry> does not dereference ByteSlice.
- Around line 255-284: Add focused tests in the existing tests module covering
the validation failures in ActivationInventory::from_inventory: incorrect
schema_version, empty tokens, more than MAX_TOKENIZER_INVENTORY_ENTRIES,
non-increasing token IDs, and an empty control descriptor. Add a
write_encode_output test that verifies OUTPUT_TOO_SMALL and the required length
is populated for retry. Test every TokenizerError variant through encode_status
against its expected TokenizerEncodeStatus, and verify
HostTokenizerCapability::new exposes identity bytes equal to the bound tokenizer
identity, retaining the returned capability while inspecting the identity slice.
In `@crates/skippy-server/src/tokenizer.rs`:
- Around line 606-631: Update encode to validate pieces in one match-based pass,
rejecting controls and invalid UTF-8 without retaining the unreachable
InputPiece::Bytes else branch. Merge the validated byte contents in order and
derive the returned token IDs from that merged input instead of always calling
tokenize with an empty string, while preserving max_tokens handling and existing
error types.
In `@crates/skippy-server/src/tokenizer/binding.rs`:
- Around line 21-34: The inventory_from_stage function currently collapses
missing files, scan failures, ID conversion overflow, token_pieces errors, and
length mismatches into None. Change its return type and each early-return path
to preserve distinct TokenizerCapabilityError reasons, including propagating the
error from source.token_pieces, so the host can distinguish transient failures
from unsupported inventories.
- Around line 1-8: Add a #[cfg(test)] module in binding.rs covering
tokenizer_binding_digest determinism for identical inputs and digest changes
when token IDs, piece kinds, or piece bytes differ. Move the source_gguf_path
fallback test from tokenizer.rs into this module, preserving its existing
behavior and using the binding module’s relevant symbols.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 9b68558b-63bc-494f-90df-d0c31ff14180
📒 Files selected for processing (4)
crates/mesh-native-serving-plugin-host/src/lib.rscrates/mesh-native-serving-plugin-host/src/tokenizer_capability.rscrates/skippy-server/src/tokenizer.rscrates/skippy-server/src/tokenizer/binding.rs
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rs`:
- Around line 792-800: The existing test only validates resolver translation via
to_embedded_openai_args and does not exercise native plugin ingress. Add a
direct-serving integration test using a fake serving_hooks_factory, send a
greedy request through the loaded single-stage backend, and assert the plugin
receives propose and report callbacks with a target-authoritative verification
receipt; retain the existing NgramProposerKind::Suffix assertion as
resolver-only coverage.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 0e3812b2-ef7a-41bd-a49b-143e104445ca
📒 Files selected for processing (2)
crates/mesh-llm-host-runtime/src/inference/skippy/resolver/native_mtp_tests.rscrates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs
michaelneale
left a comment
michaelneale
left a comment
There was a problem hiding this comment.
🤖 Reviewed on micn's behalf.
Verified all CodeRabbit findings are fixed on the branch tip: the identity ByteSlice use-after-free (now owned bytes wired via Arc), byte-piece concatenation before UTF-8 validation/tokenization, consistent *output_length = 0 on all error paths, and the module splits bringing both files under 1,000 lines. CI is fully green.
The V1 → V2 plugin ABI break is acknowledged and accepted — fine to strand V1 plugins in this case.
Looks good. ✅
2 participants
Why
Native-serving plugins need access to the active model's tokenizer through a bounded, model-bound interface. Metadata alone is insufficient for structured input, and unsupported or non-lossless input must fail explicitly rather than be silently transformed.
This PR adds that capability at the Mesh/plugin boundary while keeping tokenizer preparation outside latency-sensitive proposal callbacks.
Old API (V1)
Activation provided a borrowed native tokenizer inventory:
Plugins resolved:
The inventory was metadata only. There was no model-bound encode callback, explicit encode limits, or binding digest.
New API (V2)
Activation provides a model-bound capability:
Plugins resolve:
The structured encoder accepts bounded, tagged pieces:
The contract:
Implementation
skippy-tokenizerstructured encode API and typed errors.Validation
cargo fmt --all -- --checkcargo clippy -p skippy-tokenizer -p skippy-server -p mesh-native-serving-plugin-api -p mesh-native-serving-plugin-host --all-targets -- -D warningscargo test -p skippy-tokenizer -p skippy-server -p mesh-native-serving-plugin-api -p mesh-native-serving-plugin-hostcargo check --workspaceThis is intentionally a breaking replacement for V1. V1 plugins must migrate to
mesh_native_serving_plugin_v2; backward compatibility is not provided.Summary by CodeRabbit