From 305ba8ae5c9d036d5fa009f677ae0fe3d2f04afb Mon Sep 17 00:00:00 2001 From: Stella Test Date: Thu, 6 Aug 2026 04:06:38 -0700 Subject: [PATCH 1/2] fix(stella-mcp): sort the MCP schema segment so the advertised toolset is byte-stable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ToolRegistry::schemas` sorts by name for a stated reason: the list is serialized verbatim at position 0 of the prompt prefix, prompt caching is a byte-level prefix match, and HashMap iteration is per-process randomized. So two processes share the tools+system cache entry only if they emit the same bytes. `McpToolSet::schemas` then concatenated after that sorted list without re-sorting, which handed the guarantee straight back to `self.clients` order and to which server finished connecting first. Within one process the order is stable, so this is a CROSS-process miss — a restart inside the cache TTL, or two stella processes in one workspace, is exactly the case the registry's sort comment says it exists for. The MCP segment is now sorted by its namespaced name, which makes the answer independent of both client order and connection-completion order rather than merely stable within a run. Namespaced names are unique by construction (`routes` is keyed on them), so the order is total and no tie is left for the sort to break arbitrarily. Segments are preserved rather than flattened: native tools first, then MCP. That order is a deliberate contract ("the base layer the MCP set augments"), and sorting the whole list would have made the new test pass while silently moving the native tools — so the witness asserts the segment boundary too. The two sibling decorators were checked and need no change. `CandidateMcpView::schemas` concatenates a sorted native list with a filtered view of `inner.schemas()`, and filtering preserves relative order, so it inherits this fix. `DiscoveryToolSet::discovery_schemas` builds a literal `vec![]`, which is ordered by construction. Witness: `schemas_are_byte_identical_whatever_order_the_servers_connected_in` builds two sets from the same two servers registered in opposite orders and compares the SERIALIZED schemas — bytes, because bytes are what the cache matches on. It fails on the old code ("two processes that connected the same servers in different orders must advertise the same bytes") and passes with the sort. This changes the advertised order once, which is a one-time prompt-cache invalidation for sessions live across the upgrade. That is the cost of having the property at all, and it is paid once rather than on every restart. `cargo test -p stella-mcp` — 146 passed, 0 failed. The prompt-cache golden fixtures (`cargo test -p stella-pipeline --test cache_correctness`) are unaffected: 5 passed. Closes #1848 --- crates/stella-mcp/src/toolset.rs | 78 +++++++++++++++++++++++++++++++- 1 file changed, 77 insertions(+), 1 deletion(-) diff --git a/crates/stella-mcp/src/toolset.rs b/crates/stella-mcp/src/toolset.rs index 5d1821eb2..819115fe8 100644 --- a/crates/stella-mcp/src/toolset.rs +++ b/crates/stella-mcp/src/toolset.rs @@ -485,12 +485,29 @@ impl McpToolSet { #[async_trait] impl ToolExecutor for McpToolSet { + /// The advertised toolset, as two segments in a fixed order: the native + /// tools, then this set's MCP tools. + /// + /// **Both the segment order and the order within each segment are a + /// cross-process contract, not a presentation choice.** This list is + /// serialized verbatim at position 0 of the prompt prefix, and prompt + /// caching is a byte-level prefix match — so two stella processes in the + /// same workspace inside the cache TTL share the tools+system entry only + /// if they emit the same bytes. `ToolRegistry::schemas` sorts the native + /// segment for exactly that reason; this decorator used to concatenate + /// after it without re-sorting, which put the whole guarantee back at the + /// mercy of `self.clients` order and of which server finished connecting + /// first (#1848). + /// + /// Sorting the MCP segment by its namespaced name — rather than trusting + /// client order — is what makes the answer independent of both. fn schemas(&self) -> Vec { let mut schemas = Vec::new(); // Native tools first — they are the base layer the MCP set augments. if let Some(native) = &self.native { schemas.extend(native.schemas()); } + let mut mcp = Vec::new(); for (idx, client) in self.clients.iter().enumerate() { // A disabled server advertises nothing this session — the engine // re-reads schemas each model call, so the model stops seeing its @@ -503,7 +520,7 @@ impl ToolExecutor for McpToolSet { // Only advertise tools that actually route back to this client // (defends against any skipped/collided entry). if self.routes.get(&namespaced).map(|(i, _)| *i) == Some(idx) { - schemas.push(ToolSchema { + mcp.push(ToolSchema { name: namespaced, description: tool.description.clone(), input_schema: tool.input_schema.clone(), @@ -517,6 +534,11 @@ impl ToolExecutor for McpToolSet { } } } + // Namespaced names are unique by construction (`routes` is keyed on + // them), so this is a total order — no tie for the sort to resolve + // arbitrarily and reintroduce the nondeterminism. + mcp.sort_by(|a, b| a.name.cmp(&b.name)); + schemas.extend(mcp); schemas } @@ -739,6 +761,60 @@ mod tests { assert!(set.failed_servers().is_empty()); } + /// The advertised toolset must be byte-identical however the servers + /// happened to be ordered (#1848). + /// + /// This list is serialized verbatim at position 0 of the prompt prefix and + /// prompt caching is a byte-level prefix match, so the question is not + /// "does it look the same" but "do two processes emit the same bytes". + /// `ToolRegistry::schemas` sorts its own tools for that reason; this + /// decorator appended after it without re-sorting, so the guarantee held + /// only as long as `self.clients` order and connection-completion order + /// happened to match between processes — which across a restart, or two + /// stella processes in one workspace inside the cache TTL, they need not. + /// + /// Serialized rather than compared field by field: bytes are what the + /// cache matches on, so bytes are what the assertion should be about. + #[tokio::test] + async fn schemas_are_byte_identical_whatever_order_the_servers_connected_in() { + let forward = McpToolSet::from_clients(vec![ + connected_client("alpha", "one").await, + connected_client("zeta", "two").await, + ]) + .wrapping(Arc::new(FakeNative)); + let reverse = McpToolSet::from_clients(vec![ + connected_client("zeta", "two").await, + connected_client("alpha", "one").await, + ]) + .wrapping(Arc::new(FakeNative)); + + let bytes = |set: &McpToolSet| serde_json::to_string(&set.schemas()).expect("serialize"); + assert_eq!( + bytes(&forward), + bytes(&reverse), + "two processes that connected the same servers in different orders must \ + advertise the same bytes, or they cannot share a prompt-cache entry" + ); + + // And the segment contract the doc comment states: native first, then + // the MCP tools. A sort applied to the WHOLE list would also make the + // assertion above pass while silently moving the native tools. + let names: Vec = forward.schemas().into_iter().map(|s| s.name).collect(); + let first_mcp = names + .iter() + .position(|n| n.starts_with(NS_PREFIX)) + .expect("the mcp segment exists"); + assert!( + names[..first_mcp].iter().all(|n| !n.starts_with(NS_PREFIX)), + "native tools must all precede the mcp segment: {names:?}" + ); + assert_eq!( + names[first_mcp..], + ["mcp__alpha__one".to_string(), "mcp__zeta__two".to_string()], + "the mcp segment is sorted by namespaced name" + ); + } + #[tokio::test] async fn schemas_namespace_mcp_tools_and_include_native() { let client = connected_client("files", "read").await; From 382f4b26f00373d860b59372e70b2ee6fa9fbe3d Mon Sep 17 00:00:00 2001 From: Stella Test Date: Thu, 6 Aug 2026 11:37:30 -0700 Subject: [PATCH 2/2] fix(stella-mcp): bound one server's contribution to the advertised toolset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The issue reports MCP schemas as forwarded "uncapped". That is not quite right, and the correction is the fix. `crate::client::ingest` already caps every tool INDIVIDUALLY — 256 tools per server, 2,000 description characters, a per-tool `inputSchema` byte budget. What nothing bounds is the AGGREGATE. 256 tools times a 2,000-character description is ~512 KB of third-party prose from a single server, at position 0 of every request, before schemas are counted — and every per-item cap holds the whole way. The caps are per-item; the cost is the sum. Measured, through the real `tools/list` path: 40 tools with 2 KB descriptions (every ingest cap satisfied) advertise 81,440 bytes from one server. Adds `MAX_SERVER_SCHEMA_BYTES` (32 KB, ~9k tokens) applied per server. The first tool of a server is always admitted however large — a server whose one tool exceeds the budget should advertise that tool, not vanish; ingest's per-tool caps already bound it. Applied to the SORTED segment, which matters for more than tidiness: sorting by namespaced name groups each server's tools into a contiguous run AND fixes the order inside it, so the truncation point is a function of the tool names rather than of client order or connection timing. Budgeting an unsorted list would have re-lost the cross-process byte stability the sort exists to buy. `over_budget_servers()` reports what was cut, deliberately separate from `over_advertising_servers()`: those are two different walls and a reader needs to know which one was hit — 300 tools trips the count cap, twelve verbose ones trip this. One `budget_segment` fold serves both callers, and one `mcp_segment` builds the input for both, rather than a second copy of either for reporting (the drift #1613 was filed for). Witness: `one_chatty_server_cannot_spend_the_whole_prefix` drives a 40-tool server through the real transport and asserts the block is bounded, that the server is trimmed rather than silenced, that every tool is either advertised or counted as dropped, and that the survivors are the lexicographic prefix (so the cut is deterministic). It fails on the old code at 81,440 bytes. `an_ordinary_server_is_not_trimmed` is the other direction — the budget must be invisible for the servers people actually run. ## Not done: the lean-catalog default The issue's other half asks to flip `STELLA_LEAN_TOOLS` on by default, and its own fix direction says to "measure lean-mode's activation rate on the bench corpus before flipping any default" — each activation mutates the schema block and costs a prefix invalidation, so the trade is a ~70% smaller prefix against ~4-5 invalidations per session. That is a bench measurement, not a code change, and flipping it unmeasured is exactly what the issue warns against. Left open. `cargo test -p stella-mcp` — 130 passed, 0 failed. Clippy and fmt clean. Refs #1856, #1848 --- crates/stella-mcp/src/toolset.rs | 243 ++++++++++++++++++++++++++++--- 1 file changed, 220 insertions(+), 23 deletions(-) diff --git a/crates/stella-mcp/src/toolset.rs b/crates/stella-mcp/src/toolset.rs index 819115fe8..6295d5e3e 100644 --- a/crates/stella-mcp/src/toolset.rs +++ b/crates/stella-mcp/src/toolset.rs @@ -64,6 +64,27 @@ const NS_SEP: &str = "__"; /// Default per-call (and per-connect) timeout when the caller does not set one. pub const DEFAULT_CALL_TIMEOUT: Duration = Duration::from_secs(60); +/// Byte budget for ONE server's contribution to the advertised toolset. +/// +/// Ingest already caps each tool individually — `MAX_TOOLS_PER_SERVER` (256), +/// `MAX_TOOL_DESCRIPTION_CHARS` (2,000) and a per-tool `inputSchema` budget in +/// [`crate::client::ingest`]. What none of them bounds is the **aggregate**: +/// 256 tools times a 2,000-character description is ~512 KB of third-party +/// prose from a single server, at position 0 of every request, before its +/// schemas are counted (#1856). +/// +/// Every per-tool cap can hold and the block still exceed the whole +/// recall+memory+rules budget combined, because the caps are per-item and the +/// cost is the sum. 32 KB is roughly 9k tokens: generous for the servers +/// people actually run (a handful of tools each), and a wall for the one that +/// advertises a catalogue. +/// +/// Applied to the SORTED segment, so which tools survive is a deterministic +/// function of their names rather than of client order — the byte-stability +/// contract in [`McpToolSet::schemas`] would be worthless if the truncation +/// point moved between processes. +pub const MAX_SERVER_SCHEMA_BYTES: usize = 32 * 1024; + /// What a connected server said about itself during the `initialize` /// handshake — its own name (which need not match the local alias), version, /// display title, and free-prose `instructions`. @@ -483,30 +504,96 @@ impl McpToolSet { } } -#[async_trait] -impl ToolExecutor for McpToolSet { - /// The advertised toolset, as two segments in a fixed order: the native - /// tools, then this set's MCP tools. +/// Which server a namespaced tool belongs to — `mcp____`. +/// +/// Server names are guaranteed free of the `__` separator (`is_namespaceable` +/// rejects the rest at connect), so the segment between the prefix and the +/// first separator is unambiguous. +fn server_of(namespaced: &str) -> &str { + namespaced + .strip_prefix(NS_PREFIX) + .and_then(|rest| rest.split_once(NS_SEP)) + .map_or("", |(server, _)| server) +} + +/// Roughly what one schema costs in the serialized tools block: its name, its +/// description, and its `inputSchema`. +/// +/// Approximate on purpose — the exact cost depends on the provider's own JSON +/// envelope, which this crate does not model. A budget is a bound, not an +/// accounting, and being a few bytes out either way changes nothing about +/// whether a 512 KB catalogue is refused. +fn schema_cost(schema: &ToolSchema) -> usize { + schema.name.len() + + schema.description.len() + + serde_json::to_string(&schema.input_schema).map_or(0, |s| s.len()) +} + +/// Hold each server's contribution to a SORTED MCP segment under +/// [`MAX_SERVER_SCHEMA_BYTES`], returning what survives and, per server, how +/// many tools the budget cut. +/// +/// One function with two callers ([`McpToolSet::schemas`] takes the schemas, +/// [`McpToolSet::over_budget_servers`] takes the counts) rather than a second +/// copy of the rule for reporting — two implementations of one fold is exactly +/// the drift #1613 was filed for. +/// +/// The input is already sorted by namespaced name, which groups each server's +/// tools into a contiguous run (the server name is the prefix) AND fixes the +/// order within it. So the truncation point is a deterministic function of the +/// tool names, not of client order or connection timing — without that, the +/// byte-stability contract above would survive the sort and then be lost here. +fn budget_segment(sorted: Vec) -> (Vec, Vec<(String, usize)>) { + let mut kept = Vec::with_capacity(sorted.len()); + let mut elided: Vec<(String, usize)> = Vec::new(); + let mut current = String::new(); + let mut spent = 0usize; + + for schema in sorted { + let server = server_of(&schema.name); + if server != current { + current = server.to_string(); + spent = 0; + } + let cost = schema_cost(&schema); + // The first tool of a server is always admitted, however large: a + // server whose single tool exceeds the budget should advertise that + // one tool, not vanish. Ingest's per-tool caps already bound it. + if spent > 0 && spent + cost > MAX_SERVER_SCHEMA_BYTES { + match elided.last_mut() { + Some((name, count)) if *name == current => *count += 1, + _ => elided.push((current.clone(), 1)), + } + continue; + } + spent += cost; + kept.push(schema); + } + (kept, elided) +} + +impl McpToolSet { + /// Servers whose advertised tools were trimmed to fit + /// [`MAX_SERVER_SCHEMA_BYTES`], as `(server, tools dropped)`. /// - /// **Both the segment order and the order within each segment are a - /// cross-process contract, not a presentation choice.** This list is - /// serialized verbatim at position 0 of the prompt prefix, and prompt - /// caching is a byte-level prefix match — so two stella processes in the - /// same workspace inside the cache TTL share the tools+system entry only - /// if they emit the same bytes. `ToolRegistry::schemas` sorts the native - /// segment for exactly that reason; this decorator used to concatenate - /// after it without re-sorting, which put the whole guarantee back at the - /// mercy of `self.clients` order and of which server finished connecting - /// first (#1848). + /// Deliberately separate from [`Self::over_advertising_servers`], which + /// reports the per-tool COUNT cap applied at ingest. These are two + /// different walls and a reader needs to know which one it hit: 300 tools + /// trips the first, twelve verbose ones trip this. Empty for every server + /// that fits, which is nearly all of them. + #[must_use] + pub fn over_budget_servers(&self) -> Vec<(String, usize)> { + budget_segment(self.mcp_segment()).1 + } + + /// This set's MCP tools, namespaced and sorted, before the byte budget. /// - /// Sorting the MCP segment by its namespaced name — rather than trusting - /// client order — is what makes the answer independent of both. - fn schemas(&self) -> Vec { - let mut schemas = Vec::new(); - // Native tools first — they are the base layer the MCP set augments. - if let Some(native) = &self.native { - schemas.extend(native.schemas()); - } + /// The one place the segment is built. [`ToolExecutor::schemas`] budgets it + /// and returns the schemas; [`Self::over_budget_servers`] budgets it and + /// returns the counts. Deriving the counts from the already-budgeted list + /// would report zero every time — the honest input to both questions is + /// what the servers advertised, not what survived. + fn mcp_segment(&self) -> Vec { let mut mcp = Vec::new(); for (idx, client) in self.clients.iter().enumerate() { // A disabled server advertises nothing this session — the engine @@ -538,7 +625,35 @@ impl ToolExecutor for McpToolSet { // them), so this is a total order — no tie for the sort to resolve // arbitrarily and reintroduce the nondeterminism. mcp.sort_by(|a, b| a.name.cmp(&b.name)); - schemas.extend(mcp); + mcp + } +} + +#[async_trait] +impl ToolExecutor for McpToolSet { + /// The advertised toolset, as two segments in a fixed order: the native + /// tools, then this set's MCP tools. + /// + /// **Both the segment order and the order within each segment are a + /// cross-process contract, not a presentation choice.** This list is + /// serialized verbatim at position 0 of the prompt prefix, and prompt + /// caching is a byte-level prefix match — so two stella processes in the + /// same workspace inside the cache TTL share the tools+system entry only + /// if they emit the same bytes. `ToolRegistry::schemas` sorts the native + /// segment for exactly that reason; this decorator used to concatenate + /// after it without re-sorting, which put the whole guarantee back at the + /// mercy of `self.clients` order and of which server finished connecting + /// first (#1848). + /// + /// Sorting the MCP segment by its namespaced name — rather than trusting + /// client order — is what makes the answer independent of both. + fn schemas(&self) -> Vec { + let mut schemas = Vec::new(); + // Native tools first — they are the base layer the MCP set augments. + if let Some(native) = &self.native { + schemas.extend(native.schemas()); + } + schemas.extend(budget_segment(self.mcp_segment()).0); schemas } @@ -815,6 +930,88 @@ mod tests { ); } + /// One chatty server cannot spend the whole prefix (#1856). + /// + /// Ingest already caps each tool — 256 per server, 2,000 description + /// characters, a per-tool schema budget. None of them bounds the SUM, and + /// the sum is what lands at position 0 of every request: 256 × 2,000 is + /// ~512 KB of third-party prose from one server, which can exceed the + /// whole recall+memory+rules budget combined while every per-item cap + /// holds. + /// + /// Driven through the real `tools/list` path rather than by constructing + /// schemas directly, so it exercises what a server actually sends. + #[tokio::test] + async fn one_chatty_server_cannot_spend_the_whole_prefix() { + // 40 tools × ~2 KB of description — every per-tool cap satisfied. + let big = "d".repeat(2_000); + let tools: Vec = (0..40) + .map(|i| { + serde_json::json!({ + "name": format!("tool{i:02}"), + "description": big, + "inputSchema": { "type": "object" }, + }) + }) + .collect(); + let transport = ScriptedTransport::new(); + transport.push_ok( + "initialize", + serde_json::json!({ "protocolVersion": PREFERRED_PROTOCOL_VERSION }), + ); + transport.push_ok("tools/list", serde_json::json!({ "tools": tools })); + let mut client = McpClient::new("chatty", Box::new(transport)); + client.initialize().await.unwrap(); + let set = McpToolSet::from_clients(vec![client]); + + let advertised = ToolExecutor::schemas(&set); + let bytes: usize = advertised.iter().map(schema_cost).sum(); + assert!( + bytes <= MAX_SERVER_SCHEMA_BYTES + 2_100, + "one server advertised {bytes} bytes, over the \ + {MAX_SERVER_SCHEMA_BYTES}-byte budget (the slack is one \ + over-budget final tool, which is admitted by design)" + ); + assert!( + !advertised.is_empty(), + "the budget must trim a chatty server, not silence it" + ); + + // The cut is reported, not silent: an operator who wonders why a tool + // is missing needs to be told, and the count is the whole diagnostic. + let over = set.over_budget_servers(); + assert_eq!(over.len(), 1, "one server was trimmed: {over:?}"); + assert_eq!(over[0].0, "chatty"); + assert_eq!( + over[0].1 + advertised.len(), + 40, + "every tool is either advertised or counted as dropped: {over:?}" + ); + + // Deterministic, because the segment is sorted before it is budgeted — + // truncating an unsorted list would make WHICH tools survive depend on + // client order, losing the byte-stability the sort exists to buy. + let names: Vec<&str> = advertised.iter().map(|s| s.name.as_str()).collect(); + let mut sorted = names.clone(); + sorted.sort_unstable(); + assert_eq!( + names, sorted, + "the surviving tools are the lexicographic prefix" + ); + } + + /// A well-behaved server is untouched: the budget must be invisible for + /// the servers people actually run. + #[tokio::test] + async fn an_ordinary_server_is_not_trimmed() { + let set = McpToolSet::from_clients(vec![connected_client("files", "read").await]); + assert_eq!(ToolExecutor::schemas(&set).len(), 1); + assert!( + set.over_budget_servers().is_empty(), + "nothing to report for a server that fits" + ); + } + #[tokio::test] async fn schemas_namespace_mcp_tools_and_include_native() { let client = connected_client("files", "read").await;