[Rust TreeCore] Harden runtime and CI parity

.github/actions/download-rust-ext/action.yml

@@ -53,7 +53,7 @@ runs:
         path: |
           python/sglang/srt/rust_extensions/_*.so
           python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
-        key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
+        key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'proto/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
 
     # Job-wide, but only setup.py reads it, and only while building.
     # Whether the modules suit this interpreter is not decided here:

.github/workflows/_pr-test-rust-ext-build.yml

@@ -88,10 +88,12 @@ jobs:
           ref: ${{ inputs.git_ref || github.sha }}
           # Just what the cache key hashes, plus the action and script this job
           # runs: the workspace is cold here and the rest of the tree is mostly
-          # docs. Both jobs must hash the same rust/** set, which this preserves.
+          # docs. Both jobs must hash the same Rust extension inputs, which this
+          # preserves.
           # Cone mode off is what allows naming a single file.
           sparse-checkout: |
             rust
+            proto
             python/setup.py
             python/pyproject.toml
             python/sglang/srt/rust_extensions/torch_build.py
@@ -115,7 +117,7 @@ jobs:
           path: |
             python/sglang/srt/rust_extensions/_*.so
             python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
-          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
+          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'proto/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
 
       # On a miss: different hash = rust/setup.py moved; no entries = evicted.
       - name: Report cache lookup
@@ -324,6 +326,7 @@ jobs:
           # both have to check out the same set for hashFiles to agree.
           sparse-checkout: |
             rust
+            proto
             python/setup.py
             python/pyproject.toml
             python/sglang/srt/rust_extensions/torch_build.py
@@ -359,7 +362,7 @@ jobs:
           path: |
             python/sglang/srt/rust_extensions/_*.so
             python/sglang/srt/mem_cache/rust_tree_core/mem_cache*.so
-          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
+          key: ${{ inputs.cache_key_prefix }}-${{ hashFiles('rust/**', 'proto/**', 'python/setup.py', 'python/pyproject.toml', 'python/sglang/srt/rust_extensions/torch_build.py', '.github/workflows/_pr-test-rust-ext-build.yml', 'scripts/ci/utils/stage_rust_ext_modules.sh') }}
 
       - name: Upload extension modules
         uses: actions/upload-artifact@v4

.github/workflows/seed-rust-ext-cache.yml

@@ -9,6 +9,7 @@ on:
     branches: [main]
     paths:
       - 'rust/**'
+      - 'proto/**'
       - 'python/setup.py'
       - 'python/pyproject.toml'
       - 'python/sglang/srt/rust_extensions/torch_build.py'

python/sglang/srt/rust_extensions/loader.py

@@ -55,6 +55,7 @@ class _CrateSpec:
     manifest: Path
     workspace: Path
     features: tuple[str, ...]
+    source_inputs: tuple[Path, ...]
 
 
 @dataclass(frozen=True)
@@ -80,7 +81,8 @@ def load_rust_extension(
     The crate is discovered from the workspace under ``rust/``: the one whose
     Cargo manifest declares ``[package.metadata.sglang] python-module`` equal
     to ``python_module`` (the same metadata setup.py uses for wheel builds), so
-    new crates need no registration here.
+    new crates need no registration here. Crates may declare ``source-inputs``
+    relative to their manifest for build inputs outside the Rust workspace.
 
     ``auto`` prefers a module bundled in an installed wheel. In a source tree,
     it ignores unverified in-package artifacts and uses the fingerprinted cache
@@ -152,9 +154,12 @@ def load_rust_extension(
             features=features,
             build_environment=build_environment,
         )
-        if _source_digest(crate.workspace) != context.source_digest:
+        if (
+            _source_digest(crate.workspace, crate.source_inputs)
+            != context.source_digest
+        ):
             raise RuntimeError(
-                f"Rust sources under {crate.workspace} changed during the build; "
+                f"Rust extension sources for {crate.package} changed during the build; "
                 "the result was not cached"
             )
         _stage_atomically(artifact, extension_path)
@@ -216,6 +221,10 @@ def _discover_crate(workspace: Path, python_module: str) -> _CrateSpec:
                 manifest=manifest,
                 workspace=crate_workspace,
                 features=tuple(sglang_metadata.get("features", ())),
+                source_inputs=tuple(
+                    (manifest.parent / path).resolve()
+                    for path in sglang_metadata.get("source-inputs", ())
+                ),
             )
         )
 
