Merge branch 'main' of github.com:Dark-Alex-17/coyote
This commit is contained in:
@@ -141,6 +141,7 @@ arboard = { version = "3.3.0", default-features = false }
|
||||
|
||||
[dev-dependencies]
|
||||
pretty_assertions = "1.4.0"
|
||||
rmcp = { version = "3.1.2", features = ["server"] }
|
||||
serial_test = "3"
|
||||
|
||||
[[bin]]
|
||||
|
||||
@@ -35,7 +35,9 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g
|
||||
* [Create Custom TypeScript Tools](https://github.com/Dark-Alex-17/coyote/wiki/Custom-Tools#custom-typescript-based-tools)
|
||||
* [Create Custom Bash Tools](https://github.com/Dark-Alex-17/coyote/wiki/Custom-Bash-Tools)
|
||||
* [Bash Prompt Utilities](https://github.com/Dark-Alex-17/coyote/wiki/Bash-Prompt-Helpers)
|
||||
* [First-Class MCP Server Support](https://github.com/Dark-Alex-17/coyote/wiki/MCP-Servers): Easily connect and interact with MCP servers for advanced functionality.
|
||||
* [First-Class MCP Server Support](https://github.com/Dark-Alex-17/coyote/wiki/MCP-Servers): Easily connect and interact with MCP servers for advanced functionality. Coyote supports all three MCP capabilities: tools, resources, and prompts.
|
||||
* Models interact with each server through a compact set of capability-gated meta-tools: `mcp_search`/`mcp_describe` for discovery across tools, resources, and prompts, `mcp_invoke` for tool calls, `mcp_read` for paged and regex-filterable resource reads, and `mcp_prompt` for server-defined prompts. Binary content is spilled to disk instead of inlined, and oversized tool results are bounded before they reach the model.
|
||||
* Invoke server prompts yourself with `.prompt <server> <name> [key=value ...]` in the REPL, with live staged tab-completion (servers, then prompt names, then `key=` arguments), and discover them with `.list prompts`.
|
||||
* [Macros](https://github.com/Dark-Alex-17/coyote/wiki/Macros): Automate repetitive tasks and workflows with Coyote "scripts" (macros). Macros are Coyote's custom commands: invoke any macro directly by name (e.g. `.review main`), with tab-completion, right alongside the built-in REPL commands.
|
||||
* Give a macro a `description` (shown in `.list macros` and completions) and set `isolated: false` to run its steps on the live session, exactly as if you typed them. Note that non-isolated steps are recorded in the session, and mutating steps (`.role`, `.model`, ...) persist after the macro ends — by design. Steps are fail-fast: an error aborts the remaining steps, but completed steps' effects remain. A non-isolated macro step cannot invoke another macro, and a `.exit` step never exits the REPL.
|
||||
* Commit project-specific macros to `.coyote/macros/` in your repo — they shadow same-named global macros (opt out with `--no-workspace-macros`).
|
||||
|
||||
+6
-6
@@ -47,7 +47,7 @@ pub enum McpScopeArg {
|
||||
.args(["sandbox", "fresh"])
|
||||
.multiple(true)
|
||||
.conflicts_with_all([
|
||||
"model", "prompt", "role", "session", "agent", "rag", "rebuild_rag",
|
||||
"model", "temp_role", "role", "session", "agent", "rag", "rebuild_rag",
|
||||
"macro_name", "execute", "code", "file", "no_stream", "no_memory",
|
||||
"init_memory", "dry_run", "info", "build_tools", "install",
|
||||
"install_builtins", "sync_models", "list_models", "list_roles",
|
||||
@@ -70,9 +70,9 @@ pub struct Cli {
|
||||
/// Select a LLM model
|
||||
#[arg(short, long, add = ArgValueCompleter::new(model_completer))]
|
||||
pub model: Option<String>,
|
||||
/// Use the system prompt
|
||||
/// Set a temporary role (an ad-hoc system prompt) for this invocation
|
||||
#[arg(long)]
|
||||
pub prompt: Option<String>,
|
||||
pub temp_role: Option<String>,
|
||||
/// Select a role
|
||||
#[arg(short, long, add = ArgValueCompleter::new(role_completer))]
|
||||
pub role: Option<String>,
|
||||
@@ -705,9 +705,9 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_prompt_flag() {
|
||||
let cli = parse(&["--prompt", "be a pirate"]);
|
||||
assert_eq!(cli.prompt, Some("be a pirate".to_string()));
|
||||
fn parse_temp_role_flag() {
|
||||
let cli = parse(&["--temp-role", "be a pirate"]);
|
||||
assert_eq!(cli.temp_role, Some("be a pirate".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
+2
-1
@@ -15,6 +15,7 @@ use crate::config::prompts::{
|
||||
};
|
||||
use crate::graph::types::RagNode;
|
||||
use crate::graph::{Graph, GraphParser, NodeType};
|
||||
use crate::mcp::McpServerFeatures;
|
||||
use crate::rag::RagInitConfig;
|
||||
use crate::vault::SECRET_RE;
|
||||
use anyhow::{Context, Result};
|
||||
@@ -380,7 +381,7 @@ impl Agent {
|
||||
self.graph_rags.get(node_id).cloned()
|
||||
}
|
||||
|
||||
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) {
|
||||
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<McpServerFeatures>) {
|
||||
self.functions.append_mcp_meta_functions(mcp_servers);
|
||||
}
|
||||
|
||||
|
||||
@@ -70,7 +70,7 @@ impl AppState {
|
||||
|
||||
let mut functions = Functions::init(config.visible_tools.as_ref().unwrap_or(&Vec::new()))?;
|
||||
if !mcp_registry.is_empty() && config.mcp_server_support {
|
||||
functions.append_mcp_meta_functions(mcp_registry.list_started_servers());
|
||||
functions.append_mcp_meta_functions(mcp_registry.server_features());
|
||||
}
|
||||
|
||||
let mcp_registry = if mcp_registry.is_empty() {
|
||||
|
||||
+66
-2
@@ -442,7 +442,9 @@ impl BundleStore {
|
||||
self.ensure_bundle_exists(bundle)?;
|
||||
for (name, record) in self.bundles.iter_mut() {
|
||||
if name != bundle {
|
||||
record.files.retain(|owned| owned.path != file.path);
|
||||
record
|
||||
.files
|
||||
.retain(|owned| !same_installed_path(&owned.path, &file.path));
|
||||
}
|
||||
}
|
||||
let record = self
|
||||
@@ -450,7 +452,9 @@ impl BundleStore {
|
||||
.get_mut(bundle)
|
||||
.expect("bundle existence checked above");
|
||||
|
||||
record.files.retain(|owned| owned.path != file.path);
|
||||
record
|
||||
.files
|
||||
.retain(|owned| !same_installed_path(&owned.path, &file.path));
|
||||
record.files.push(file);
|
||||
|
||||
self.save()
|
||||
@@ -577,6 +581,17 @@ pub(crate) struct BundleListRow {
|
||||
pub(crate) drift: DriftSummary,
|
||||
}
|
||||
|
||||
/// NTFS and default APFS resolve file names case-insensitively, so records
|
||||
/// differing only in case denote the same physical file there. Linux keeps
|
||||
/// exact matching because case variants are genuinely distinct files.
|
||||
fn same_installed_path(a: &str, b: &str) -> bool {
|
||||
if cfg!(any(windows, target_os = "macos")) {
|
||||
a.eq_ignore_ascii_case(b)
|
||||
} else {
|
||||
a == b
|
||||
}
|
||||
}
|
||||
|
||||
/// An unreadable file counts as locally modified: it exists but its integrity
|
||||
/// cannot be verified.
|
||||
pub(crate) fn bundle_list_rows(store: &BundleStore, config_dir: &Path) -> Vec<BundleListRow> {
|
||||
@@ -1489,4 +1504,53 @@ mod tests {
|
||||
|
||||
assert!(rows.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn same_installed_path_matches_filesystem_case_semantics() {
|
||||
assert!(same_installed_path("macros/a.yaml", "macros/a.yaml"));
|
||||
assert!(!same_installed_path("macros/a.yaml", "macros/b.yaml"));
|
||||
assert_eq!(
|
||||
same_installed_path("macros/Foo.yaml", "macros/foo.yaml"),
|
||||
cfg!(any(windows, target_os = "macos"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_file_transfers_case_variant_ownership_on_case_insensitive_hosts() {
|
||||
let dir = TempStoreDir::new("bundles-case-variant");
|
||||
let mut store = dir.store();
|
||||
store
|
||||
.upsert_bundle("alpha", metadata("https://x/a", "aaa"))
|
||||
.unwrap();
|
||||
store
|
||||
.upsert_bundle("beta", metadata("https://x/b", "bbb"))
|
||||
.unwrap();
|
||||
store
|
||||
.record_file("alpha", file_record("macros/Shared.yaml", "one"))
|
||||
.unwrap();
|
||||
|
||||
store
|
||||
.record_file("beta", file_record("macros/shared.yaml", "two"))
|
||||
.unwrap();
|
||||
|
||||
let alpha_still_owns = store
|
||||
.get("alpha")
|
||||
.unwrap()
|
||||
.files
|
||||
.iter()
|
||||
.any(|f| f.path == "macros/Shared.yaml");
|
||||
assert_eq!(
|
||||
alpha_still_owns,
|
||||
!cfg!(any(windows, target_os = "macos")),
|
||||
"case-variant paths are one physical file on case-insensitive hosts"
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.get("beta")
|
||||
.unwrap()
|
||||
.files
|
||||
.iter()
|
||||
.any(|f| f.path == "macros/shared.yaml")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+131
-21
@@ -752,9 +752,29 @@ fn select_uninstall_candidate(store: &BundleStore, spec: &str) -> Result<Option<
|
||||
fn is_safe_relative_path(path: &str) -> bool {
|
||||
let recorded = Path::new(path);
|
||||
!recorded.is_absolute()
|
||||
&& recorded
|
||||
.components()
|
||||
.all(|c| matches!(c, Component::Normal(_)))
|
||||
&& recorded.components().all(|c| match c {
|
||||
Component::Normal(name) => is_safe_component(&name.to_string_lossy()),
|
||||
_ => false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Rejects names Windows refuses or silently rewrites (alternate data stream
|
||||
/// colons, reserved device names, trailing dots or spaces) so a recorded path
|
||||
/// denotes the same regular file on every platform.
|
||||
fn is_safe_component(name: &str) -> bool {
|
||||
!name.contains(':')
|
||||
&& !name.ends_with('.')
|
||||
&& !name.ends_with(' ')
|
||||
&& !is_windows_reserved_name(name)
|
||||
}
|
||||
|
||||
fn is_windows_reserved_name(name: &str) -> bool {
|
||||
let stem = name.split('.').next().unwrap_or("");
|
||||
let lower = stem.to_ascii_lowercase();
|
||||
matches!(lower.as_str(), "con" | "prn" | "aux" | "nul")
|
||||
|| (lower.len() == 4
|
||||
&& (lower.starts_with("com") || lower.starts_with("lpt"))
|
||||
&& matches!(lower.as_bytes()[3], b'1'..=b'9'))
|
||||
}
|
||||
|
||||
fn uninstall_owned_files(
|
||||
@@ -1067,7 +1087,12 @@ impl TempRepoDir {
|
||||
|
||||
impl Drop for TempRepoDir {
|
||||
fn drop(&mut self) {
|
||||
let _ = fs::remove_dir_all(&self.path);
|
||||
if let Err(error) = fs::remove_dir_all(&self.path) {
|
||||
log::warn!(
|
||||
"failed to remove temp clone {}: {error}",
|
||||
self.path.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1079,13 +1104,37 @@ fn is_commit_sha(reference: &str) -> bool {
|
||||
|
||||
fn clone_to_temp(url: &str, reference: Option<&str>) -> Result<TempRepoDir> {
|
||||
let dest = utils::temp_file("coyote-remote-install-", "");
|
||||
match clone_into(&dest, url, reference) {
|
||||
Ok(head_sha) => Ok(TempRepoDir {
|
||||
path: dest,
|
||||
head_sha,
|
||||
}),
|
||||
Err(error) => {
|
||||
let _ = fs::remove_dir_all(&dest);
|
||||
Err(error)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Checked-out bytes must not depend on the machine's git configuration:
|
||||
/// recorded sha256 provenance would otherwise drift with autocrlf settings.
|
||||
/// Long paths are opted into for deep bundle trees on Windows.
|
||||
fn git_content_config() -> Vec<OsString> {
|
||||
["core.autocrlf=false", "core.eol=lf", "core.longpaths=true"]
|
||||
.iter()
|
||||
.flat_map(|setting| ["-c".into(), (*setting).into()])
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn clone_into(dest: &Path, url: &str, reference: Option<&str>) -> Result<String> {
|
||||
let dest_arg: OsString = dest.as_os_str().into();
|
||||
|
||||
let is_sha = reference.is_some_and(is_commit_sha);
|
||||
|
||||
match reference {
|
||||
Some(r) if !is_sha => {
|
||||
run_git(vec![
|
||||
let mut args = git_content_config();
|
||||
args.extend([
|
||||
"clone".into(),
|
||||
"--depth".into(),
|
||||
"1".into(),
|
||||
@@ -1094,26 +1143,28 @@ fn clone_to_temp(url: &str, reference: Option<&str>) -> Result<TempRepoDir> {
|
||||
"--".into(),
|
||||
url.into(),
|
||||
dest_arg,
|
||||
])?;
|
||||
]);
|
||||
run_git(args)?;
|
||||
}
|
||||
Some(r) => {
|
||||
run_git(vec![
|
||||
"clone".into(),
|
||||
"--".into(),
|
||||
url.into(),
|
||||
dest_arg.clone(),
|
||||
])?;
|
||||
run_git(vec!["-C".into(), dest_arg, "checkout".into(), r.into()])?;
|
||||
let mut args = git_content_config();
|
||||
args.extend(["clone".into(), "--".into(), url.into(), dest_arg.clone()]);
|
||||
run_git(args)?;
|
||||
let mut args = git_content_config();
|
||||
args.extend(["-C".into(), dest_arg, "checkout".into(), r.into()]);
|
||||
run_git(args)?;
|
||||
}
|
||||
None => {
|
||||
run_git(vec![
|
||||
let mut args = git_content_config();
|
||||
args.extend([
|
||||
"clone".into(),
|
||||
"--depth".into(),
|
||||
"1".into(),
|
||||
"--".into(),
|
||||
url.into(),
|
||||
dest_arg,
|
||||
])?;
|
||||
]);
|
||||
run_git(args)?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1123,11 +1174,7 @@ fn clone_to_temp(url: &str, reference: Option<&str>) -> Result<TempRepoDir> {
|
||||
"rev-parse".into(),
|
||||
"HEAD".into(),
|
||||
])?;
|
||||
|
||||
Ok(TempRepoDir {
|
||||
path: dest,
|
||||
head_sha,
|
||||
})
|
||||
Ok(head_sha)
|
||||
}
|
||||
|
||||
fn run_git(args: Vec<OsString>) -> Result<()> {
|
||||
@@ -1816,7 +1863,17 @@ fn record_written_file(
|
||||
}
|
||||
|
||||
fn provenance_path(dst: &Path) -> String {
|
||||
let rel = dst.strip_prefix(paths::config_dir()).unwrap_or(dst);
|
||||
let rel = match dst.strip_prefix(paths::config_dir()) {
|
||||
Ok(rel) => rel,
|
||||
Err(_) => {
|
||||
log::warn!(
|
||||
"bundle file {} lies outside the config dir (an asset dir override?); \
|
||||
it will not be uninstallable and drift checks may misreport it",
|
||||
dst.display()
|
||||
);
|
||||
dst
|
||||
}
|
||||
};
|
||||
rel.to_string_lossy().replace('\\', "/")
|
||||
}
|
||||
|
||||
@@ -2242,6 +2299,53 @@ mod tests {
|
||||
use std::env;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
#[test]
|
||||
fn safe_relative_path_accepts_plain_portable_components() {
|
||||
assert!(is_safe_relative_path("macros/a.yaml"));
|
||||
assert!(is_safe_relative_path("skills/deep/nested/file.md"));
|
||||
assert!(is_safe_relative_path("functions/tools/console.sh"));
|
||||
assert!(is_safe_relative_path("roles/common.md"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn safe_relative_path_rejects_escapes_and_windows_hazards() {
|
||||
assert!(!is_safe_relative_path("../outside.yaml"));
|
||||
assert!(!is_safe_relative_path("/abs/path.yaml"));
|
||||
assert!(!is_safe_relative_path("macros/../../evil.yaml"));
|
||||
assert!(!is_safe_relative_path("macros/a.yaml:stream"));
|
||||
assert!(!is_safe_relative_path("macros/trailing."));
|
||||
assert!(!is_safe_relative_path("macros/trailing "));
|
||||
assert!(!is_safe_relative_path("macros/nul"));
|
||||
assert!(!is_safe_relative_path("macros/NUL.yaml"));
|
||||
assert!(!is_safe_relative_path("con/a.yaml"));
|
||||
assert!(!is_safe_relative_path("macros/COM1.txt"));
|
||||
assert!(!is_safe_relative_path("macros/lpt9"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn windows_reserved_name_check_is_stem_based() {
|
||||
assert!(is_windows_reserved_name("nul"));
|
||||
assert!(is_windows_reserved_name("NUL.txt"));
|
||||
assert!(is_windows_reserved_name("com1"));
|
||||
assert!(!is_windows_reserved_name("com0"));
|
||||
assert!(!is_windows_reserved_name("com10"));
|
||||
assert!(!is_windows_reserved_name("console"));
|
||||
assert!(!is_windows_reserved_name("nullable.yaml"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn git_content_config_pins_line_endings_and_long_paths() {
|
||||
let args = git_content_config();
|
||||
let rendered: Vec<String> = args
|
||||
.iter()
|
||||
.map(|a| a.to_string_lossy().into_owned())
|
||||
.collect();
|
||||
assert_eq!(rendered.len(), 6);
|
||||
assert!(rendered.contains(&"core.autocrlf=false".to_string()));
|
||||
assert!(rendered.contains(&"core.eol=lf".to_string()));
|
||||
assert!(rendered.contains(&"core.longpaths=true".to_string()));
|
||||
}
|
||||
|
||||
struct TestVaultConfigGuard {
|
||||
dir_key: String,
|
||||
file_key: String,
|
||||
@@ -4286,6 +4390,12 @@ mod tests {
|
||||
#[test]
|
||||
#[serial]
|
||||
fn uninstall_shorthand_with_multiple_matches_bails_non_interactively() {
|
||||
if *IS_STDOUT_TERMINAL {
|
||||
eprintln!(
|
||||
"Skipping uninstall_shorthand_with_multiple_matches_bails_non_interactively: requires non-TTY stdout"
|
||||
);
|
||||
return;
|
||||
}
|
||||
let _guard = TestVaultConfigGuard::new("uninst-short-multi");
|
||||
let mut store = BundleStore::load().unwrap();
|
||||
store
|
||||
|
||||
@@ -287,6 +287,7 @@ fn discover_macros_in(dirs: &[(MacroSource, PathBuf)]) -> Vec<DiscoveredMacro> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::repl;
|
||||
use crate::utils::get_env_name;
|
||||
use serial_test::serial;
|
||||
use std::path::Path;
|
||||
@@ -565,6 +566,34 @@ mod tests {
|
||||
assert_eq!(state_of(&policy, "a"), &MacroState::Enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_macro_is_shadowed_by_builtin() {
|
||||
let policy = MacroPolicy::effective_with(
|
||||
globals(&["prompt"]),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
&repl::builtin_command_names(),
|
||||
);
|
||||
|
||||
assert_eq!(state_of(&policy, "prompt"), &MacroState::ShadowedBuiltin);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn temp_role_macro_is_shadowed_by_builtin() {
|
||||
let policy = MacroPolicy::effective_with(
|
||||
globals(&["temp-role"]),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
&repl::builtin_command_names(),
|
||||
);
|
||||
|
||||
assert_eq!(state_of(&policy, "temp-role"), &MacroState::ShadowedBuiltin);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn locked_wins_over_shadowed_builtin() {
|
||||
let l = list(&["a"]);
|
||||
|
||||
@@ -22,6 +22,8 @@ pub(crate) mod todo;
|
||||
mod tool_scope;
|
||||
mod update;
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) use self::agent::AgentConfig;
|
||||
pub use self::agent::{
|
||||
Agent, AgentVariable, AgentVariables, complete_agent_variables, list_agents,
|
||||
list_agents_with_descriptions,
|
||||
@@ -51,6 +53,11 @@ pub use self::skill::Skill;
|
||||
pub use self::skill_policy::SkillPolicy;
|
||||
#[allow(unused_imports)]
|
||||
pub use self::skill_registry::SkillRegistry;
|
||||
#[cfg(test)]
|
||||
pub(crate) use self::tool_scope::test_fixtures;
|
||||
pub use self::tool_scope::{
|
||||
McpPromptCompletion, flatten_prompt_messages, resolve_prompt_args, sanitize_display_text,
|
||||
};
|
||||
pub use self::update::run_self_update;
|
||||
use crate::client::{
|
||||
self, ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS,
|
||||
|
||||
+168
-91
@@ -5,7 +5,7 @@ use super::skill::{SKILL_SCAFFOLD, Skill};
|
||||
use super::skill_policy::SkillPolicy;
|
||||
use super::skill_registry::SkillRegistry;
|
||||
use super::todo::TodoList;
|
||||
use super::tool_scope::{McpRuntime, ToolScope};
|
||||
use super::tool_scope::{McpPromptCompletion, McpRuntime, ToolScope, format_prompt_arguments};
|
||||
use super::{
|
||||
AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, AssetCategory, CREATE_TITLE_ROLE,
|
||||
Input, InstallFilter, LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, MacroAllowlistLevel,
|
||||
@@ -23,8 +23,8 @@ use crate::function::{
|
||||
user_interaction::USER_FUNCTION_PREFIX,
|
||||
};
|
||||
use crate::mcp::{
|
||||
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
||||
MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, is_auth_required_error,
|
||||
CatalogItem, MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired,
|
||||
McpServersConfig, is_auth_required_error, is_mcp_meta_function, mcp_meta_function_names,
|
||||
};
|
||||
use crate::rag::Rag;
|
||||
use crate::supervisor::Supervisor;
|
||||
@@ -77,6 +77,31 @@ fn installed_bundle_names() -> Vec<String> {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn expand_enabled_mcp_server_ids(
|
||||
app: &AppConfig,
|
||||
mcp_config: &McpServersConfig,
|
||||
enabled_mcp_servers: &[String],
|
||||
) -> Vec<String> {
|
||||
if enabled_mcp_servers.iter().any(|s| s.trim() == "all") {
|
||||
return mcp_config.mcp_servers.keys().cloned().collect();
|
||||
}
|
||||
|
||||
let mut ids = Vec::new();
|
||||
for item in enabled_mcp_servers.iter().map(|s| s.trim()) {
|
||||
if mcp_config.mcp_servers.contains_key(item) {
|
||||
ids.push(item.to_string());
|
||||
} else if let Some(mapped) = app.mapping_mcp_servers.get(item) {
|
||||
for mapped_id in mapped.split(',').map(|s| s.trim()) {
|
||||
if mcp_config.mcp_servers.contains_key(mapped_id) {
|
||||
ids.push(mapped_id.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
ids
|
||||
}
|
||||
|
||||
pub struct AutoContinueConfig {
|
||||
pub enabled: bool,
|
||||
pub max_continues: usize,
|
||||
@@ -139,6 +164,20 @@ pub(crate) fn asset_table(header: &[&str]) -> Table {
|
||||
table
|
||||
}
|
||||
|
||||
fn mcp_prompt_rows(items: &[CatalogItem]) -> Vec<[String; 4]> {
|
||||
items
|
||||
.iter()
|
||||
.map(|item| {
|
||||
[
|
||||
item.server.clone(),
|
||||
item.name.clone(),
|
||||
item.description.clone(),
|
||||
format_prompt_arguments(item.arguments.as_deref().unwrap_or_default()),
|
||||
]
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn complete_skills_with_descriptions(names: Vec<String>) -> Vec<(String, Option<String>)> {
|
||||
names
|
||||
.into_iter()
|
||||
@@ -2011,11 +2050,7 @@ impl RequestContext {
|
||||
.functions
|
||||
.declarations()
|
||||
.iter()
|
||||
.filter(|v| {
|
||||
!v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||
&& !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||
})
|
||||
.filter(|v| !is_mcp_meta_function(&v.name))
|
||||
.map(|v| v.name.to_string())
|
||||
.collect();
|
||||
|
||||
@@ -2025,11 +2060,7 @@ impl RequestContext {
|
||||
.functions()
|
||||
.declarations()
|
||||
.iter()
|
||||
.filter(|v| {
|
||||
!v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||
&& !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||
})
|
||||
.filter(|v| !is_mcp_meta_function(&v.name))
|
||||
.map(|v| v.name.to_string()),
|
||||
);
|
||||
}
|
||||
@@ -2102,11 +2133,7 @@ impl RequestContext {
|
||||
.declarations()
|
||||
.to_vec()
|
||||
.into_iter()
|
||||
.filter(|v| {
|
||||
!v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||
&& !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||
})
|
||||
.filter(|v| !is_mcp_meta_function(&v.name))
|
||||
.collect();
|
||||
|
||||
if let Some(ref tool_names) = role_filter {
|
||||
@@ -2155,11 +2182,7 @@ impl RequestContext {
|
||||
.functions
|
||||
.declarations()
|
||||
.iter()
|
||||
.filter(|v| {
|
||||
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||
|| v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||
|| v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||
})
|
||||
.filter(|v| is_mcp_meta_function(&v.name))
|
||||
.map(|v| v.name.to_string())
|
||||
.collect();
|
||||
if let Some(agent) = &self.agent {
|
||||
@@ -2168,12 +2191,7 @@ impl RequestContext {
|
||||
.functions()
|
||||
.declarations()
|
||||
.iter()
|
||||
.filter(|v| {
|
||||
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||
|| v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||
|| v.name
|
||||
.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||
})
|
||||
.filter(|v| is_mcp_meta_function(&v.name))
|
||||
.map(|v| v.name.to_string()),
|
||||
);
|
||||
}
|
||||
@@ -2188,41 +2206,17 @@ impl RequestContext {
|
||||
continue;
|
||||
}
|
||||
|
||||
let item_invoke_name =
|
||||
format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX);
|
||||
let item_search_name =
|
||||
format!("{}_{item}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX);
|
||||
let item_describe_name =
|
||||
format!("{}_{item}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX);
|
||||
if let Some(values) = app.mapping_mcp_servers.get(item) {
|
||||
server_names.extend(
|
||||
values
|
||||
.split(',')
|
||||
.flat_map(|v| {
|
||||
vec![
|
||||
format!(
|
||||
"{}_{}",
|
||||
MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
||||
v.to_string()
|
||||
),
|
||||
format!(
|
||||
"{}_{}",
|
||||
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
|
||||
v.to_string()
|
||||
),
|
||||
format!(
|
||||
"{}_{}",
|
||||
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX,
|
||||
v.to_string()
|
||||
),
|
||||
]
|
||||
})
|
||||
.flat_map(mcp_meta_function_names)
|
||||
.filter(|v| mcp_declaration_names.contains(v)),
|
||||
)
|
||||
} else if mcp_declaration_names.contains(&item_invoke_name) {
|
||||
server_names.insert(item_invoke_name);
|
||||
server_names.insert(item_search_name);
|
||||
server_names.insert(item_describe_name);
|
||||
} else if mcp_declaration_names.contains(&item_search_name) {
|
||||
server_names.extend(mcp_meta_function_names(item));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2251,11 +2245,7 @@ impl RequestContext {
|
||||
.declarations()
|
||||
.to_vec()
|
||||
.into_iter()
|
||||
.filter(|v| {
|
||||
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||
|| v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||
|| v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||
})
|
||||
.filter(|v| is_mcp_meta_function(&v.name))
|
||||
.collect();
|
||||
|
||||
if let Some(ref server_names) = role_filter {
|
||||
@@ -2351,7 +2341,7 @@ impl RequestContext {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn use_prompt(&mut self, _app: &AppConfig, prompt: &str) -> Result<()> {
|
||||
pub fn use_temp_role(&mut self, _app: &AppConfig, prompt: &str) -> Result<()> {
|
||||
let mut role = Role::new(TEMP_ROLE_NAME, prompt);
|
||||
role.set_model(self.current_model().clone());
|
||||
self.use_role_obj(role)
|
||||
@@ -2828,7 +2818,7 @@ impl RequestContext {
|
||||
}
|
||||
"bundles" => bundles::list_installed_bundles(),
|
||||
_ => bail!(
|
||||
"Unknown kind '{kind}'. Valid kinds: roles, sessions, agents, rags, macros, skills, tools, mcp-servers, bundles"
|
||||
"Unknown kind '{kind}'. Valid kinds: roles, sessions, agents, rags, macros, skills, prompts, tools, mcp-servers, bundles"
|
||||
),
|
||||
}
|
||||
}
|
||||
@@ -3396,6 +3386,7 @@ impl RequestContext {
|
||||
"rags",
|
||||
"macros",
|
||||
"skills",
|
||||
"prompts",
|
||||
"tools",
|
||||
"mcp-servers",
|
||||
"bundles",
|
||||
@@ -3720,6 +3711,40 @@ impl RequestContext {
|
||||
fuzzy_filter(values, |v| v.0.as_str(), filter)
|
||||
}
|
||||
|
||||
pub fn mcp_prompt_completion(&self, args: &[&str]) -> McpPromptCompletion {
|
||||
let app = self.app.config.as_ref();
|
||||
let enabled_ids = match &self.app.mcp_config {
|
||||
Some(mcp_config) => {
|
||||
let mut servers = self
|
||||
.enabled_mcp_servers_for_current_scope(app, true)
|
||||
.unwrap_or_default();
|
||||
servers.extend(self.skill_registry.loaded_mcp_servers());
|
||||
expand_enabled_mcp_server_ids(app, mcp_config, &servers)
|
||||
}
|
||||
None => vec![],
|
||||
};
|
||||
self.tool_scope
|
||||
.mcp_runtime
|
||||
.prompt_completion(&enabled_ids, args)
|
||||
}
|
||||
|
||||
pub async fn list_mcp_prompts(&self) -> Result<()> {
|
||||
let items = self.tool_scope.mcp_runtime.prompt_catalog().await;
|
||||
if items.is_empty() {
|
||||
println!("No prompts found.");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let mut table = asset_table(&["server", "name", "description", "args"]);
|
||||
for row in mcp_prompt_rows(&items) {
|
||||
table.add_row(row.to_vec());
|
||||
}
|
||||
|
||||
println!("Prompts:");
|
||||
println!("{table}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn rebuild_tool_scope(
|
||||
&mut self,
|
||||
app: &AppConfig,
|
||||
@@ -3763,24 +3788,7 @@ impl RequestContext {
|
||||
&& let Some(mcp_config) = &self.app.mcp_config
|
||||
{
|
||||
let server_ids: Vec<String> = match &enabled_mcp_servers {
|
||||
Some(servers) if servers.iter().any(|s| s.trim() == "all") => {
|
||||
mcp_config.mcp_servers.keys().cloned().collect()
|
||||
}
|
||||
Some(servers) => {
|
||||
let mut ids = Vec::new();
|
||||
for item in servers.iter().map(|s| s.trim()) {
|
||||
if mcp_config.mcp_servers.contains_key(item) {
|
||||
ids.push(item.to_string());
|
||||
} else if let Some(mapped) = app.mapping_mcp_servers.get(item) {
|
||||
for mapped_id in mapped.split(',').map(|s| s.trim()) {
|
||||
if mcp_config.mcp_servers.contains_key(mapped_id) {
|
||||
ids.push(mapped_id.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
ids
|
||||
}
|
||||
Some(servers) => expand_enabled_mcp_server_ids(app, mcp_config, servers),
|
||||
None => vec![],
|
||||
};
|
||||
|
||||
@@ -3836,7 +3844,7 @@ impl RequestContext {
|
||||
functions.append_todo_functions();
|
||||
}
|
||||
if !mcp_runtime.is_empty() {
|
||||
functions.append_mcp_meta_functions(mcp_runtime.server_names());
|
||||
functions.append_mcp_meta_functions(mcp_runtime.server_features());
|
||||
}
|
||||
if app.function_calling_support && policy.skills_enabled {
|
||||
functions.append_skill_functions();
|
||||
@@ -4685,10 +4693,11 @@ mod tests {
|
||||
use crate::config::AppState;
|
||||
use crate::config::agent::AgentConfig;
|
||||
use crate::function::{ToolCall, skill};
|
||||
use crate::mcp::{McpServer, McpServersConfig, McpTransportType};
|
||||
use crate::mcp::{McpServer, McpServerFeatures, McpServersConfig, McpTransportType};
|
||||
use crate::utils;
|
||||
use crate::utils::get_env_name;
|
||||
use crate::vault::Vault;
|
||||
use rmcp::model::PromptArgument;
|
||||
use serde_json::json;
|
||||
use serial_test::serial;
|
||||
use std::env;
|
||||
@@ -4746,6 +4755,15 @@ mod tests {
|
||||
RequestContext::new(default_app_state(), WorkingMode::Cmd)
|
||||
}
|
||||
|
||||
fn tools_only_features(name: &str) -> McpServerFeatures {
|
||||
McpServerFeatures {
|
||||
name: name.to_string(),
|
||||
tools: true,
|
||||
resources: false,
|
||||
prompts: false,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_creates_clean_state() {
|
||||
let ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd);
|
||||
@@ -4834,10 +4852,10 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn use_prompt_creates_temp_role() {
|
||||
fn use_temp_role_creates_temp_role() {
|
||||
let mut ctx = create_test_ctx();
|
||||
let app = ctx.app.config.clone();
|
||||
ctx.use_prompt(&app, "you are a pirate").unwrap();
|
||||
ctx.use_temp_role(&app, "you are a pirate").unwrap();
|
||||
assert!(ctx.role.is_some());
|
||||
assert_eq!(ctx.role.as_ref().unwrap().name(), "temp");
|
||||
assert!(
|
||||
@@ -4849,6 +4867,41 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_prompt_rows_assembles_columns() {
|
||||
let items = vec![
|
||||
CatalogItem {
|
||||
name: "summarize".to_string(),
|
||||
server: "docs".to_string(),
|
||||
description: "Summarize a document".to_string(),
|
||||
arguments: Some(vec![
|
||||
PromptArgument::new("path").with_required(true),
|
||||
PromptArgument::new("style"),
|
||||
]),
|
||||
..Default::default()
|
||||
},
|
||||
CatalogItem {
|
||||
name: "greet".to_string(),
|
||||
server: "misc".to_string(),
|
||||
..Default::default()
|
||||
},
|
||||
];
|
||||
|
||||
let rows = mcp_prompt_rows(&items);
|
||||
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(
|
||||
rows[0],
|
||||
[
|
||||
"docs",
|
||||
"summarize",
|
||||
"Summarize a document",
|
||||
"path (required), style"
|
||||
]
|
||||
);
|
||||
assert_eq!(rows[1], ["misc", "greet", "", ""]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_role_returns_standalone_role() {
|
||||
let mut ctx = create_test_ctx();
|
||||
@@ -5753,9 +5806,10 @@ mod tests {
|
||||
#[test]
|
||||
fn select_enabled_mcp_servers_all_returns_all_mcp_functions() {
|
||||
let mut ctx = create_test_ctx();
|
||||
ctx.tool_scope
|
||||
.functions
|
||||
.append_mcp_meta_functions(vec!["github".into(), "slack".into()]);
|
||||
ctx.tool_scope.functions.append_mcp_meta_functions(vec![
|
||||
tools_only_features("github"),
|
||||
tools_only_features("slack"),
|
||||
]);
|
||||
|
||||
let mut role = Role::new("r", "p");
|
||||
role.set_enabled_mcp_servers(Some(vec!["all".to_string()]));
|
||||
@@ -5770,9 +5824,10 @@ mod tests {
|
||||
#[test]
|
||||
fn select_enabled_mcp_servers_comma_filters() {
|
||||
let mut ctx = create_test_ctx();
|
||||
ctx.tool_scope
|
||||
.functions
|
||||
.append_mcp_meta_functions(vec!["github".into(), "slack".into()]);
|
||||
ctx.tool_scope.functions.append_mcp_meta_functions(vec![
|
||||
tools_only_features("github"),
|
||||
tools_only_features("slack"),
|
||||
]);
|
||||
|
||||
let mut role = Role::new("r", "p");
|
||||
role.set_enabled_mcp_servers(Some(vec!["github".to_string()]));
|
||||
@@ -5783,6 +5838,28 @@ mod tests {
|
||||
assert!(!names.contains(&"mcp_invoke_slack"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_enabled_mcp_servers_keeps_resources_only_server() {
|
||||
let mut ctx = create_test_ctx();
|
||||
ctx.tool_scope
|
||||
.functions
|
||||
.append_mcp_meta_functions(vec![McpServerFeatures {
|
||||
name: "res".to_string(),
|
||||
tools: false,
|
||||
resources: true,
|
||||
prompts: false,
|
||||
}]);
|
||||
|
||||
let mut role = Role::new("r", "p");
|
||||
role.set_enabled_mcp_servers(Some(vec!["res".to_string()]));
|
||||
|
||||
let fns = ctx.select_enabled_mcp_servers(&role);
|
||||
let names: Vec<&str> = fns.iter().map(|f| f.name.as_str()).collect();
|
||||
assert!(names.contains(&"mcp_search_res"));
|
||||
assert!(names.contains(&"mcp_describe_res"));
|
||||
assert!(!names.contains(&"mcp_invoke_res"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_empty_context_has_no_context_flags() {
|
||||
let ctx = create_test_ctx();
|
||||
|
||||
+1474
-20
File diff suppressed because it is too large
Load Diff
+1532
-35
File diff suppressed because it is too large
Load Diff
@@ -630,14 +630,14 @@ async fn populate_agent_mcp_runtime(ctx: &mut RequestContext, server_ids: &[Stri
|
||||
}
|
||||
|
||||
fn sync_agent_functions_to_ctx(ctx: &mut RequestContext) -> Result<()> {
|
||||
let server_names = ctx.tool_scope.mcp_runtime.server_names();
|
||||
let server_features = ctx.tool_scope.mcp_runtime.server_features();
|
||||
let functions = {
|
||||
let agent = ctx
|
||||
.agent
|
||||
.as_mut()
|
||||
.with_context(|| "Agent should be initialized")?;
|
||||
if !server_names.is_empty() {
|
||||
agent.append_mcp_meta_functions(server_names);
|
||||
if !server_features.is_empty() {
|
||||
agent.append_mcp_meta_functions(server_features);
|
||||
}
|
||||
agent.functions().clone()
|
||||
};
|
||||
@@ -1453,7 +1453,8 @@ async fn summarize_output(ctx: &RequestContext, agent_name: &str, output: &str)
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::{AppState, WorkingMode};
|
||||
use crate::config::test_fixtures::{FixtureServer, fixture_runtime};
|
||||
use crate::config::{AgentConfig, AppState, WorkingMode};
|
||||
use crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
|
||||
use serde_json::json;
|
||||
use serial_test::serial;
|
||||
@@ -1510,6 +1511,28 @@ mod tests {
|
||||
.block_on(f)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_agent_functions_gates_meta_functions_on_live_capabilities() {
|
||||
let mut ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd);
|
||||
ctx.agent = Some(Agent::test_new(AgentConfig::default()));
|
||||
let (runtime, _server) = fixture_runtime(FixtureServer {
|
||||
tools_capability: false,
|
||||
resources_capability: true,
|
||||
..FixtureServer::default()
|
||||
})
|
||||
.await;
|
||||
ctx.tool_scope.mcp_runtime = runtime;
|
||||
|
||||
sync_agent_functions_to_ctx(&mut ctx).unwrap();
|
||||
|
||||
let functions = &ctx.tool_scope.functions;
|
||||
assert_eq!(functions.declarations().len(), 3);
|
||||
assert!(functions.contains("mcp_search_fixture"));
|
||||
assert!(functions.contains("mcp_describe_fixture"));
|
||||
assert!(functions.contains("mcp_read_fixture"));
|
||||
assert!(!functions.contains("mcp_invoke_fixture"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handle_list_running_empty_supervisor() {
|
||||
let mut ctx = ctx_with_supervisor(4, 3);
|
||||
|
||||
+2
-2
@@ -380,8 +380,8 @@ async fn run(
|
||||
.await?;
|
||||
} else {
|
||||
let app: Arc<AppConfig> = Arc::clone(&ctx.app.config);
|
||||
if let Some(prompt) = &cli.prompt {
|
||||
ctx.use_prompt(app.as_ref(), prompt)?;
|
||||
if let Some(prompt) = &cli.temp_role {
|
||||
ctx.use_temp_role(app.as_ref(), prompt)?;
|
||||
} else if let Some(name) = &cli.role {
|
||||
ctx.use_role(app.as_ref(), name, abort_signal.clone())
|
||||
.await?;
|
||||
|
||||
+151
-47
@@ -1,6 +1,7 @@
|
||||
mod auth_client;
|
||||
pub(crate) mod manage;
|
||||
pub(crate) mod oauth;
|
||||
pub(crate) mod render;
|
||||
mod sse_transport;
|
||||
|
||||
use crate::config::AppConfig;
|
||||
@@ -15,6 +16,7 @@ use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use http::{HeaderName, HeaderValue};
|
||||
use indexmap::IndexMap;
|
||||
use indoc::formatdoc;
|
||||
use rmcp::model::{PromptArgument, ServerCapabilities};
|
||||
use rmcp::service::RunningService;
|
||||
use rmcp::transport::StreamableHttpClientTransport;
|
||||
use rmcp::transport::TokioChildProcess;
|
||||
@@ -34,27 +36,97 @@ use tokio::process::Command;
|
||||
pub const MCP_INVOKE_META_FUNCTION_NAME_PREFIX: &str = "mcp_invoke";
|
||||
pub const MCP_SEARCH_META_FUNCTION_NAME_PREFIX: &str = "mcp_search";
|
||||
pub const MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX: &str = "mcp_describe";
|
||||
pub const MCP_READ_META_FUNCTION_NAME_PREFIX: &str = "mcp_read";
|
||||
pub const MCP_PROMPT_META_FUNCTION_NAME_PREFIX: &str = "mcp_prompt";
|
||||
|
||||
pub const MCP_META_FUNCTION_PREFIXES: [&str; 5] = [
|
||||
MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
||||
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
|
||||
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX,
|
||||
MCP_READ_META_FUNCTION_NAME_PREFIX,
|
||||
MCP_PROMPT_META_FUNCTION_NAME_PREFIX,
|
||||
];
|
||||
|
||||
pub fn is_mcp_meta_function(name: &str) -> bool {
|
||||
MCP_META_FUNCTION_PREFIXES
|
||||
.iter()
|
||||
.any(|prefix| name.starts_with(prefix))
|
||||
}
|
||||
|
||||
pub fn mcp_meta_function_names(server: &str) -> Vec<String> {
|
||||
MCP_META_FUNCTION_PREFIXES
|
||||
.iter()
|
||||
.map(|prefix| format!("{prefix}_{server}"))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub type ConnectedServer = RunningService<RoleClient, ()>;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct McpServerFeatures {
|
||||
pub name: String,
|
||||
pub tools: bool,
|
||||
pub resources: bool,
|
||||
pub prompts: bool,
|
||||
}
|
||||
|
||||
impl McpServerFeatures {
|
||||
pub fn from_capabilities(
|
||||
name: impl Into<String>,
|
||||
capabilities: Option<&ServerCapabilities>,
|
||||
) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
tools: capabilities.is_none_or(|c| c.tools.is_some()),
|
||||
resources: capabilities.is_some_and(|c| c.resources.is_some()),
|
||||
prompts: capabilities.is_some_and(|c| c.prompts.is_some()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum CatalogItemKind {
|
||||
#[default]
|
||||
Tool,
|
||||
Resource,
|
||||
ResourceTemplate,
|
||||
Prompt,
|
||||
}
|
||||
|
||||
impl CatalogItemKind {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Tool => "tool",
|
||||
Self::Resource => "resource",
|
||||
Self::ResourceTemplate => "resource_template",
|
||||
Self::Prompt => "prompt",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for CatalogItemKind {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, Serialize)]
|
||||
pub struct CatalogItem {
|
||||
pub kind: CatalogItemKind,
|
||||
pub name: String,
|
||||
pub server: String,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct ServerCatalog {
|
||||
items: HashMap<String, CatalogItem>,
|
||||
}
|
||||
|
||||
impl Clone for ServerCatalog {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
items: self.items.clone(),
|
||||
}
|
||||
}
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub uri: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub mime_type: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub size: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub arguments: Option<Vec<PromptArgument>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub audience: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
@@ -160,7 +232,6 @@ pub struct McpRegistry {
|
||||
log_path: Option<PathBuf>,
|
||||
config: Option<McpServersConfig>,
|
||||
servers: HashMap<String, Arc<ConnectedServer>>,
|
||||
catalogs: HashMap<String, ServerCatalog>,
|
||||
}
|
||||
|
||||
impl McpRegistry {
|
||||
@@ -303,7 +374,7 @@ impl McpRegistry {
|
||||
|
||||
debug!("Starting selected MCP servers: {:?}", ids_to_start);
|
||||
|
||||
let results: Vec<Option<(String, Arc<ConnectedServer>, ServerCatalog)>> = stream::iter(
|
||||
let results: Vec<Option<(String, Arc<ConnectedServer>)>> = stream::iter(
|
||||
ids_to_start
|
||||
.into_iter()
|
||||
.map(|id| async { self.start_server(id).await }),
|
||||
@@ -312,18 +383,14 @@ impl McpRegistry {
|
||||
.try_collect()
|
||||
.await?;
|
||||
|
||||
for (id, server, catalog) in results.into_iter().flatten() {
|
||||
self.servers.insert(id.clone(), server);
|
||||
self.catalogs.insert(id, catalog);
|
||||
for (id, server) in results.into_iter().flatten() {
|
||||
self.servers.insert(id, server);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn start_server(
|
||||
&self,
|
||||
id: String,
|
||||
) -> Result<Option<(String, Arc<ConnectedServer>, ServerCatalog)>> {
|
||||
async fn start_server(&self, id: String) -> Result<Option<(String, Arc<ConnectedServer>)>> {
|
||||
let spec = self
|
||||
.config
|
||||
.as_ref()
|
||||
@@ -347,30 +414,9 @@ impl McpRegistry {
|
||||
Err(e) => return Err(e),
|
||||
};
|
||||
|
||||
let tools = service.list_tools(None).await?;
|
||||
debug!("Available tools for MCP server {id}: {tools:?}");
|
||||
|
||||
let mut items_vec = Vec::new();
|
||||
for t in tools.tools {
|
||||
let name = t.name.to_string();
|
||||
let description = t.description.unwrap_or_default().to_string();
|
||||
items_vec.push(CatalogItem {
|
||||
name,
|
||||
server: id.clone(),
|
||||
description,
|
||||
});
|
||||
}
|
||||
|
||||
let mut items_map = HashMap::new();
|
||||
items_vec.into_iter().for_each(|it| {
|
||||
items_map.insert(it.name.clone(), it);
|
||||
});
|
||||
|
||||
let catalog = ServerCatalog { items: items_map };
|
||||
|
||||
info!("Started MCP server: {id}");
|
||||
|
||||
Ok(Some((id.to_string(), service, catalog)))
|
||||
Ok(Some((id, service)))
|
||||
}
|
||||
|
||||
fn resolve_server_ids(&self, enabled_mcp_servers: Option<Vec<String>>) -> Vec<String> {
|
||||
@@ -398,8 +444,21 @@ impl McpRegistry {
|
||||
&self.servers
|
||||
}
|
||||
|
||||
pub fn list_started_servers(&self) -> Vec<String> {
|
||||
self.servers.keys().cloned().collect()
|
||||
pub fn server_features(&self) -> Vec<McpServerFeatures> {
|
||||
let mut features: Vec<McpServerFeatures> = self
|
||||
.servers
|
||||
.iter()
|
||||
.map(|(name, handle)| {
|
||||
let info = handle.peer_info();
|
||||
McpServerFeatures::from_capabilities(
|
||||
name.as_str(),
|
||||
info.as_ref().map(|info| &info.capabilities),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
features.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
features
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
@@ -1161,7 +1220,7 @@ mod tests {
|
||||
let registry = McpRegistry::default();
|
||||
|
||||
assert!(registry.is_empty());
|
||||
assert!(registry.list_started_servers().is_empty());
|
||||
assert!(registry.server_features().is_empty());
|
||||
assert!(registry.mcp_config().is_none());
|
||||
assert!(registry.log_path().is_none());
|
||||
}
|
||||
@@ -1185,6 +1244,51 @@ mod tests {
|
||||
assert_eq!(MCP_INVOKE_META_FUNCTION_NAME_PREFIX, "mcp_invoke");
|
||||
assert_eq!(MCP_SEARCH_META_FUNCTION_NAME_PREFIX, "mcp_search");
|
||||
assert_eq!(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, "mcp_describe");
|
||||
assert_eq!(MCP_READ_META_FUNCTION_NAME_PREFIX, "mcp_read");
|
||||
assert_eq!(MCP_PROMPT_META_FUNCTION_NAME_PREFIX, "mcp_prompt");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_mcp_meta_function_classifies_names() {
|
||||
assert!(is_mcp_meta_function("mcp_invoke_github"));
|
||||
assert!(is_mcp_meta_function("mcp_search_github"));
|
||||
assert!(is_mcp_meta_function("mcp_describe_github"));
|
||||
assert!(is_mcp_meta_function("mcp_read_github"));
|
||||
assert!(is_mcp_meta_function("mcp_prompt_github"));
|
||||
assert!(!is_mcp_meta_function("mcp_gateway_tool"));
|
||||
assert!(!is_mcp_meta_function("fs_read"));
|
||||
assert!(!is_mcp_meta_function(""));
|
||||
assert!(!is_mcp_meta_function("mcp_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn meta_function_prefixes_are_not_prefixes_of_each_other() {
|
||||
for (i, a) in MCP_META_FUNCTION_PREFIXES.iter().enumerate() {
|
||||
for (j, b) in MCP_META_FUNCTION_PREFIXES.iter().enumerate() {
|
||||
if i != j {
|
||||
assert!(!b.starts_with(a), "{a} is a prefix of {b}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_mcp_meta_function_preserves_lax_prefix_matching() {
|
||||
assert!(is_mcp_meta_function("mcp_invoker_x"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_meta_function_names_returns_all_prefixes_in_order() {
|
||||
assert_eq!(
|
||||
mcp_meta_function_names("github"),
|
||||
vec![
|
||||
"mcp_invoke_github",
|
||||
"mcp_search_github",
|
||||
"mcp_describe_github",
|
||||
"mcp_read_github",
|
||||
"mcp_prompt_github",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -0,0 +1,896 @@
|
||||
//! Content policy for MCP resource and tool content: UTF-8-boundary-safe text
|
||||
//! paging, grep-style pattern filtering, and spill-to-disk for binary blobs.
|
||||
|
||||
use crate::config::paths;
|
||||
use base64::engine::general_purpose::STANDARD;
|
||||
use base64::read::DecoderReader;
|
||||
use fancy_regex::Regex;
|
||||
use serde::Serialize;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::error::Error;
|
||||
use std::fs::{self, OpenOptions};
|
||||
use std::io::{ErrorKind, Read, Write};
|
||||
#[cfg(unix)]
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::SystemTime;
|
||||
use std::{fmt, io};
|
||||
|
||||
/// Default page size when the caller does not specify `max_bytes`.
|
||||
pub const DEFAULT_TEXT_MAX_BYTES: usize = 51_200;
|
||||
/// Hard upper bound on a single text slice regardless of requested `max_bytes`.
|
||||
pub const TEXT_MAX_BYTES_CLAMP: usize = 204_800;
|
||||
/// Maximum decoded size of a base64 blob before rendering is refused.
|
||||
pub const BLOB_DECODE_CEILING_BYTES: usize = 50 * 1024 * 1024;
|
||||
/// Total size bound for the spill tree; oldest files are evicted beyond it.
|
||||
pub const SPILL_DIR_MAX_BYTES: u64 = 512 * 1024 * 1024;
|
||||
/// Byte bound on server-supplied metadata strings (uri, mime type) copied into output.
|
||||
pub const METADATA_MAX_BYTES: usize = 4096;
|
||||
|
||||
const PATTERN_CONTEXT_LINES: usize = 2;
|
||||
const HUNK_SEPARATOR: &str = "--";
|
||||
|
||||
const MIME_EXTENSIONS: &[(&str, &str)] = &[
|
||||
("application/gzip", "gz"),
|
||||
("application/json", "json"),
|
||||
("application/pdf", "pdf"),
|
||||
("application/zip", "zip"),
|
||||
("audio/mpeg", "mp3"),
|
||||
("image/gif", "gif"),
|
||||
("image/jpeg", "jpg"),
|
||||
("image/png", "png"),
|
||||
("image/webp", "webp"),
|
||||
("text/csv", "csv"),
|
||||
("video/mp4", "mp4"),
|
||||
];
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum RenderError {
|
||||
InvalidPattern { pattern: String, error: String },
|
||||
DecodedSizeExceeded,
|
||||
InvalidBase64(String),
|
||||
Io(io::Error),
|
||||
}
|
||||
|
||||
impl fmt::Display for RenderError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Self::InvalidPattern { pattern, error } => write!(
|
||||
f,
|
||||
"Invalid filter pattern '{pattern}': {error}. Provide a valid regex; \
|
||||
lines matching it are returned with {PATTERN_CONTEXT_LINES} lines of context."
|
||||
),
|
||||
Self::DecodedSizeExceeded => write!(
|
||||
f,
|
||||
"Decoded blob exceeds BLOB_DECODE_CEILING_BYTES ({} MiB); refusing to render it",
|
||||
BLOB_DECODE_CEILING_BYTES / (1024 * 1024)
|
||||
),
|
||||
Self::InvalidBase64(error) => write!(f, "Invalid base64 in blob content: {error}"),
|
||||
Self::Io(error) => write!(f, "Failed to spill blob to disk: {error}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Error for RenderError {
|
||||
fn source(&self) -> Option<&(dyn Error + 'static)> {
|
||||
match self {
|
||||
Self::Io(error) => Some(error),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<io::Error> for RenderError {
|
||||
fn from(error: io::Error) -> Self {
|
||||
Self::Io(error)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RenderedText {
|
||||
pub text: String,
|
||||
pub truncated: bool,
|
||||
pub total_bytes: usize,
|
||||
pub next_offset: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum RenderedBlob {
|
||||
Text(String),
|
||||
Spilled(SpillMetadata),
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SpillMetadata {
|
||||
pub spilled: bool,
|
||||
pub path: PathBuf,
|
||||
pub mime_type: Option<String>,
|
||||
pub sniffed: bool,
|
||||
pub size_bytes: u64,
|
||||
pub sha256: String,
|
||||
}
|
||||
|
||||
/// Pages `text` with UTF-8-boundary-safe slicing. When `pattern` is set, the
|
||||
/// text is first reduced to matching lines plus context (grep-style, with
|
||||
/// 1-based line-number prefixes), and all offset/size math operates on that
|
||||
/// filtered stream.
|
||||
pub fn render_text(
|
||||
text: &str,
|
||||
pattern: Option<&str>,
|
||||
offset: usize,
|
||||
max_bytes: Option<usize>,
|
||||
) -> Result<RenderedText, RenderError> {
|
||||
let filtered = match pattern {
|
||||
Some(pattern) => Some(filter_lines(text, pattern)?),
|
||||
None => None,
|
||||
};
|
||||
let stream = filtered.as_deref().unwrap_or(text);
|
||||
let max_bytes = max_bytes
|
||||
.unwrap_or(DEFAULT_TEXT_MAX_BYTES)
|
||||
.min(TEXT_MAX_BYTES_CLAMP);
|
||||
let total_bytes = stream.len();
|
||||
let mut start = offset.min(total_bytes);
|
||||
while !stream.is_char_boundary(start) {
|
||||
start += 1;
|
||||
}
|
||||
let mut end = start.saturating_add(max_bytes).min(total_bytes);
|
||||
while !stream.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
// A max_bytes smaller than one codepoint would produce an empty page with
|
||||
// next_offset == offset, stalling paging; always advance by at least one.
|
||||
if end == start && start < total_bytes {
|
||||
end += 1;
|
||||
while !stream.is_char_boundary(end) {
|
||||
end += 1;
|
||||
}
|
||||
}
|
||||
|
||||
let truncated = end < total_bytes;
|
||||
Ok(RenderedText {
|
||||
text: stream[start..end].to_string(),
|
||||
truncated,
|
||||
total_bytes,
|
||||
next_offset: truncated.then_some(end),
|
||||
})
|
||||
}
|
||||
|
||||
/// Decodes a base64 blob, returning it as text when it is valid UTF-8 and
|
||||
/// spilling it under `cache_dir()/mcp-resources/<server>/` otherwise.
|
||||
pub fn render_blob(
|
||||
b64: &str,
|
||||
claimed_mime: Option<&str>,
|
||||
server: &str,
|
||||
) -> Result<RenderedBlob, RenderError> {
|
||||
let spill_base = paths::cache_dir().join("mcp-resources");
|
||||
render_blob_at(b64, claimed_mime, server, &spill_base)
|
||||
}
|
||||
|
||||
pub fn render_blob_at(
|
||||
b64: &str,
|
||||
claimed_mime: Option<&str>,
|
||||
server: &str,
|
||||
spill_base: &Path,
|
||||
) -> Result<RenderedBlob, RenderError> {
|
||||
let decoded = decode_base64_bounded(b64)?;
|
||||
let decoded = match String::from_utf8(decoded) {
|
||||
Ok(text) => return Ok(RenderedBlob::Text(text)),
|
||||
Err(error) => error.into_bytes(),
|
||||
};
|
||||
let sha256 = format!("{:x}", Sha256::digest(&decoded));
|
||||
let dir = spill_base.join(sanitize_server(server));
|
||||
fs::create_dir_all(&dir)?;
|
||||
let path = dir.join(format!("{sha256}.{}", extension_for_mime(claimed_mime)));
|
||||
|
||||
// Writes land in a temp file and are renamed into place, so a visible
|
||||
// file at the final path is always complete and the dedup check below is
|
||||
// race-safe across processes (same sha means same content).
|
||||
if !path.exists() {
|
||||
static TEMP_COUNTER: AtomicUsize = AtomicUsize::new(0);
|
||||
let temp = dir.join(format!(
|
||||
"{sha256}.tmp-{}-{}",
|
||||
std::process::id(),
|
||||
TEMP_COUNTER.fetch_add(1, Ordering::Relaxed)
|
||||
));
|
||||
let mut options = OpenOptions::new();
|
||||
options.write(true).create_new(true);
|
||||
#[cfg(unix)]
|
||||
options.mode(0o600);
|
||||
let written = options
|
||||
.open(&temp)
|
||||
.and_then(|mut file| file.write_all(&decoded))
|
||||
.and_then(|()| fs::rename(&temp, &path));
|
||||
if let Err(error) = written {
|
||||
let _ = fs::remove_file(&temp);
|
||||
return Err(RenderError::Io(error));
|
||||
}
|
||||
}
|
||||
enforce_spill_bound(spill_base, SPILL_DIR_MAX_BYTES, &path);
|
||||
|
||||
Ok(RenderedBlob::Spilled(SpillMetadata {
|
||||
spilled: true,
|
||||
path,
|
||||
mime_type: claimed_mime.map(str::to_string),
|
||||
sniffed: false,
|
||||
size_bytes: decoded.len() as u64,
|
||||
sha256,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Truncates `text` to at most `max_bytes`, rounding the cut point back to a
|
||||
/// UTF-8 character boundary.
|
||||
pub fn truncate_utf8(text: &str, max_bytes: usize) -> &str {
|
||||
if text.len() <= max_bytes {
|
||||
return text;
|
||||
}
|
||||
let mut end = max_bytes;
|
||||
while !text.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
&text[..end]
|
||||
}
|
||||
|
||||
/// Bounds a server-supplied metadata string to [`METADATA_MAX_BYTES`],
|
||||
/// appending a marker citing the constant when the input is truncated.
|
||||
pub fn clamp_metadata(text: &str) -> String {
|
||||
if text.len() <= METADATA_MAX_BYTES {
|
||||
return text.to_string();
|
||||
}
|
||||
|
||||
let clamped = truncate_utf8(text, METADATA_MAX_BYTES);
|
||||
format!("{clamped} [truncated: exceeds METADATA_MAX_BYTES ({METADATA_MAX_BYTES} bytes)]")
|
||||
}
|
||||
|
||||
fn filter_lines(text: &str, pattern: &str) -> Result<String, RenderError> {
|
||||
let regex = Regex::new(pattern).map_err(|error| RenderError::InvalidPattern {
|
||||
pattern: pattern.to_string(),
|
||||
error: error.to_string(),
|
||||
})?;
|
||||
let lines: Vec<&str> = text.lines().collect();
|
||||
// fancy_regex can also fail at match time (backtracking limits); treat
|
||||
// that as a non-match rather than failing the whole render.
|
||||
let is_match: Vec<bool> = lines
|
||||
.iter()
|
||||
.map(|line| regex.is_match(line).unwrap_or(false))
|
||||
.collect();
|
||||
|
||||
let mut keep = vec![false; lines.len()];
|
||||
for (i, _) in is_match.iter().enumerate().filter(|&(_, matched)| *matched) {
|
||||
let start = i.saturating_sub(PATTERN_CONTEXT_LINES);
|
||||
let end = (i + PATTERN_CONTEXT_LINES).min(lines.len() - 1);
|
||||
keep[start..=end].fill(true);
|
||||
}
|
||||
|
||||
let mut out: Vec<String> = Vec::new();
|
||||
let mut prev_kept: Option<usize> = None;
|
||||
for (i, line) in lines.iter().enumerate() {
|
||||
if !keep[i] {
|
||||
continue;
|
||||
}
|
||||
if prev_kept.is_some_and(|prev| i > prev + 1) {
|
||||
out.push(HUNK_SEPARATOR.to_string());
|
||||
}
|
||||
let marker = if is_match[i] { ':' } else { '-' };
|
||||
out.push(format!("{}{marker}{line}", i + 1));
|
||||
prev_kept = Some(i);
|
||||
}
|
||||
|
||||
Ok(out.join("\n"))
|
||||
}
|
||||
|
||||
fn decode_base64_bounded(b64: &str) -> Result<Vec<u8>, RenderError> {
|
||||
// The encoded length puts a lower bound on the decoded size; reject
|
||||
// inputs that bound already proves oversized before decoding anything.
|
||||
let min_decoded = (b64.len() / 4).saturating_mul(3).saturating_sub(2);
|
||||
if min_decoded > BLOB_DECODE_CEILING_BYTES {
|
||||
return Err(RenderError::DecodedSizeExceeded);
|
||||
}
|
||||
|
||||
let mut reader = DecoderReader::new(b64.as_bytes(), &STANDARD);
|
||||
let mut decoded = Vec::new();
|
||||
let mut chunk = [0u8; 8192];
|
||||
loop {
|
||||
match reader.read(&mut chunk) {
|
||||
Ok(0) => return Ok(decoded),
|
||||
Ok(n) => {
|
||||
if decoded.len() + n > BLOB_DECODE_CEILING_BYTES {
|
||||
return Err(RenderError::DecodedSizeExceeded);
|
||||
}
|
||||
decoded.extend_from_slice(&chunk[..n]);
|
||||
}
|
||||
Err(error) => return Err(RenderError::InvalidBase64(error.to_string())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Maps a server-controlled mime type to a spill-file extension via an exact
|
||||
/// allowlist lookup; anything unrecognized falls back to `bin`.
|
||||
fn extension_for_mime(mime: Option<&str>) -> &'static str {
|
||||
let Some(mime) = mime else {
|
||||
return "bin";
|
||||
};
|
||||
let bare = mime
|
||||
.split(';')
|
||||
.next()
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
let ext = MIME_EXTENSIONS
|
||||
.iter()
|
||||
.find(|(known, _)| *known == bare)
|
||||
.map(|(_, ext)| *ext)
|
||||
.unwrap_or("bin");
|
||||
let safe = !ext.is_empty()
|
||||
&& ext.len() <= 8
|
||||
&& ext
|
||||
.bytes()
|
||||
.all(|b| b.is_ascii_lowercase() || b.is_ascii_digit());
|
||||
|
||||
if safe { ext } else { "bin" }
|
||||
}
|
||||
|
||||
fn sanitize_server(server: &str) -> String {
|
||||
let mut sanitized: String = server
|
||||
.chars()
|
||||
.map(|c| {
|
||||
if c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') {
|
||||
c
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.take(64)
|
||||
.collect();
|
||||
// Windows strips trailing dots at create time, which would make the
|
||||
// constructed path disagree with the on-disk name.
|
||||
while sanitized.ends_with('.') {
|
||||
sanitized.pop();
|
||||
}
|
||||
if sanitized.is_empty() {
|
||||
return "_".to_string();
|
||||
}
|
||||
// Windows reserves device names (bare or with any extension).
|
||||
let stem = sanitized.split('.').next().unwrap_or("");
|
||||
if is_windows_reserved(stem) {
|
||||
sanitized.insert(0, '_');
|
||||
}
|
||||
sanitized
|
||||
}
|
||||
|
||||
fn is_windows_reserved(stem: &str) -> bool {
|
||||
let lower = stem.to_ascii_lowercase();
|
||||
matches!(lower.as_str(), "con" | "prn" | "aux" | "nul")
|
||||
|| (lower.len() == 4
|
||||
&& (lower.starts_with("com") || lower.starts_with("lpt"))
|
||||
&& matches!(lower.as_bytes()[3], b'1'..=b'9'))
|
||||
}
|
||||
|
||||
struct SpillEntry {
|
||||
path: PathBuf,
|
||||
size: u64,
|
||||
modified: SystemTime,
|
||||
}
|
||||
|
||||
fn enforce_spill_bound(base: &Path, max_total: u64, protect: &Path) {
|
||||
let mut entries = Vec::new();
|
||||
collect_spill_files(base, &mut entries);
|
||||
evict_oldest(entries, max_total, protect);
|
||||
}
|
||||
|
||||
/// Best-effort eviction: the spill dir is shared across processes, so a file
|
||||
/// vanishing underneath us (`NotFound`) is expected and never fails the spill.
|
||||
fn evict_oldest(mut entries: Vec<SpillEntry>, max_total: u64, protect: &Path) {
|
||||
let mut total: u64 = entries.iter().map(|entry| entry.size).sum();
|
||||
if total <= max_total {
|
||||
return;
|
||||
}
|
||||
entries.sort_by_key(|entry| entry.modified);
|
||||
for entry in &entries {
|
||||
if total <= max_total {
|
||||
break;
|
||||
}
|
||||
|
||||
// Filenames are content-hashed, so name equality is sufficient and
|
||||
// survives filesystems that normalize directory names (case folding,
|
||||
// trailing-dot stripping) where a full-path comparison would miss.
|
||||
if entry.path.file_name() == protect.file_name() {
|
||||
continue;
|
||||
}
|
||||
|
||||
match fs::remove_file(&entry.path) {
|
||||
Ok(()) => total -= entry.size,
|
||||
Err(error) if error.kind() == ErrorKind::NotFound => total -= entry.size,
|
||||
Err(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_spill_files(dir: &Path, out: &mut Vec<SpillEntry>) {
|
||||
let Ok(entries) = fs::read_dir(dir) else {
|
||||
return;
|
||||
};
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
let Ok(metadata) = entry.metadata() else {
|
||||
continue;
|
||||
};
|
||||
if metadata.is_dir() {
|
||||
collect_spill_files(&path, out);
|
||||
} else if metadata.is_file() {
|
||||
out.push(SpillEntry {
|
||||
path,
|
||||
size: metadata.len(),
|
||||
modified: metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use base64::Engine;
|
||||
use std::env;
|
||||
#[cfg(unix)]
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
use std::process;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
fn with_spill_base<F: FnOnce(&Path)>(f: F) {
|
||||
static COUNTER: AtomicUsize = AtomicUsize::new(0);
|
||||
let unique = format!(
|
||||
"{}-{}",
|
||||
process::id(),
|
||||
COUNTER.fetch_add(1, Ordering::Relaxed)
|
||||
);
|
||||
let base = env::temp_dir().join(format!("coyote-render-test-{unique}"));
|
||||
fs::create_dir_all(&base).unwrap();
|
||||
f(&base);
|
||||
let _ = fs::remove_dir_all(&base);
|
||||
}
|
||||
|
||||
fn set_mtime(path: &Path, secs_after_epoch: u64) {
|
||||
let file = OpenOptions::new().write(true).open(path).unwrap();
|
||||
file.set_modified(SystemTime::UNIX_EPOCH + Duration::from_secs(secs_after_epoch))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
fn write_spill_file(dir: &Path, name: &str, len: usize, mtime_secs: u64) -> PathBuf {
|
||||
let path = dir.join(name);
|
||||
fs::write(&path, vec![0u8; len]).unwrap();
|
||||
set_mtime(&path, mtime_secs);
|
||||
path
|
||||
}
|
||||
|
||||
const TEN_LINES: &str = "one\ntwo\nthree\nfour\nfive\nsix\nseven\neight\nnine\nten";
|
||||
|
||||
#[test]
|
||||
fn slices_basic_ascii_page() {
|
||||
let rendered = render_text("hello world", None, 0, Some(5)).unwrap();
|
||||
|
||||
assert_eq!(rendered.text, "hello");
|
||||
assert!(rendered.truncated);
|
||||
assert_eq!(rendered.total_bytes, 11);
|
||||
assert_eq!(rendered.next_offset, Some(5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn offset_mid_codepoint_rounds_forward() {
|
||||
// 'é' occupies bytes 1..3; offset 2 lands inside it.
|
||||
let rendered = render_text("héllo", None, 2, None).unwrap();
|
||||
|
||||
assert_eq!(rendered.text, "llo");
|
||||
assert!(!rendered.truncated);
|
||||
assert_eq!(rendered.next_offset, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn end_mid_codepoint_rounds_backward() {
|
||||
// 'é' occupies bytes 1..3; offset 0 + max_bytes 2 lands inside it.
|
||||
let rendered = render_text("aé", None, 0, Some(2)).unwrap();
|
||||
|
||||
assert_eq!(rendered.text, "a");
|
||||
assert!(rendered.truncated);
|
||||
assert_eq!(rendered.total_bytes, 3);
|
||||
assert_eq!(rendered.next_offset, Some(1));
|
||||
|
||||
let rest = render_text("aé", None, 1, Some(2)).unwrap();
|
||||
|
||||
assert_eq!(rest.text, "é");
|
||||
assert!(!rest.truncated);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn max_bytes_below_one_codepoint_still_advances() {
|
||||
// 'é' is 2 bytes; max_bytes 1 must not stall at next_offset == offset.
|
||||
let rendered = render_text("éa", None, 0, Some(1)).unwrap();
|
||||
|
||||
assert_eq!(rendered.text, "é");
|
||||
assert!(rendered.truncated);
|
||||
assert_eq!(rendered.total_bytes, 3);
|
||||
assert_eq!(rendered.next_offset, Some(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn offset_past_eof_returns_empty() {
|
||||
let rendered = render_text("short", None, 100, None).unwrap();
|
||||
|
||||
assert_eq!(rendered.text, "");
|
||||
assert!(!rendered.truncated);
|
||||
assert_eq!(rendered.total_bytes, 5);
|
||||
assert_eq!(rendered.next_offset, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exact_fit_is_not_truncated() {
|
||||
let rendered = render_text("exact", None, 0, Some(5)).unwrap();
|
||||
|
||||
assert_eq!(rendered.text, "exact");
|
||||
assert!(!rendered.truncated);
|
||||
assert_eq!(rendered.next_offset, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_max_bytes_is_default_text_max_bytes() {
|
||||
let text = "a".repeat(DEFAULT_TEXT_MAX_BYTES + 1);
|
||||
|
||||
let rendered = render_text(&text, None, 0, None).unwrap();
|
||||
|
||||
assert_eq!(rendered.text.len(), DEFAULT_TEXT_MAX_BYTES);
|
||||
assert!(rendered.truncated);
|
||||
assert_eq!(rendered.next_offset, Some(DEFAULT_TEXT_MAX_BYTES));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn max_bytes_above_clamp_is_clamped() {
|
||||
let text = "a".repeat(TEXT_MAX_BYTES_CLAMP + 1);
|
||||
|
||||
let rendered = render_text(&text, None, 0, Some(usize::MAX)).unwrap();
|
||||
|
||||
assert_eq!(rendered.text.len(), TEXT_MAX_BYTES_CLAMP);
|
||||
assert!(rendered.truncated);
|
||||
assert_eq!(rendered.next_offset, Some(TEXT_MAX_BYTES_CLAMP));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_utf8_rounds_back_to_char_boundary() {
|
||||
// 'é' occupies bytes 1..3; a cut at byte 2 lands inside it.
|
||||
assert_eq!(truncate_utf8("aé", 2), "a");
|
||||
assert_eq!(truncate_utf8("aé", 3), "aé");
|
||||
assert_eq!(truncate_utf8("abc", 10), "abc");
|
||||
assert_eq!(truncate_utf8("abc", 0), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clamp_metadata_appends_marker_only_when_oversized() {
|
||||
assert_eq!(clamp_metadata("text/plain"), "text/plain");
|
||||
|
||||
let long = "u".repeat(METADATA_MAX_BYTES + 1);
|
||||
|
||||
let clamped = clamp_metadata(&long);
|
||||
|
||||
assert!(clamped.starts_with(&"u".repeat(METADATA_MAX_BYTES)));
|
||||
assert!(clamped.contains("METADATA_MAX_BYTES"));
|
||||
assert!(clamped.contains(&METADATA_MAX_BYTES.to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pattern_emits_matches_with_context_and_line_numbers() {
|
||||
let rendered = render_text(TEN_LINES, Some("^five$"), 0, None).unwrap();
|
||||
|
||||
assert_eq!(rendered.text, "3-three\n4-four\n5:five\n6-six\n7-seven");
|
||||
assert!(!rendered.truncated);
|
||||
assert_eq!(rendered.total_bytes, rendered.text.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pattern_separates_disjoint_hunks() {
|
||||
let rendered = render_text(TEN_LINES, Some("^(two|nine)$"), 0, None).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
rendered.text,
|
||||
"1-one\n2:two\n3-three\n4-four\n--\n7-seven\n8-eight\n9:nine\n10-ten"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pattern_merges_adjacent_hunks_without_duplicates() {
|
||||
let rendered = render_text(TEN_LINES, Some("^(two|six)$"), 0, None).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
rendered.text,
|
||||
"1-one\n2:two\n3-three\n4-four\n5-five\n6:six\n7-seven\n8-eight"
|
||||
);
|
||||
assert!(!rendered.text.contains(HUNK_SEPARATOR));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pattern_paging_walks_the_filtered_stream() {
|
||||
let full = render_text(TEN_LINES, Some("^t"), 0, None).unwrap();
|
||||
assert!(!full.truncated);
|
||||
|
||||
let mut assembled = String::new();
|
||||
let mut offset = 0;
|
||||
loop {
|
||||
let page = render_text(TEN_LINES, Some("^t"), offset, Some(7)).unwrap();
|
||||
assert_eq!(page.total_bytes, full.text.len());
|
||||
assembled.push_str(&page.text);
|
||||
match page.next_offset {
|
||||
Some(next) => offset = next,
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(assembled, full.text);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pattern_with_no_matches_returns_empty() {
|
||||
let rendered = render_text(TEN_LINES, Some("^zebra$"), 0, None).unwrap();
|
||||
|
||||
assert_eq!(rendered.text, "");
|
||||
assert_eq!(rendered.total_bytes, 0);
|
||||
assert!(!rendered.truncated);
|
||||
assert_eq!(rendered.next_offset, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_pattern_is_a_teaching_error() {
|
||||
let parse_error = Regex::new("(").unwrap_err().to_string();
|
||||
|
||||
let err = render_text("text", Some("("), 0, None).unwrap_err();
|
||||
|
||||
assert!(matches!(err, RenderError::InvalidPattern { .. }));
|
||||
let message = err.to_string();
|
||||
assert!(message.contains("'('"));
|
||||
assert!(message.contains(&parse_error));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn utf8_blob_decodes_to_text_without_spilling() {
|
||||
with_spill_base(|base| {
|
||||
let b64 = STANDARD.encode("hello ✓ world");
|
||||
|
||||
let rendered = render_blob_at(&b64, Some("text/plain"), "srv", base).unwrap();
|
||||
|
||||
let RenderedBlob::Text(text) = rendered else {
|
||||
panic!("expected text variant");
|
||||
};
|
||||
assert_eq!(text, "hello ✓ world");
|
||||
assert_eq!(fs::read_dir(base).unwrap().count(), 0);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binary_blob_spills_with_metadata_and_0600_perms() {
|
||||
with_spill_base(|base| {
|
||||
let data: &[u8] = &[0xff, 0xfe, 0x00, 0x88, 0x01];
|
||||
let b64 = STANDARD.encode(data);
|
||||
|
||||
let rendered = render_blob_at(&b64, Some("application/pdf"), "docs", base).unwrap();
|
||||
|
||||
let RenderedBlob::Spilled(meta) = rendered else {
|
||||
panic!("expected spilled variant");
|
||||
};
|
||||
let expected_sha = format!("{:x}", Sha256::digest(data));
|
||||
assert_eq!(meta.sha256, expected_sha);
|
||||
assert_eq!(
|
||||
meta.path,
|
||||
base.join("docs").join(format!("{expected_sha}.pdf"))
|
||||
);
|
||||
assert_eq!(meta.size_bytes, data.len() as u64);
|
||||
assert_eq!(meta.mime_type.as_deref(), Some("application/pdf"));
|
||||
assert!(!meta.sniffed);
|
||||
assert!(meta.spilled);
|
||||
assert_eq!(fs::read(&meta.path).unwrap(), data);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let mode = fs::metadata(&meta.path).unwrap().permissions().mode();
|
||||
assert_eq!(mode & 0o777, 0o600);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_ceiling_rejects_oversized_blob() {
|
||||
with_spill_base(|base| {
|
||||
// base64 of 51 MiB of zero bytes is just a repeated-'A' string.
|
||||
let encoded = "A".repeat(51 * 1024 * 1024 / 3 * 4);
|
||||
|
||||
let err = render_blob_at(&encoded, None, "srv", base).unwrap_err();
|
||||
|
||||
assert!(matches!(err, RenderError::DecodedSizeExceeded));
|
||||
assert!(err.to_string().contains("BLOB_DECODE_CEILING_BYTES"));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_base64_is_rejected() {
|
||||
with_spill_base(|base| {
|
||||
let err = render_blob_at("!!!not base64!!!", None, "srv", base).unwrap_err();
|
||||
|
||||
assert!(matches!(err, RenderError::InvalidBase64(_)));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn spill_dedup_returns_same_path_without_rewriting() {
|
||||
with_spill_base(|base| {
|
||||
let data: &[u8] = &[0xff, 0x01, 0x02];
|
||||
let b64 = STANDARD.encode(data);
|
||||
|
||||
let RenderedBlob::Spilled(first) = render_blob_at(&b64, None, "srv", base).unwrap()
|
||||
else {
|
||||
panic!("expected spilled variant");
|
||||
};
|
||||
fs::write(&first.path, b"sentinel").unwrap();
|
||||
|
||||
let RenderedBlob::Spilled(second) = render_blob_at(&b64, None, "srv", base).unwrap()
|
||||
else {
|
||||
panic!("expected spilled variant");
|
||||
};
|
||||
|
||||
assert_eq!(second.path, first.path);
|
||||
assert_eq!(second.sha256, first.sha256);
|
||||
assert_eq!(fs::read(&second.path).unwrap(), b"sentinel");
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn spill_metadata_serializes_spilled_true() {
|
||||
with_spill_base(|base| {
|
||||
let b64 = STANDARD.encode([0xffu8, 0x00]);
|
||||
|
||||
let RenderedBlob::Spilled(meta) =
|
||||
render_blob_at(&b64, Some("image/png"), "srv", base).unwrap()
|
||||
else {
|
||||
panic!("expected spilled variant");
|
||||
};
|
||||
|
||||
let value = serde_json::to_value(&meta).unwrap();
|
||||
assert_eq!(value["spilled"], serde_json::Value::Bool(true));
|
||||
assert_eq!(value["sniffed"], serde_json::Value::Bool(false));
|
||||
assert_eq!(value["sha256"].as_str(), Some(meta.sha256.as_str()));
|
||||
assert_eq!(value["mime_type"].as_str(), Some("image/png"));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extension_allowlist_normalizes_and_defaults_to_bin() {
|
||||
assert_eq!(extension_for_mime(Some("application/pdf")), "pdf");
|
||||
assert_eq!(extension_for_mime(Some("image/png")), "png");
|
||||
assert_eq!(extension_for_mime(Some(" TEXT/CSV ; charset=utf-8")), "csv");
|
||||
assert_eq!(extension_for_mime(Some("../../evil")), "bin");
|
||||
assert_eq!(extension_for_mime(Some("image/png/../../x")), "bin");
|
||||
assert_eq!(extension_for_mime(Some("application/x-∞")), "bin");
|
||||
assert_eq!(extension_for_mime(Some("text/plain")), "bin");
|
||||
assert_eq!(extension_for_mime(None), "bin");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_server_strips_path_separators() {
|
||||
assert_eq!(sanitize_server("../evil/srv"), ".._evil_srv");
|
||||
assert_eq!(sanitize_server("srv name!"), "srv_name_");
|
||||
assert_eq!(sanitize_server(""), "_");
|
||||
assert_eq!(sanitize_server("."), "_");
|
||||
assert_eq!(sanitize_server(".."), "_");
|
||||
assert_eq!(sanitize_server("good-server_1.0"), "good-server_1.0");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_server_escapes_windows_reserved_names() {
|
||||
assert_eq!(sanitize_server("con"), "_con");
|
||||
assert_eq!(sanitize_server("CON"), "_CON");
|
||||
assert_eq!(sanitize_server("nul.txt"), "_nul.txt");
|
||||
assert_eq!(sanitize_server("COM1"), "_COM1");
|
||||
assert_eq!(sanitize_server("lpt9"), "_lpt9");
|
||||
assert_eq!(sanitize_server("com0"), "com0");
|
||||
assert_eq!(sanitize_server("com10"), "com10");
|
||||
assert_eq!(sanitize_server("consul"), "consul");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_server_strips_trailing_dots_and_caps_length() {
|
||||
assert_eq!(sanitize_server("srv."), "srv");
|
||||
assert_eq!(sanitize_server("srv..."), "srv");
|
||||
assert_eq!(sanitize_server("..."), "_");
|
||||
let long = "a".repeat(100);
|
||||
assert_eq!(sanitize_server(&long).len(), 64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn spill_path_confines_crafted_server_and_mime() {
|
||||
with_spill_base(|base| {
|
||||
let b64 = STANDARD.encode([0xffu8, 0x00, 0x11]);
|
||||
|
||||
let RenderedBlob::Spilled(meta) =
|
||||
render_blob_at(&b64, Some("../../evil"), "../evil/srv", base).unwrap()
|
||||
else {
|
||||
panic!("expected spilled variant");
|
||||
};
|
||||
|
||||
assert!(meta.path.starts_with(base));
|
||||
let dir_name = meta.path.parent().unwrap().file_name().unwrap();
|
||||
assert_eq!(dir_name, ".._evil_srv");
|
||||
assert_eq!(meta.path.extension().unwrap(), "bin");
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn eviction_removes_oldest_files_first_across_server_dirs() {
|
||||
with_spill_base(|base| {
|
||||
let srv_a = base.join("srv-a");
|
||||
let srv_b = base.join("srv-b");
|
||||
fs::create_dir_all(&srv_a).unwrap();
|
||||
fs::create_dir_all(&srv_b).unwrap();
|
||||
let oldest = write_spill_file(&srv_a, "a.bin", 100, 100);
|
||||
let middle = write_spill_file(&srv_b, "b.bin", 100, 200);
|
||||
let newest = write_spill_file(&srv_b, "c.bin", 100, 300);
|
||||
|
||||
enforce_spill_bound(base, 150, &newest);
|
||||
|
||||
assert!(!oldest.exists());
|
||||
assert!(!middle.exists());
|
||||
assert!(newest.exists());
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn eviction_skips_protected_file() {
|
||||
with_spill_base(|base| {
|
||||
let srv = base.join("srv");
|
||||
fs::create_dir_all(&srv).unwrap();
|
||||
let oldest = write_spill_file(&srv, "a.bin", 100, 100);
|
||||
let middle = write_spill_file(&srv, "b.bin", 100, 200);
|
||||
let newest = write_spill_file(&srv, "c.bin", 100, 300);
|
||||
|
||||
enforce_spill_bound(base, 250, &oldest);
|
||||
|
||||
assert!(oldest.exists());
|
||||
assert!(!middle.exists());
|
||||
assert!(newest.exists());
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn eviction_under_bound_is_noop() {
|
||||
with_spill_base(|base| {
|
||||
let srv = base.join("srv");
|
||||
fs::create_dir_all(&srv).unwrap();
|
||||
let first = write_spill_file(&srv, "a.bin", 100, 100);
|
||||
let second = write_spill_file(&srv, "b.bin", 100, 200);
|
||||
|
||||
enforce_spill_bound(base, 1000, &second);
|
||||
|
||||
assert!(first.exists());
|
||||
assert!(second.exists());
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn eviction_tolerates_already_removed_entries() {
|
||||
with_spill_base(|base| {
|
||||
let srv = base.join("srv");
|
||||
fs::create_dir_all(&srv).unwrap();
|
||||
let real = write_spill_file(&srv, "real.bin", 100, 200);
|
||||
let entries = vec![
|
||||
SpillEntry {
|
||||
path: srv.join("ghost.bin"),
|
||||
size: 100,
|
||||
modified: SystemTime::UNIX_EPOCH + Duration::from_secs(100),
|
||||
},
|
||||
SpillEntry {
|
||||
path: real.clone(),
|
||||
size: 100,
|
||||
modified: SystemTime::UNIX_EPOCH + Duration::from_secs(200),
|
||||
},
|
||||
];
|
||||
|
||||
evict_oldest(entries, 50, &base.join("untouched"));
|
||||
|
||||
assert!(!real.exists());
|
||||
});
|
||||
}
|
||||
}
|
||||
+293
-1
@@ -1,11 +1,17 @@
|
||||
use super::{REPL_COMMANDS, ReplCommand};
|
||||
|
||||
use crate::{config::RequestContext, utils::fuzzy_filter};
|
||||
use crate::config::{McpPromptCompletion, RequestContext, sanitize_display_text};
|
||||
use crate::mcp::ConnectedServer;
|
||||
use crate::utils::fuzzy_filter;
|
||||
|
||||
use parking_lot::RwLock;
|
||||
use reedline::{Completer, Span, Suggestion};
|
||||
use rmcp::model::Prompt;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
const PROMPT_COMPLETION_RPC_TIMEOUT: Duration = Duration::from_secs(2);
|
||||
|
||||
impl Completer for ReplCompleter {
|
||||
fn complete(&mut self, line: &str, pos: usize) -> Vec<Suggestion> {
|
||||
@@ -29,6 +35,22 @@ impl Completer for ReplCompleter {
|
||||
return suggestions;
|
||||
}
|
||||
|
||||
if cmd == ".prompt" && parts_len > 1 {
|
||||
let span = Span::new(parts[parts_len - 1].1, pos);
|
||||
let args: Vec<&str> = parts.iter().skip(1).map(|(v, _)| *v).collect();
|
||||
let filter = args.last().copied().unwrap_or_default().to_string();
|
||||
let stage = {
|
||||
let ctx = self.ctx.read();
|
||||
ctx.mcp_prompt_completion(&args)
|
||||
};
|
||||
return complete_prompt_stage(stage, &filter, PROMPT_COMPLETION_RPC_TIMEOUT)
|
||||
.iter()
|
||||
.map(|(value, description)| {
|
||||
create_suggestion(value, description.as_deref().unwrap_or_default(), span)
|
||||
})
|
||||
.collect();
|
||||
}
|
||||
|
||||
let ctx = self.ctx.read();
|
||||
let state = ctx.state();
|
||||
let model_has_reasoning = !ctx.current_model().reasoning_levels().is_empty();
|
||||
@@ -141,6 +163,71 @@ fn create_suggestion(value: &str, description: &str, span: Span) -> Suggestion {
|
||||
}
|
||||
}
|
||||
|
||||
fn complete_prompt_stage(
|
||||
stage: McpPromptCompletion,
|
||||
filter: &str,
|
||||
rpc_timeout: Duration,
|
||||
) -> Vec<(String, Option<String>)> {
|
||||
let values = match stage {
|
||||
McpPromptCompletion::Ready(values) => values,
|
||||
McpPromptCompletion::PromptNames { server } => list_prompts_blocking(server, rpc_timeout)
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|prompt| {
|
||||
(
|
||||
sanitize_display_text(&prompt.name),
|
||||
prompt
|
||||
.description
|
||||
.map(|description| sanitize_display_text(&description)),
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
McpPromptCompletion::ArgumentKeys {
|
||||
server,
|
||||
prompt,
|
||||
typed_keys,
|
||||
} => list_prompts_blocking(server, rpc_timeout)
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.find(|candidate| candidate.name == prompt)
|
||||
.and_then(|candidate| candidate.arguments)
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.filter(|arg| !typed_keys.contains(&arg.name))
|
||||
.map(|arg| {
|
||||
let description = arg
|
||||
.description
|
||||
.map(|description| sanitize_display_text(&description));
|
||||
let description = match (arg.required == Some(true), description) {
|
||||
(true, Some(description)) => Some(format!("{description} (required)")),
|
||||
(true, None) => Some("(required)".to_string()),
|
||||
(false, description) => description,
|
||||
};
|
||||
(
|
||||
format!("{}=", sanitize_display_text(&arg.name)),
|
||||
description,
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
fuzzy_filter(values, |(value, _)| value.as_str(), filter)
|
||||
}
|
||||
|
||||
fn list_prompts_blocking(
|
||||
server: Arc<ConnectedServer>,
|
||||
rpc_timeout: Duration,
|
||||
) -> Option<Vec<Prompt>> {
|
||||
let fut = async move { tokio::time::timeout(rpc_timeout, server.list_all_prompts()).await };
|
||||
// block_in_place is only sound because the REPL's read_line runs inside the
|
||||
// main-thread block_on of the multi-thread runtime.
|
||||
let result = match tokio::runtime::Handle::try_current().ok() {
|
||||
Some(handle) => tokio::task::block_in_place(|| handle.block_on(fut)),
|
||||
None => tokio::runtime::Runtime::new().ok()?.block_on(fut),
|
||||
};
|
||||
|
||||
result.ok()?.ok()
|
||||
}
|
||||
|
||||
fn split_line(line: &str) -> Vec<(&str, usize)> {
|
||||
let mut parts = vec![];
|
||||
let mut part_start = None;
|
||||
@@ -178,3 +265,208 @@ fn test_split_line() {
|
||||
vec![(".set", 0), ("highlight", 5), ("t", 15)],
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod prompt_completion_tests {
|
||||
use super::*;
|
||||
use crate::config::test_fixtures::{FixtureServer, fixture_runtime};
|
||||
use std::sync::atomic::Ordering;
|
||||
|
||||
fn prompts_fixture() -> FixtureServer {
|
||||
FixtureServer {
|
||||
prompts_capability: true,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn stage_two_lists_prompt_names_with_descriptions() {
|
||||
let (runtime, _server) = fixture_runtime(prompts_fixture()).await;
|
||||
let server = runtime.get("fixture").cloned().unwrap();
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::PromptNames { server },
|
||||
"",
|
||||
Duration::from_secs(2),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
values,
|
||||
vec![(
|
||||
"summarize".to_string(),
|
||||
Some("Summarize a document".to_string())
|
||||
)]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn stage_three_suggests_argument_keys_with_required_marker() {
|
||||
let (runtime, _server) = fixture_runtime(prompts_fixture()).await;
|
||||
let server = runtime.get("fixture").cloned().unwrap();
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::ArgumentKeys {
|
||||
server,
|
||||
prompt: "summarize".to_string(),
|
||||
typed_keys: vec![],
|
||||
},
|
||||
"",
|
||||
Duration::from_secs(2),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
values,
|
||||
vec![
|
||||
(
|
||||
"path=".to_string(),
|
||||
Some("Document path (required)".to_string())
|
||||
),
|
||||
("style=".to_string(), None),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn stage_three_excludes_typed_keys_and_fuzzy_filters() {
|
||||
let (runtime, _server) = fixture_runtime(prompts_fixture()).await;
|
||||
let server = runtime.get("fixture").cloned().unwrap();
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::ArgumentKeys {
|
||||
server: Arc::clone(&server),
|
||||
prompt: "summarize".to_string(),
|
||||
typed_keys: vec!["path".to_string()],
|
||||
},
|
||||
"",
|
||||
Duration::from_secs(2),
|
||||
);
|
||||
assert_eq!(values, vec![("style=".to_string(), None)]);
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::ArgumentKeys {
|
||||
server,
|
||||
prompt: "summarize".to_string(),
|
||||
typed_keys: vec![],
|
||||
},
|
||||
"sty",
|
||||
Duration::from_secs(2),
|
||||
);
|
||||
assert_eq!(values, vec![("style=".to_string(), None)]);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn stage_three_unknown_prompt_is_empty() {
|
||||
let (runtime, _server) = fixture_runtime(prompts_fixture()).await;
|
||||
let server = runtime.get("fixture").cloned().unwrap();
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::ArgumentKeys {
|
||||
server,
|
||||
prompt: "ghost".to_string(),
|
||||
typed_keys: vec![],
|
||||
},
|
||||
"",
|
||||
Duration::from_secs(2),
|
||||
);
|
||||
|
||||
assert!(values.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn hostile_prompt_strings_are_sanitized_in_suggestions() {
|
||||
let fixture = FixtureServer {
|
||||
hostile_prompt: true,
|
||||
..prompts_fixture()
|
||||
};
|
||||
let (runtime, _server) = fixture_runtime(fixture).await;
|
||||
let server = runtime.get("fixture").cloned().unwrap();
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::PromptNames {
|
||||
server: Arc::clone(&server),
|
||||
},
|
||||
"evil",
|
||||
Duration::from_secs(2),
|
||||
);
|
||||
assert_eq!(
|
||||
values,
|
||||
vec![(
|
||||
"summarize-evil".to_string(),
|
||||
Some("Runs hostile text".to_string())
|
||||
)]
|
||||
);
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::ArgumentKeys {
|
||||
server,
|
||||
prompt: "sum\u{1b}[31mmarize-evil".to_string(),
|
||||
typed_keys: vec![],
|
||||
},
|
||||
"",
|
||||
Duration::from_secs(2),
|
||||
);
|
||||
assert_eq!(
|
||||
values,
|
||||
vec![("path=".to_string(), Some("Doc path (required)".to_string()))]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn slow_listing_times_out_to_empty() {
|
||||
let fixture = FixtureServer {
|
||||
prompt_delay: Some(Duration::from_millis(200)),
|
||||
..prompts_fixture()
|
||||
};
|
||||
let (runtime, _server) = fixture_runtime(fixture).await;
|
||||
let server = runtime.get("fixture").cloned().unwrap();
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::PromptNames { server },
|
||||
"",
|
||||
Duration::from_millis(20),
|
||||
);
|
||||
|
||||
assert!(values.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn failed_listing_is_swallowed_without_retry() {
|
||||
let fixture = FixtureServer {
|
||||
fail_prompt_listings: true,
|
||||
..prompts_fixture()
|
||||
};
|
||||
let list_prompts_calls = Arc::clone(&fixture.list_prompts_calls);
|
||||
let (runtime, _server) = fixture_runtime(fixture).await;
|
||||
let server = runtime.get("fixture").cloned().unwrap();
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::PromptNames { server },
|
||||
"",
|
||||
Duration::from_secs(2),
|
||||
);
|
||||
|
||||
assert!(values.is_empty());
|
||||
assert_eq!(list_prompts_calls.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bridge_without_ambient_runtime_uses_fallback_runtime() {
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
let (runtime, _server) = rt.block_on(fixture_runtime(prompts_fixture()));
|
||||
let server = runtime.get("fixture").cloned().unwrap();
|
||||
|
||||
let values = complete_prompt_stage(
|
||||
McpPromptCompletion::PromptNames { server },
|
||||
"",
|
||||
Duration::from_secs(2),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
values,
|
||||
vec![(
|
||||
"summarize".to_string(),
|
||||
Some("Summarize a document".to_string())
|
||||
)]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+139
-12
@@ -13,7 +13,7 @@ use crate::client::{
|
||||
};
|
||||
use crate::config::{
|
||||
AgentVariables, AppConfig, AssertState, Input, LastMessage, MacroState, RequestContext,
|
||||
StateFlags, macro_execute,
|
||||
StateFlags, flatten_prompt_messages, macro_execute, resolve_prompt_args, sanitize_display_text,
|
||||
};
|
||||
use crate::config::{AssetCategory, paths};
|
||||
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
|
||||
@@ -29,6 +29,7 @@ use anyhow::{Context, Result, bail};
|
||||
use crossterm::cursor::SetCursorStyle;
|
||||
use fancy_regex::Regex;
|
||||
use indoc::indoc;
|
||||
use inquire::Text;
|
||||
use log::warn;
|
||||
use parking_lot::RwLock;
|
||||
use reedline::CursorConfig;
|
||||
@@ -38,6 +39,7 @@ use reedline::{
|
||||
default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
|
||||
};
|
||||
use reedline::{MenuBuilder, Signal};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::LazyLock;
|
||||
use std::{env, process, sync::Arc};
|
||||
use tokio::task;
|
||||
@@ -53,7 +55,7 @@ pub const DEFAULT_CONTINUATION_PROMPT: &str = indoc! {"
|
||||
4. Continue with the next pending item now. Call tools immediately."
|
||||
};
|
||||
|
||||
static REPL_COMMANDS: LazyLock<[ReplCommand; 61]> = LazyLock::new(|| {
|
||||
static REPL_COMMANDS: LazyLock<[ReplCommand; 62]> = LazyLock::new(|| {
|
||||
[
|
||||
ReplCommand::new(".help", "Show this help guide", AssertState::pass()),
|
||||
ReplCommand::new(".info", "Show system info", AssertState::pass()),
|
||||
@@ -105,6 +107,11 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 61]> = LazyLock::new(|| {
|
||||
ReplCommand::new(".model", "Switch LLM model", AssertState::pass()),
|
||||
ReplCommand::new(
|
||||
".prompt",
|
||||
"Invoke an MCP prompt and submit the result as chat input",
|
||||
AssertState::pass(),
|
||||
),
|
||||
ReplCommand::new(
|
||||
".temp-role",
|
||||
"Set a temporary role using a prompt",
|
||||
AssertState::False(StateFlags::SESSION | StateFlags::AGENT),
|
||||
),
|
||||
@@ -307,7 +314,7 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 61]> = LazyLock::new(|| {
|
||||
),
|
||||
ReplCommand::new(
|
||||
".list",
|
||||
"List roles, sessions, agents, RAGs, macros, skills, tools, MCP servers, or bundles",
|
||||
"List roles, sessions, agents, RAGs, macros, skills, prompts, tools, MCP servers, or bundles",
|
||||
AssertState::pass(),
|
||||
),
|
||||
ReplCommand::new(
|
||||
@@ -774,12 +781,47 @@ pub async fn run_repl_command(
|
||||
.tool disable <name> # Disable a single tool in the current context"#
|
||||
),
|
||||
},
|
||||
".prompt" => match args {
|
||||
".prompt" => {
|
||||
let (words, _) = split_args_text(args.unwrap_or_default(), cfg!(windows));
|
||||
match words.as_slice() {
|
||||
[server, name, rest @ ..] => {
|
||||
let provided = parse_prompt_call_args(rest)?;
|
||||
let prompts = ctx.tool_scope.mcp_runtime.list_prompts(server).await?;
|
||||
let declared = prompts
|
||||
.into_iter()
|
||||
.find(|prompt| prompt.name == *name)
|
||||
.with_context(|| {
|
||||
format!("Prompt '{name}' not found on MCP server '{server}'")
|
||||
})?
|
||||
.arguments
|
||||
.unwrap_or_default();
|
||||
let (mut arguments, missing) = resolve_prompt_args(&declared, provided);
|
||||
for key in missing {
|
||||
let value = Text::new(&prompt_arg_inquire_label(server, name, &key))
|
||||
.prompt()
|
||||
.with_context(|| {
|
||||
format!("Failed to read prompt argument '{key}'")
|
||||
})?;
|
||||
arguments.insert(key, value);
|
||||
}
|
||||
let result = ctx
|
||||
.tool_scope
|
||||
.mcp_runtime
|
||||
.prompt(server, name, arguments)
|
||||
.await?;
|
||||
let flattened = flatten_prompt_messages(&result.messages);
|
||||
let input = Input::from_str(ctx, &flattened, None)?;
|
||||
ask(ctx, abort_signal.clone(), input, true).await?;
|
||||
}
|
||||
_ => println!("Usage: .prompt <server> <name> [key=value ...]"),
|
||||
}
|
||||
}
|
||||
".temp-role" => match args {
|
||||
Some(text) => {
|
||||
let app = Arc::clone(&ctx.app.config);
|
||||
ctx.use_prompt(app.as_ref(), text)?;
|
||||
ctx.use_temp_role(app.as_ref(), text)?;
|
||||
}
|
||||
None => println!("Usage: .prompt <text>..."),
|
||||
None => println!("Usage: .temp-role <text>..."),
|
||||
},
|
||||
".role" => match args {
|
||||
Some(args) => match args.split_once(['\n', ' ']) {
|
||||
@@ -1207,13 +1249,16 @@ pub async fn run_repl_command(
|
||||
println!("Usage: .uninstall <bundle-name> [--yes] (see `.uninstall --help`)")
|
||||
}
|
||||
},
|
||||
".list" => match args {
|
||||
".list" => match args.map(str::trim) {
|
||||
Some("prompts") => {
|
||||
ctx.list_mcp_prompts().await?;
|
||||
}
|
||||
Some(args) => {
|
||||
ctx.list_assets(args.trim())?;
|
||||
ctx.list_assets(args)?;
|
||||
}
|
||||
_ => {
|
||||
println!(
|
||||
"Usage: .list <roles|sessions|agents|rags|macros|skills|tools|mcp-servers|bundles>"
|
||||
"Usage: .list <roles|sessions|agents|rags|macros|skills|prompts|tools|mcp-servers|bundles>"
|
||||
)
|
||||
}
|
||||
},
|
||||
@@ -1737,6 +1782,37 @@ fn split_first_arg(args: Option<&str>) -> Option<(&str, Option<&str>)> {
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_prompt_call_args(words: &[String]) -> Result<HashMap<String, String>> {
|
||||
let mut args = HashMap::new();
|
||||
for word in words {
|
||||
let Some((key, value)) = word.split_once('=') else {
|
||||
bail!("Invalid prompt argument '{word}': arguments must be key=value pairs");
|
||||
};
|
||||
args.insert(key.to_string(), unquote_prompt_value(value).to_string());
|
||||
}
|
||||
Ok(args)
|
||||
}
|
||||
|
||||
fn unquote_prompt_value(value: &str) -> &str {
|
||||
let quoted = value.len() >= 2
|
||||
&& ((value.starts_with('"') && value.ends_with('"'))
|
||||
|| (value.starts_with('\'') && value.ends_with('\'')));
|
||||
if quoted {
|
||||
&value[1..value.len() - 1]
|
||||
} else {
|
||||
value
|
||||
}
|
||||
}
|
||||
|
||||
fn prompt_arg_inquire_label(server: &str, prompt: &str, arg: &str) -> String {
|
||||
format!(
|
||||
"Prompt '{}' on '{}' requires '{}':",
|
||||
sanitize_display_text(prompt),
|
||||
sanitize_display_text(server),
|
||||
sanitize_display_text(arg)
|
||||
)
|
||||
}
|
||||
|
||||
pub fn split_args_text(line: &str, is_win: bool) -> (Vec<String>, &str) {
|
||||
let mut words = Vec::new();
|
||||
let mut word = String::new();
|
||||
@@ -1888,8 +1964,47 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repl_commands_has_61_entries() {
|
||||
assert_eq!(REPL_COMMANDS.len(), 61);
|
||||
fn repl_commands_has_62_entries() {
|
||||
assert_eq!(REPL_COMMANDS.len(), 62);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_prompt_call_args_splits_on_first_equals_and_unquotes() {
|
||||
let words = vec![
|
||||
"path=notes.txt".to_string(),
|
||||
r#"style="a b""#.to_string(),
|
||||
"expr=a=b".to_string(),
|
||||
];
|
||||
|
||||
let args = parse_prompt_call_args(&words).unwrap();
|
||||
|
||||
assert_eq!(args["path"], "notes.txt");
|
||||
assert_eq!(args["style"], "a b");
|
||||
assert_eq!(args["expr"], "a=b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_prompt_call_args_rejects_words_without_equals() {
|
||||
let err = parse_prompt_call_args(&["positional".to_string()])
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
|
||||
assert_eq!(
|
||||
err,
|
||||
"Invalid prompt argument 'positional': arguments must be key=value pairs"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_arg_inquire_label_sanitizes_all_components() {
|
||||
assert_eq!(
|
||||
prompt_arg_inquire_label("srv", "summarize", "path"),
|
||||
"Prompt 'summarize' on 'srv' requires 'path':"
|
||||
);
|
||||
assert_eq!(
|
||||
prompt_arg_inquire_label("s\u{1b}[31mrv", "sum\u{1b}]0;x\u{7}marize", "pa\u{7}th"),
|
||||
"Prompt 'summarize' on 'srv' requires 'pa th':"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -2105,10 +2220,22 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repl_commands_prompt_blocked_in_session_or_agent() {
|
||||
fn repl_commands_prompt_always_available() {
|
||||
let cmd = REPL_COMMANDS.iter().find(|c| c.name == ".prompt").unwrap();
|
||||
assert!(cmd.is_valid(StateFlags::empty()));
|
||||
assert!(cmd.is_valid(StateFlags::ROLE));
|
||||
assert!(cmd.is_valid(StateFlags::SESSION));
|
||||
assert!(cmd.is_valid(StateFlags::AGENT));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn repl_commands_temp_role_blocked_in_session_or_agent() {
|
||||
let cmd = REPL_COMMANDS
|
||||
.iter()
|
||||
.find(|c| c.name == ".temp-role")
|
||||
.unwrap();
|
||||
assert!(cmd.is_valid(StateFlags::empty()));
|
||||
assert!(cmd.is_valid(StateFlags::ROLE));
|
||||
assert!(!cmd.is_valid(StateFlags::SESSION));
|
||||
assert!(!cmd.is_valid(StateFlags::AGENT));
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user