feat(skippy): expose the loaded tokenizer capability - #1147
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (18)
🚧 Files skipped from review as they are similar to previous changes (14)
📝 WalkthroughWalkthroughChangesThe PR adds a tokenizer protocol and runtime-backed Tokenizer serving
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant instrumented_openai_router
participant tokenizer_http_router
participant TokenizerCapability
participant LoadedRuntime
Client->>instrumented_openai_router: POST /v1/tokenize
instrumented_openai_router->>tokenizer_http_router: Route tokenizer request
tokenizer_http_router->>TokenizerCapability: Validate identity and limits
TokenizerCapability->>LoadedRuntime: Tokenize text
LoadedRuntime-->>TokenizerCapability: Token IDs and token pieces
TokenizerCapability-->>Client: TokenizeResponse or structured error
Possibly related PRs
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 |
|
This pull request is currently a draft. Reviews will not take place until the PR is marked as ready for review. |
4a6203c to
b089857
Compare
b089857 to
85244ff
Compare
85244ff to
1dcef38
Compare
|
@coderabbitai full review |
✅ Action performedFull review finished. |
![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 (5)
crates/skippy-server/src/lib.rs (1)
18-18: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe crate-root re-export duplicates an already-public path.
Line 18 declares
pub mod tokenizer. Line 41 re-exports three of its items at the crate root. Consumers can then reachTokenizerCapabilitythrough two paths, and neither is canonical.Pick one. Either keep
pub mod tokenizerand drop Line 41, or make the module private and keep the re-export as the single path, matchingcrates/skippy-protocol/src/lib.rs.As per coding guidelines: "Minimize crate-root re-exports. Temporary compatibility re-exports are allowed during refactors, but new code should import from the owning module directly."
Also applies to: 41-41
🤖 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/lib.rs` at line 18, Remove the crate-root re-export at the line containing the three tokenizer items, and keep `pub mod tokenizer` as the canonical public path. Update any internal imports or consumers that rely on the crate-root names to import them from `tokenizer` instead.Source: Coding guidelines
crates/mesh-llm-host-runtime/src/network/openai/request_parse.rs (2)
173-185: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueAlign the emptiness check with the server-side identity check.
Line 179 filters on
!model_id.is_empty().tokenizer_identity_from_stageincrates/skippy-server/src/tokenizer.rsat Line 63 usesmodel_id.trim().is_empty().A whitespace-only
model_idpasses here and becomes the routing key. Routing then fails as model-unavailable rather than as a malformed request. The outcome is safe, but the two layers disagree on what an empty identity is.Use
!model_id.trim().is_empty()so both layers reject the same inputs.♻️ Proposed change
- .filter(|model_id| !model_id.is_empty()) + .filter(|model_id| !model_id.trim().is_empty())🤖 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-llm-host-runtime/src/network/openai/request_parse.rs` around lines 173 - 185, Update the model_id validation in the tokenize branch of the request parsing flow to use trimmed emptiness checking, matching tokenizer_identity_from_stage. Ensure whitespace-only expected_identity.model_id values are rejected as malformed requests while preserving the existing non-empty model handling.
969-1014: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a negative test for the missing-identity rejection.
Line 180 introduces a new fail-closed path. A
/v1/tokenizerequest without a non-emptyexpected_identity.model_idmakesread_http_request_with_limitsreturn an error, and the caller answers 400. No test covers it.
read_request_from_partsunwraps, so the test needs a variant that returns theResult. Add one and assert the error for a body that omitsexpected_identity, and for a body wheremodel_idis an empty string.🤖 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-llm-host-runtime/src/network/openai/request_parse.rs` around lines 969 - 1014, Add a Result-returning test helper alongside read_request_from_parts, then add negative async tests for /v1/tokenize requests whose body omits expected_identity or provides an empty expected_identity.model_id. Assert that read_http_request_with_limits returns an error in both cases, while preserving the existing successful routing tests.crates/skippy-server/src/tokenizer.rs (1)
359-376: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winAdd a
source_model_sha256mismatch case.The test varies only
model_id. The source-file hash is the strongest part of the identity contract and the stated reason the capability fails closed. Add a case that changesexpected_identity.source_model_sha256while keepingmodel_idequal, and assertIdentityMismatch.Also consider a case for a non-hex
source_model_sha256passed totokenizer_identity_from_stage, which should yieldIdentityUnavailablethrough theis_sha256filter at Line 67.💚 Proposed additional assertions
let (capability, _) = capability(Vec::new()); let mut request = request("x".to_string()); request.expected_identity.model_id = "another-model".to_string(); assert_eq!( capability.tokenize(request).unwrap_err(), TokenizerCapabilityError::IdentityMismatch ); + + assert_eq!( + tokenizer_identity_from_stage(0, "model", Some(&"z".repeat(64))).unwrap_err(), + TokenizerCapabilityError::IdentityUnavailable + ); + let (capability, _) = capability(Vec::new()); + let mut request = request("x".to_string()); + request.expected_identity.source_model_sha256 = "f".repeat(64); + assert_eq!( + capability.tokenize(request).unwrap_err(), + TokenizerCapabilityError::IdentityMismatch + ); }🤖 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 359 - 376, Extend the test function identity_is_authoritative_and_fail_closed with an IdentityMismatch assertion that changes only expected_identity.source_model_sha256 while retaining the matching model_id. Also add a tokenizer_identity_from_stage case using a non-hex source SHA-256 value and assert it returns TokenizerCapabilityError::IdentityUnavailable through the existing validation.crates/mesh-llm-host-runtime/src/network/openai/transport.rs (1)
274-296: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffCorrect, but the tokenizer exemption is now scattered across five call sites.
The guards at Line 275 and Line 294 are right, and they match the guards at
request_parse.rsLine 698 andingress.rsLine 606. Each one turns off a generation-only behavior for tokenizer requests.The risk is additive. Any future generation-only step must remember to add a sixth guard, and nothing in the type system enforces it. Consider classifying the request once, for example as a
RequestKind::CapabilityversusRequestKind::Generation, and branching on that classification at the top of the plan.request_context_budgetat Line 48 already demonstrates the chokepoint pattern that works well here.This is structural advice, not a defect in the current change.
🤖 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-llm-host-runtime/src/network/openai/transport.rs` around lines 274 - 296, Centralize the tokenizer-versus-generation classification instead of repeating is_tokenize_request guards across rewrite_effective_model, prepare_mesh_targets, request_parse, and ingress. Introduce or reuse a RequestKind-style classification at the planning chokepoint, following request_context_budget, and branch generation-only behavior from that single value while preserving tokenizer behavior.
🤖 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/skippy-server/src/embedded.rs`:
- Around line 371-375: Restore the public start_openai_backend function to its
legacy two-argument signature (bind_addr, backend). Add a separate
tokenizer-aware entry point for the current implementation, and have the legacy
function delegate using the existing/default tokenizer so callers such as
mesh-llm-host-runtime remain source-compatible.
In `@crates/skippy-server/src/tokenizer.rs`:
- Around line 241-244: Fix the vacuous mutation-safety test around
RecordingTokenizer::generation_mutations by removing the unwritten counter
assertion and renaming the test to match its actual behavior, or preferably
exercise LoadedStageZeroTokenizer and compare observable RuntimeState before and
after tokenization. Ensure the test verifies that tokenization does not mutate
generation state through the real runtime path.
- Around line 91-113: Update TokenizerSource and LoadedStageZeroTokenizer to add
token_pieces, acquiring the runtime lock once while detokenizing all requested
IDs. Change TokenizerCapability::tokenize to use the batched method instead of
calling token_piece per token, and run the capability call from
tokenize_entrypoint inside tokio::task::spawn_blocking so mutex contention
cannot block Tokio worker threads.
---
Nitpick comments:
In `@crates/mesh-llm-host-runtime/src/network/openai/request_parse.rs`:
- Around line 173-185: Update the model_id validation in the tokenize branch of
the request parsing flow to use trimmed emptiness checking, matching
tokenizer_identity_from_stage. Ensure whitespace-only expected_identity.model_id
values are rejected as malformed requests while preserving the existing
non-empty model handling.
- Around line 969-1014: Add a Result-returning test helper alongside
read_request_from_parts, then add negative async tests for /v1/tokenize requests
whose body omits expected_identity or provides an empty
expected_identity.model_id. Assert that read_http_request_with_limits returns an
error in both cases, while preserving the existing successful routing tests.
In `@crates/mesh-llm-host-runtime/src/network/openai/transport.rs`:
- Around line 274-296: Centralize the tokenizer-versus-generation classification
instead of repeating is_tokenize_request guards across rewrite_effective_model,
prepare_mesh_targets, request_parse, and ingress. Introduce or reuse a
RequestKind-style classification at the planning chokepoint, following
request_context_budget, and branch generation-only behavior from that single
value while preserving tokenizer behavior.
In `@crates/skippy-server/src/lib.rs`:
- Line 18: Remove the crate-root re-export at the line containing the three
tokenizer items, and keep `pub mod tokenizer` as the canonical public path.
Update any internal imports or consumers that rely on the crate-root names to
import them from `tokenizer` instead.
In `@crates/skippy-server/src/tokenizer.rs`:
- Around line 359-376: Extend the test function
identity_is_authoritative_and_fail_closed with an IdentityMismatch assertion
that changes only expected_identity.source_model_sha256 while retaining the
matching model_id. Also add a tokenizer_identity_from_stage case using a non-hex
source SHA-256 value and assert it returns
TokenizerCapabilityError::IdentityUnavailable through the existing validation.
🪄 Autofix (Beta)
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: 01720c7b-122f-4e61-a537-03eb560969f3
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (16)
crates/mesh-llm-host-runtime/src/inference/skippy/mod.rscrates/mesh-llm-host-runtime/src/network/openai/ingress.rscrates/mesh-llm-host-runtime/src/network/openai/request_parse.rscrates/mesh-llm-host-runtime/src/network/openai/transport.rscrates/mesh-llm-host-runtime/src/runtime/local.rscrates/mesh-llm-host-runtime/src/runtime/local_split/loading.rscrates/skippy-protocol/Cargo.tomlcrates/skippy-protocol/src/lib.rscrates/skippy-protocol/src/tokenizer.rscrates/skippy-server/Cargo.tomlcrates/skippy-server/src/embedded.rscrates/skippy-server/src/frontend/generation/server.rscrates/skippy-server/src/http.rscrates/skippy-server/src/lib.rscrates/skippy-server/src/tokenizer.rsscripts/skippy-ci-smoke.sh
michaelneale
left a comment
michaelneale
left a comment
There was a problem hiding this comment.
makes sense - clippy failures to fix up before merge.
This will lock up generation (which I guess is ok as it is like generation)
f2cc0bc to
ac25145
Compare
ac25145 to
4fb088b
Compare
2cba0cb to
ddd4287
Compare
3 participants
Use case
An application embedding the local OpenAI server may need to count tokens, split documents on the model's actual token boundaries, or inspect token pieces. A benchmark runner may also need exact token counts from the serving tokenizer so throughput and latency measurements use the same tokenization as generation.
Loading a second tokenizer for that work wastes memory and can silently diverge from the tokenizer used by the model.
Change
The stage-zero server exposes a bounded, read-only tokenizer capability backed by the tokenizer already loaded with the model. Callers must provide the expected model and source-file identity, and mismatches fail closed.
The capability is local to the serving boundary and is not exposed over remote stage transport. Reusing the loaded tokenizer guarantees that token accounting and generation use the same vocabulary without introducing another model load.
Validation
Skippy Server tests pass (353/353), including tokenizer identity, input/output bounds, token-piece alignment, and generation-state mutation safety.
Summary by CodeRabbit
New Features
/v1/tokenizesupport with token IDs and optional token pieces.Bug Fixes