469 lines
17 KiB
Rust
469 lines
17 KiB
Rust
//! Shared size-bounding for MCP/text tool output.
|
|
//!
|
|
//! Large payloads (e.g. Sentry attachment base64 resources) must not land
|
|
//! fully in chat state: they inflate the token estimate and trigger premature
|
|
//! auto-compact.
|
|
//!
|
|
//! # Configurable limit
|
|
//!
|
|
//! Default [`MCP_MAX_OUTPUT_BYTES`] (20_000). Effective limit (highest first):
|
|
//!
|
|
//! 1. [`TruncationCfg`](crate::types::resources::TruncationCfg) per-tool /
|
|
//! MCP-specific (`mcp_max_output_bytes` — e.g. a winning repo-level
|
|
//! `[mcp] max_output_bytes`, seeded per session by the shell) / default,
|
|
//! when present in resources
|
|
//! 2. Host-seeded effective limit via [`set_mcp_max_output_bytes`] (host
|
|
//! resolves requirements > env > config > remote config > default once at
|
|
//! bootstrap / remote-config refresh and stores the result)
|
|
//! 3. When host has not seeded (`0`): env
|
|
//! [`ENV_GROK_MAX_MCP_OUTPUT_BYTES`] / [`ENV_MAX_MCP_OUTPUT_BYTES`]
|
|
//! 4. Built-in default
|
|
|
|
use std::path::PathBuf;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
use xai_tool_runtime::ToolCallContext;
|
|
|
|
use crate::types::output::{MCPOutputDetails, ToolOutput};
|
|
use crate::types::tool::ToolKind;
|
|
use crate::util::query_tools::{QueryTools, examples_clause};
|
|
use crate::util::truncate::format_bytes;
|
|
|
|
/// Default inline limit for MCP tool output in chat state (bytes, not tokens).
|
|
pub const MCP_MAX_OUTPUT_BYTES: usize = 20_000;
|
|
|
|
/// Env override for the MCP inline output cap (bytes).
|
|
/// Some agents use `MAX_MCP_OUTPUT_TOKENS`; we bound by **bytes** because
|
|
/// truncation is byte-oriented (`truncate_str`).
|
|
pub const ENV_MAX_MCP_OUTPUT_BYTES: &str = "MAX_MCP_OUTPUT_BYTES";
|
|
|
|
/// Grok-native env override for the MCP inline output cap (bytes).
|
|
pub const ENV_GROK_MAX_MCP_OUTPUT_BYTES: &str = "GROK_MAX_MCP_OUTPUT_BYTES";
|
|
|
|
/// Process-wide effective limit. `0` = host has not seeded; fall through to
|
|
/// env / default. The shell writes the *fully resolved* stack here so free-
|
|
/// function tool dispatch (no live `Config`) sees the same value.
|
|
static EFFECTIVE_MCP_MAX_OUTPUT_BYTES: AtomicUsize = AtomicUsize::new(0);
|
|
|
|
/// Host (shell) sets the fully-resolved MCP output cap in bytes.
|
|
///
|
|
/// Pass the already-resolved limit (requirements > env > config > remote config >
|
|
/// default). Pass `0` only in tests to clear and fall through to env / default.
|
|
pub fn set_mcp_max_output_bytes(bytes: usize) {
|
|
EFFECTIVE_MCP_MAX_OUTPUT_BYTES.store(bytes, Ordering::Relaxed);
|
|
}
|
|
|
|
/// Parse a positive byte limit from an env var. Zero / unparseable → `None`.
|
|
fn parse_positive_bytes_env(name: &str) -> Option<usize> {
|
|
let raw = std::env::var(name).ok()?;
|
|
let n = raw.trim().parse::<u64>().ok()?;
|
|
usize::try_from(n).ok().filter(|n| *n > 0)
|
|
}
|
|
|
|
/// Env tier: `GROK_MAX_MCP_OUTPUT_BYTES` then `MAX_MCP_OUTPUT_BYTES`.
|
|
///
|
|
/// Grok-native wins when both are set. Positive integers only. Used by the
|
|
/// shell resolver and as the standalone fallback when the host has not called
|
|
/// [`set_mcp_max_output_bytes`].
|
|
pub fn mcp_max_output_bytes_from_env() -> Option<usize> {
|
|
parse_positive_bytes_env(ENV_GROK_MAX_MCP_OUTPUT_BYTES)
|
|
.or_else(|| parse_positive_bytes_env(ENV_MAX_MCP_OUTPUT_BYTES))
|
|
}
|
|
|
|
/// Effective MCP inline output cap for this process.
|
|
///
|
|
/// Host-seeded value if set; otherwise env; otherwise [`MCP_MAX_OUTPUT_BYTES`].
|
|
pub fn mcp_max_output_bytes() -> usize {
|
|
match EFFECTIVE_MCP_MAX_OUTPUT_BYTES.load(Ordering::Relaxed) {
|
|
0 => mcp_max_output_bytes_from_env().unwrap_or(MCP_MAX_OUTPUT_BYTES),
|
|
n => n,
|
|
}
|
|
}
|
|
|
|
pub(crate) const LONG_LINE_BYTES: usize = 2_000;
|
|
|
|
/// How a truncated MCP payload is saved and described to the model.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub(crate) enum McpDumpKind {
|
|
LongLineJson,
|
|
Json,
|
|
LongLineText,
|
|
Other,
|
|
}
|
|
|
|
impl McpDumpKind {
|
|
pub(crate) fn classify(text: &str) -> Self {
|
|
let trimmed = text.trim();
|
|
let is_json = matches!(trimmed.as_bytes().first(), Some(b'{' | b'['))
|
|
&& serde_json::from_str::<serde::de::IgnoredAny>(trimmed).is_ok();
|
|
let has_long_line = text.lines().map(str::len).max().unwrap_or(0) > LONG_LINE_BYTES;
|
|
match (is_json, has_long_line) {
|
|
(true, true) => Self::LongLineJson,
|
|
(true, false) => Self::Json,
|
|
(false, true) => Self::LongLineText,
|
|
(false, false) => Self::Other,
|
|
}
|
|
}
|
|
|
|
pub(crate) fn extension(self) -> &'static str {
|
|
match self {
|
|
Self::LongLineJson | Self::Json => "json",
|
|
Self::LongLineText | Self::Other => "txt",
|
|
}
|
|
}
|
|
|
|
pub(crate) fn steer(self, shell: &str, tools: QueryTools) -> String {
|
|
match self {
|
|
Self::LongLineJson => format!(
|
|
" The full output is valid JSON with a very long line, so \
|
|
grep/read_file are ineffective on it — use `{shell}` to query the \
|
|
saved file{eg}.",
|
|
eg = examples_clause(&tools.json_tools()),
|
|
),
|
|
Self::Json => format!(
|
|
" The full output is valid JSON saved to the file above; use \
|
|
`{shell}` to query it{eg}.",
|
|
eg = examples_clause(&tools.json_tools()),
|
|
),
|
|
Self::LongLineText => format!(
|
|
" The full output has a very long line, so grep/read_file are \
|
|
ineffective on it — use `{shell}` to slice/search the saved \
|
|
file{eg}.",
|
|
eg = examples_clause(&tools.text_tools()),
|
|
),
|
|
Self::Other => String::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Resolved settings for truncating one MCP payload (inline limit, dump dir,
|
|
/// shell tool name, call id). Build with [`McpTruncateContext::from_tool_ctx`].
|
|
pub struct McpTruncateContext {
|
|
pub(crate) max_output_bytes: usize,
|
|
pub(crate) session_folder: Option<PathBuf>,
|
|
pub(crate) shell_tool: String,
|
|
pub(crate) call_id: String,
|
|
}
|
|
|
|
impl McpTruncateContext {
|
|
pub async fn from_tool_ctx(ctx: &ToolCallContext, tool_key: &str) -> Self {
|
|
let call_id = ctx.call_id.as_str().to_string();
|
|
let resolved_default = mcp_max_output_bytes();
|
|
match crate::types::tool_metadata::shared_resources(ctx) {
|
|
Ok(res) => {
|
|
let guard = res.lock().await;
|
|
let max_output_bytes = guard
|
|
.get::<crate::types::resources::TruncationCfg>()
|
|
.map(|cfg| cfg.0.mcp_max_output_bytes_for(tool_key, resolved_default))
|
|
.unwrap_or(resolved_default);
|
|
let session_folder = guard
|
|
.get::<crate::types::resources::SessionFolder>()
|
|
.map(|f| f.0.clone());
|
|
let shell_tool = guard
|
|
.get::<crate::types::template_renderer::TemplateRenderer>()
|
|
.and_then(|r| r.tool_for_kind(ToolKind::Execute))
|
|
.map(str::to_string)
|
|
.unwrap_or_else(|| "bash".to_string());
|
|
Self {
|
|
max_output_bytes,
|
|
session_folder,
|
|
shell_tool,
|
|
call_id,
|
|
}
|
|
}
|
|
Err(_) => Self {
|
|
max_output_bytes: resolved_default,
|
|
session_folder: None,
|
|
shell_tool: "bash".to_string(),
|
|
call_id,
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Map a `call_id` to safe filename chars so a `/` or `..` in a wire-supplied
|
|
/// id (only validated as non-empty) cannot escape the session `mcp/` dir.
|
|
fn sanitized_stem(call_id: &str) -> String {
|
|
call_id
|
|
.chars()
|
|
.map(|c| {
|
|
if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
|
|
c
|
|
} else {
|
|
'_'
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Truncate `text` in place when over the limit, dumping the full payload to
|
|
/// the session `mcp/` dir (when available) with a pointer appended.
|
|
async fn truncate_mcp_text(text: &mut String, trunc_ctx: &McpTruncateContext) {
|
|
if text.len() <= trunc_ctx.max_output_bytes {
|
|
return;
|
|
}
|
|
|
|
let total_bytes = text.len();
|
|
let kind = McpDumpKind::classify(text.as_str());
|
|
|
|
let output_file_path = trunc_ctx.session_folder.as_ref().map(|folder| {
|
|
folder.join("mcp").join(format!(
|
|
"{}.{}",
|
|
sanitized_stem(&trunc_ctx.call_id),
|
|
kind.extension()
|
|
))
|
|
});
|
|
|
|
let file_hint = if let Some(ref path) = output_file_path {
|
|
if let Some(parent) = path.parent() {
|
|
let _ = tokio::fs::create_dir_all(parent).await;
|
|
}
|
|
match tokio::fs::write(path, text.as_bytes()).await {
|
|
Ok(()) => format!(" Full output written to: {}.", path.to_string_lossy()),
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
path = %path.display(),
|
|
error = %e,
|
|
"Failed to write full MCP output to file"
|
|
);
|
|
String::new()
|
|
}
|
|
}
|
|
} else {
|
|
String::new()
|
|
};
|
|
|
|
let truncated =
|
|
crate::util::truncate::truncate_str(text.as_str(), trunc_ctx.max_output_bytes).to_owned();
|
|
let steer = if file_hint.is_empty() {
|
|
String::new()
|
|
} else {
|
|
kind.steer(&trunc_ctx.shell_tool, QueryTools::detect())
|
|
};
|
|
*text = format!(
|
|
"{}\n\n[MCP output truncated: showing first {} of {}.{}{}]",
|
|
truncated,
|
|
format_bytes(trunc_ctx.max_output_bytes),
|
|
format_bytes(total_bytes),
|
|
file_hint,
|
|
steer,
|
|
);
|
|
}
|
|
|
|
/// Bound the `MCP`/`Text` variants to the inline size limit, keeping a preview
|
|
/// and dumping the full payload to disk. Other variants are returned untouched.
|
|
pub async fn truncate_tool_output(
|
|
mut output: ToolOutput,
|
|
trunc_ctx: &McpTruncateContext,
|
|
) -> ToolOutput {
|
|
match &mut output {
|
|
ToolOutput::MCP(mcp) => {
|
|
let text = match mcp.output_mut() {
|
|
MCPOutputDetails::OkayOutput(t) | MCPOutputDetails::Error(t) => t,
|
|
};
|
|
truncate_mcp_text(text, trunc_ctx).await;
|
|
}
|
|
ToolOutput::Text(text_out) => {
|
|
truncate_mcp_text(&mut text_out.text, trunc_ctx).await;
|
|
}
|
|
_ => {}
|
|
}
|
|
output
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn cfg_with_folder(folder: PathBuf, max: usize) -> McpTruncateContext {
|
|
McpTruncateContext {
|
|
max_output_bytes: max,
|
|
session_folder: Some(folder),
|
|
shell_tool: "bash".to_string(),
|
|
call_id: "call-test".to_string(),
|
|
}
|
|
}
|
|
|
|
/// Serialize tests that mutate the process-global effective limit / env.
|
|
fn with_mcp_limit_lock<R>(f: impl FnOnce() -> R) -> R {
|
|
static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
|
|
let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
|
|
f()
|
|
}
|
|
|
|
#[test]
|
|
fn host_set_overrides_env_fallback() {
|
|
with_mcp_limit_lock(|| {
|
|
let prev = EFFECTIVE_MCP_MAX_OUTPUT_BYTES.load(Ordering::Relaxed);
|
|
// Clear host seed; with no env, effective limit is the built-in default.
|
|
set_mcp_max_output_bytes(0);
|
|
let prev_max = std::env::var(ENV_MAX_MCP_OUTPUT_BYTES).ok();
|
|
let prev_grok = std::env::var(ENV_GROK_MAX_MCP_OUTPUT_BYTES).ok();
|
|
unsafe {
|
|
std::env::remove_var(ENV_MAX_MCP_OUTPUT_BYTES);
|
|
std::env::remove_var(ENV_GROK_MAX_MCP_OUTPUT_BYTES);
|
|
}
|
|
assert_eq!(
|
|
mcp_max_output_bytes(),
|
|
MCP_MAX_OUTPUT_BYTES,
|
|
"unset host + unset env → built-in default"
|
|
);
|
|
|
|
set_mcp_max_output_bytes(10_000);
|
|
assert_eq!(mcp_max_output_bytes(), 10_000, "host seed wins over env");
|
|
|
|
set_mcp_max_output_bytes(0);
|
|
assert_eq!(
|
|
mcp_max_output_bytes(),
|
|
MCP_MAX_OUTPUT_BYTES,
|
|
"cleared host falls through to default"
|
|
);
|
|
|
|
unsafe {
|
|
match prev_max {
|
|
Some(v) => std::env::set_var(ENV_MAX_MCP_OUTPUT_BYTES, v),
|
|
None => std::env::remove_var(ENV_MAX_MCP_OUTPUT_BYTES),
|
|
}
|
|
match prev_grok {
|
|
Some(v) => std::env::set_var(ENV_GROK_MAX_MCP_OUTPUT_BYTES, v),
|
|
None => std::env::remove_var(ENV_GROK_MAX_MCP_OUTPUT_BYTES),
|
|
}
|
|
}
|
|
set_mcp_max_output_bytes(prev);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn env_parser_rejects_zero_and_junk() {
|
|
with_mcp_limit_lock(|| {
|
|
let prev_max = std::env::var(ENV_MAX_MCP_OUTPUT_BYTES).ok();
|
|
let prev_grok = std::env::var(ENV_GROK_MAX_MCP_OUTPUT_BYTES).ok();
|
|
unsafe {
|
|
std::env::remove_var(ENV_MAX_MCP_OUTPUT_BYTES);
|
|
std::env::remove_var(ENV_GROK_MAX_MCP_OUTPUT_BYTES);
|
|
}
|
|
assert_eq!(mcp_max_output_bytes_from_env(), None);
|
|
|
|
unsafe { std::env::set_var(ENV_MAX_MCP_OUTPUT_BYTES, "0") };
|
|
assert_eq!(mcp_max_output_bytes_from_env(), None);
|
|
|
|
unsafe { std::env::set_var(ENV_MAX_MCP_OUTPUT_BYTES, "not-a-number") };
|
|
assert_eq!(mcp_max_output_bytes_from_env(), None);
|
|
|
|
unsafe { std::env::set_var(ENV_MAX_MCP_OUTPUT_BYTES, "12345") };
|
|
assert_eq!(mcp_max_output_bytes_from_env(), Some(12_345));
|
|
|
|
// GROK_* wins over MAX_* when both set.
|
|
unsafe { std::env::set_var(ENV_GROK_MAX_MCP_OUTPUT_BYTES, "99999") };
|
|
assert_eq!(mcp_max_output_bytes_from_env(), Some(99_999));
|
|
|
|
unsafe { std::env::remove_var(ENV_GROK_MAX_MCP_OUTPUT_BYTES) };
|
|
assert_eq!(mcp_max_output_bytes_from_env(), Some(12_345));
|
|
|
|
unsafe {
|
|
match prev_max {
|
|
Some(v) => std::env::set_var(ENV_MAX_MCP_OUTPUT_BYTES, v),
|
|
None => std::env::remove_var(ENV_MAX_MCP_OUTPUT_BYTES),
|
|
}
|
|
match prev_grok {
|
|
Some(v) => std::env::set_var(ENV_GROK_MAX_MCP_OUTPUT_BYTES, v),
|
|
None => std::env::remove_var(ENV_GROK_MAX_MCP_OUTPUT_BYTES),
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn text_over_limit_truncates_and_dumps_full_payload() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let cfg = cfg_with_folder(dir.path().to_path_buf(), 100);
|
|
let full = "x".repeat(5_000);
|
|
|
|
let out = truncate_tool_output(ToolOutput::Text(full.clone().into()), &cfg).await;
|
|
|
|
let ToolOutput::Text(t) = out else {
|
|
panic!("expected Text");
|
|
};
|
|
assert!(t.text.len() < full.len());
|
|
assert!(t.text.starts_with(&"x".repeat(100)), "preview prefix kept");
|
|
assert!(t.text.contains("[MCP output truncated:"));
|
|
assert!(t.text.contains("Full output written to:"));
|
|
|
|
let dump = dir.path().join("mcp").join("call-test.txt");
|
|
assert_eq!(tokio::fs::read_to_string(&dump).await.unwrap(), full);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn boundary_exact_limit_untouched_one_over_truncates() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let cfg = cfg_with_folder(dir.path().to_path_buf(), 100);
|
|
|
|
let at = truncate_tool_output(ToolOutput::Text("a".repeat(100).into()), &cfg).await;
|
|
let ToolOutput::Text(t) = at else {
|
|
panic!("expected Text")
|
|
};
|
|
assert_eq!(t.text, "a".repeat(100), "exactly at limit is untouched");
|
|
assert!(
|
|
!dir.path().join("mcp").exists(),
|
|
"no dump when not truncated"
|
|
);
|
|
|
|
let over = truncate_tool_output(ToolOutput::Text("b".repeat(101).into()), &cfg).await;
|
|
let ToolOutput::Text(t) = over else {
|
|
panic!("expected Text")
|
|
};
|
|
assert!(
|
|
t.text.contains("[MCP output truncated:"),
|
|
"one over truncates"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn traversal_in_call_id_cannot_escape_session_dir() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let cfg = McpTruncateContext {
|
|
max_output_bytes: 100,
|
|
session_folder: Some(dir.path().to_path_buf()),
|
|
shell_tool: "bash".to_string(),
|
|
call_id: "../../evil".to_string(),
|
|
};
|
|
|
|
let out = truncate_tool_output(ToolOutput::Text("x".repeat(5_000).into()), &cfg).await;
|
|
|
|
let ToolOutput::Text(t) = out else {
|
|
panic!("expected Text");
|
|
};
|
|
let mcp_dir = dir.path().join("mcp");
|
|
assert!(!t.text.contains(".."), "no traversal sequence in pointer");
|
|
let entries: Vec<_> = std::fs::read_dir(&mcp_dir)
|
|
.unwrap()
|
|
.map(|e| e.unwrap().path())
|
|
.collect();
|
|
assert_eq!(entries.len(), 1, "exactly one dump file");
|
|
assert!(entries[0].starts_with(&mcp_dir), "dump stayed inside mcp/");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn non_text_variant_passes_through() {
|
|
let cfg = McpTruncateContext {
|
|
max_output_bytes: 1,
|
|
session_folder: None,
|
|
shell_tool: "bash".to_string(),
|
|
call_id: "call-test".to_string(),
|
|
};
|
|
|
|
let out = truncate_tool_output(
|
|
ToolOutput::SearchTool(crate::types::output::SearchToolOutput {
|
|
result_count: 1,
|
|
content: "anything".to_string(),
|
|
}),
|
|
&cfg,
|
|
)
|
|
.await;
|
|
|
|
let ToolOutput::SearchTool(s) = out else {
|
|
panic!("expected SearchTool");
|
|
};
|
|
assert_eq!(s.content, "anything", "passthrough leaves content intact");
|
|
}
|
|
}
|