@@ -245,7 +254,7 @@ def _build_context(
         features = crate.features
     if extension_module is None:
         extension_module = crate.python_module
-    source_digest = _source_digest(crate.workspace)
+    source_digest = _source_digest(crate.workspace, crate.source_inputs)
     toolchain = {
         "cargo": _command_version(
             "cargo", "--version", "--verbose", cwd=crate.workspace
@@ -289,10 +298,16 @@ def _build_context(
     )
 
 
-def _source_digest(workspace: Path) -> str:
+def _source_digest(workspace: Path, source_inputs: tuple[Path, ...] = ()) -> str:
     digest = hashlib.sha256()
-    for path in _source_files(workspace):
-        relative_path = path.relative_to(workspace).as_posix().encode()
+    paths = set(_source_files(workspace))
+    for source_input in source_inputs:
+        if source_input.is_dir():
+            paths.update(_source_files(source_input))
+        else:
+            paths.add(source_input)
+    for path in sorted(paths, key=lambda item: os.path.relpath(item, workspace)):
+        relative_path = Path(os.path.relpath(path, workspace)).as_posix().encode()
         digest.update(len(relative_path).to_bytes(8, "big"))
         digest.update(relative_path)
         if path.is_symlink():

rust/sglang-grpc/Cargo.toml

@@ -9,6 +9,9 @@ license.workspace = true
 # of the main sglang wheel at the given import path.
 [package.metadata.sglang]
 python-module = "sglang.srt.rust_extensions._grpc"
+# build.rs compiles schemas from outside the Rust workspace; include them in
+# the source-loader artifact fingerprint.
+source-inputs = ["../../proto"]
 # Always build optimized, even for an editable install.
 debug = false
 

rust/sglang-radix-tree/src/components/full.rs

@@ -49,16 +49,17 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
         &self,
         tree_core: &UnifiedTreeCore<K>,
         mut result: MatchResult,
+        last_device_node_idx: NodeIdx_,
+        best_match_node_idx: NodeIdx_,
         params: &MatchPrefixParams<'_, K>,
         value_chunks: &[Tensor],
         best_value_len: usize,
     ) -> MatchResult {
         // Compute Full KV host hit length: walk from last_host_node up to
         // last_device_node, summing host_value lengths of evicted nodes.
         let mut kv_host_hit = 0;
-        let mut node_idx = tree_core.arena.resolve(result.best_match_node_id);
-        let last_device_idx = tree_core.arena.resolve(result.last_device_node_id);
-        while node_idx != last_device_idx {
+        let mut node_idx = best_match_node_idx;
+        while node_idx != last_device_node_idx {
             let node = tree_core.arena.node(node_idx);
             let parent = node.try_parent().unwrap_or_else(|| {
                 panic!(
@@ -469,7 +470,10 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
                 {
                     let mut offset = 0i64;
                     for &loaded_id in transfer.nodes_to_load.iter().flatten() {
-                        let loaded_idx = tree_core.arena.resolve(loaded_id);
+                        let loaded_idx = tree_core
+                            .arena
+                            .resolve(loaded_id)
+                            .expect("load-back transfers must reference live nodes");
                         let loaded = tree_core.arena.node_mut(loaded_idx);
                         let n_len = loaded.host_value_len(FULL) as i64;
                         loaded

rust/sglang-radix-tree/src/components/mamba.rs

@@ -127,6 +127,8 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
         &self,
         tree_core: &UnifiedTreeCore<K>,
         mut result: MatchResult,
+        _last_device_node_idx: NodeIdx_,
+        best_match_node_idx: NodeIdx_,
         _params: &MatchPrefixParams<'_, K>,
         _value_chunks: &[Tensor],
         _best_value_len: usize,
@@ -143,9 +145,7 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
 
         // HiCache: if mamba was evicted from device but has host backup,
         // ensure mamba_host_hit_length >= 1 so load_back is triggered.
-        let last_node = tree_core
-            .arena
-            .node(tree_core.arena.resolve(result.best_match_node_id));
+        let last_node = tree_core.arena.node(best_match_node_idx);
         if !last_node.has_device_value(MAMBA) && last_node.has_host_value(MAMBA) {
             result.mamba_host_hit_length = result.mamba_host_hit_length.max(1);
         }
@@ -643,7 +643,12 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
                 let target_node_id = insert_result
                     .as_deref()
                     .and_then(|result| result.inserted_host_node)
-                    .map(|id| tree_core.arena.resolve(id));
+                    .map(|id| {
+                        tree_core
+                            .arena
+                            .resolve(id)
+                            .expect("prefetch insert results must reference live nodes")
+                    });
                 let attach_target = match (host_indices, target_node_id) {
                     (Some(_), Some(target))
                         if loaded && !tree_core.arena.has_host_value(target, MAMBA) =>

rust/sglang-radix-tree/src/components/mod.rs

@@ -130,6 +130,8 @@ pub trait TreeComponent<K: ChildKeyType> {
         &self,
         tree_core: &UnifiedTreeCore<K>,
         result: MatchResult,
+        _last_device_node_idx: NodeIdx_,
+        _best_match_node_idx: NodeIdx_,
         params: &MatchPrefixParams<'_, K>,
         value_chunks: &[Tensor],
         best_value_len: usize,

rust/sglang-radix-tree/src/components/swa.rs

@@ -225,14 +225,28 @@ impl SwaComponent {
         });
         let target = insert_result
             .and_then(|result| result.inserted_host_node)
-            .map(|id| tree_core.arena.resolve(id));
+            .map(|id| {
+                tree_core
+                    .arena
+                    .resolve(id)
+                    .expect("prefetch insert results must reference live nodes")
+            });
 
         let (Some(target), Some(host_indices)) = (target, transfer.host_indices.as_ref()) else {
             if let Some(host_indices) = &transfer.host_indices {
                 self.release_swa_host_(host_indices.shallow_clone(), cache_actions);
             }
             return;
         };
+        // Cache-mode graft commit only (buffer fills never reach here):
+        // a hit-shrunk window mid-tree is missing its head, so drop it.
+        // Root anchors are complete windows of their own.
+        if node_id != tree_core.arena.root()
+            && window_require_pages < self.sliding_window_size.div_ceil(page_size)
+        {
+            self.release_swa_host_(host_indices.shallow_clone(), cache_actions);
+            return;
+        }
         if window_require_pages == 0 || loaded_pages < window_require_pages {
             self.release_swa_host_(host_indices.shallow_clone(), cache_actions);
             return;
@@ -338,6 +352,8 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
         &self,
         tree_core: &UnifiedTreeCore<K>,
         mut result: MatchResult,
+        _last_device_node_idx: NodeIdx_,
+        best_match_node_idx: NodeIdx_,
         params: &MatchPrefixParams<'_, K>,
         value_chunks: &[Tensor],
         best_value_len: usize,
@@ -347,9 +363,7 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
         // toward the SWA host hit.
         let mut n_swa = 0;
         let mut swa_host_hit = 0;
-        let mut node = tree_core
-            .arena
-            .node(tree_core.arena.resolve(result.best_match_node_id));
+        let mut node = tree_core.arena.node(best_match_node_idx);
         while !node.is_root() && n_swa < self.sliding_window_size {
             if node.has_device_value(SWA) {
                 n_swa += node.device_value_len(SWA);
@@ -965,7 +979,10 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
                 let mut swa_chunks: Vec<Tensor> = Vec::new();
                 let mut offset = 0i64;
                 for &loaded_id in transfer.nodes_to_load.iter().flatten() {
-                    let loaded_idx = tree_core.arena.resolve(loaded_id);
+                    let loaded_idx = tree_core
+                        .arena
+                        .resolve(loaded_id)
+                        .expect("load-back transfers must reference live nodes");
                     let n_tokens = tree_core.arena.host_value_len(loaded_idx, SWA) as i64;
                     let swa_chunk = device_indices.narrow(0, offset, n_tokens).copy();
                     tree_core.set_component_device_value_(

rust/sglang-radix-tree/src/node.rs

@@ -688,14 +688,21 @@ pub struct ValueState {
 
 // Tree-core runtime errors.
 
+/// A public node handle does not name a live arena node.
+#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
+#[error("node {node_id} is not allocated")]
+pub struct NodeAccessError {
+    pub node_id: NodeId,
+}
+
 /// Errors surfaced from the tree-core runtime API when a caller violates a documented
 /// contract (freeing an unallocated node, allocating under a freed parent).
 #[allow(clippy::enum_variant_names)]
 #[derive(Debug, thiserror::Error)]
 pub enum TreeCoreRuntimeError {
     /// A public NodeId no longer names a live arena node.
-    #[error("node {node_id} is not allocated")]
-    NodeNotAllocated { node_id: NodeId },
+    #[error(transparent)]
+    NodeAccess(#[from] NodeAccessError),
     /// `begin_insert`/`insert` called while a resumable insert is suspended.
     #[error("concurrent insert walks")]
     ConcurrentInsertWalk,
@@ -736,6 +743,10 @@ pub enum TreeCoreRuntimeError {
     /// A host insert below a non-root anchor must remain in that anchor's namespace.
     #[error("insert_host namespace does not match non-root anchor {node_id}")]
     InsertHostNamespaceMismatch { node_id: NodeId },
+    /// An inspection-only invariant check failed without mutating the tree.
+    #[cfg(any(test, feature = "inspection"))]
+    #[error("{0}")]
+    InspectionAssertion(String),
 }
 
 // Unigram and bigram child keys.
@@ -1058,18 +1069,12 @@ impl<K: ChildKeyType> NodeArena<K> {
         self.root = self.alloc_root();
     }
 
-    /// The live slot for an external handle; panics on a freed or unknown id.
-    #[track_caller]
-    pub fn resolve(&self, id: NodeId) -> NodeIdx_ {
-        *self
-            .id_map
+    /// The live slot for an external handle.
+    pub fn resolve(&self, id: NodeId) -> Result<NodeIdx_, NodeAccessError> {
+        self.id_map
             .get(&id)
-            .unwrap_or_else(|| panic!("node {id} is not allocated"))
-    }
-
-    /// The live slot for an external handle, or None if freed/unknown.
-    pub fn try_resolve(&self, id: NodeId) -> Option<NodeIdx_> {
-        self.id_map.get(&id).copied()
+            .copied()
+            .ok_or(NodeAccessError { node_id: id })
     }
 
     /// Mint the next external handle for the slot and index it.