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]
|
[dev-dependencies]
|
||||||
pretty_assertions = "1.4.0"
|
pretty_assertions = "1.4.0"
|
||||||
|
rmcp = { version = "3.1.2", features = ["server"] }
|
||||||
serial_test = "3"
|
serial_test = "3"
|
||||||
|
|
||||||
[[bin]]
|
[[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 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)
|
* [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)
|
* [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.
|
* [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.
|
* 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`).
|
* 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"])
|
.args(["sandbox", "fresh"])
|
||||||
.multiple(true)
|
.multiple(true)
|
||||||
.conflicts_with_all([
|
.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",
|
"macro_name", "execute", "code", "file", "no_stream", "no_memory",
|
||||||
"init_memory", "dry_run", "info", "build_tools", "install",
|
"init_memory", "dry_run", "info", "build_tools", "install",
|
||||||
"install_builtins", "sync_models", "list_models", "list_roles",
|
"install_builtins", "sync_models", "list_models", "list_roles",
|
||||||
@@ -70,9 +70,9 @@ pub struct Cli {
|
|||||||
/// Select a LLM model
|
/// Select a LLM model
|
||||||
#[arg(short, long, add = ArgValueCompleter::new(model_completer))]
|
#[arg(short, long, add = ArgValueCompleter::new(model_completer))]
|
||||||
pub model: Option<String>,
|
pub model: Option<String>,
|
||||||
/// Use the system prompt
|
/// Set a temporary role (an ad-hoc system prompt) for this invocation
|
||||||
#[arg(long)]
|
#[arg(long)]
|
||||||
pub prompt: Option<String>,
|
pub temp_role: Option<String>,
|
||||||
/// Select a role
|
/// Select a role
|
||||||
#[arg(short, long, add = ArgValueCompleter::new(role_completer))]
|
#[arg(short, long, add = ArgValueCompleter::new(role_completer))]
|
||||||
pub role: Option<String>,
|
pub role: Option<String>,
|
||||||
@@ -705,9 +705,9 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parse_prompt_flag() {
|
fn parse_temp_role_flag() {
|
||||||
let cli = parse(&["--prompt", "be a pirate"]);
|
let cli = parse(&["--temp-role", "be a pirate"]);
|
||||||
assert_eq!(cli.prompt, Some("be a pirate".to_string()));
|
assert_eq!(cli.temp_role, Some("be a pirate".to_string()));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
+2
-1
@@ -15,6 +15,7 @@ use crate::config::prompts::{
|
|||||||
};
|
};
|
||||||
use crate::graph::types::RagNode;
|
use crate::graph::types::RagNode;
|
||||||
use crate::graph::{Graph, GraphParser, NodeType};
|
use crate::graph::{Graph, GraphParser, NodeType};
|
||||||
|
use crate::mcp::McpServerFeatures;
|
||||||
use crate::rag::RagInitConfig;
|
use crate::rag::RagInitConfig;
|
||||||
use crate::vault::SECRET_RE;
|
use crate::vault::SECRET_RE;
|
||||||
use anyhow::{Context, Result};
|
use anyhow::{Context, Result};
|
||||||
@@ -380,7 +381,7 @@ impl Agent {
|
|||||||
self.graph_rags.get(node_id).cloned()
|
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);
|
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()))?;
|
let mut functions = Functions::init(config.visible_tools.as_ref().unwrap_or(&Vec::new()))?;
|
||||||
if !mcp_registry.is_empty() && config.mcp_server_support {
|
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() {
|
let mcp_registry = if mcp_registry.is_empty() {
|
||||||
|
|||||||
+66
-2
@@ -442,7 +442,9 @@ impl BundleStore {
|
|||||||
self.ensure_bundle_exists(bundle)?;
|
self.ensure_bundle_exists(bundle)?;
|
||||||
for (name, record) in self.bundles.iter_mut() {
|
for (name, record) in self.bundles.iter_mut() {
|
||||||
if name != bundle {
|
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
|
let record = self
|
||||||
@@ -450,7 +452,9 @@ impl BundleStore {
|
|||||||
.get_mut(bundle)
|
.get_mut(bundle)
|
||||||
.expect("bundle existence checked above");
|
.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);
|
record.files.push(file);
|
||||||
|
|
||||||
self.save()
|
self.save()
|
||||||
@@ -577,6 +581,17 @@ pub(crate) struct BundleListRow {
|
|||||||
pub(crate) drift: DriftSummary,
|
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
|
/// An unreadable file counts as locally modified: it exists but its integrity
|
||||||
/// cannot be verified.
|
/// cannot be verified.
|
||||||
pub(crate) fn bundle_list_rows(store: &BundleStore, config_dir: &Path) -> Vec<BundleListRow> {
|
pub(crate) fn bundle_list_rows(store: &BundleStore, config_dir: &Path) -> Vec<BundleListRow> {
|
||||||
@@ -1489,4 +1504,53 @@ mod tests {
|
|||||||
|
|
||||||
assert!(rows.is_empty());
|
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 {
|
fn is_safe_relative_path(path: &str) -> bool {
|
||||||
let recorded = Path::new(path);
|
let recorded = Path::new(path);
|
||||||
!recorded.is_absolute()
|
!recorded.is_absolute()
|
||||||
&& recorded
|
&& recorded.components().all(|c| match c {
|
||||||
.components()
|
Component::Normal(name) => is_safe_component(&name.to_string_lossy()),
|
||||||
.all(|c| matches!(c, Component::Normal(_)))
|
_ => 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(
|
fn uninstall_owned_files(
|
||||||
@@ -1067,7 +1087,12 @@ impl TempRepoDir {
|
|||||||
|
|
||||||
impl Drop for TempRepoDir {
|
impl Drop for TempRepoDir {
|
||||||
fn drop(&mut self) {
|
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> {
|
fn clone_to_temp(url: &str, reference: Option<&str>) -> Result<TempRepoDir> {
|
||||||
let dest = utils::temp_file("coyote-remote-install-", "");
|
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 dest_arg: OsString = dest.as_os_str().into();
|
||||||
|
|
||||||
let is_sha = reference.is_some_and(is_commit_sha);
|
let is_sha = reference.is_some_and(is_commit_sha);
|
||||||
|
|
||||||
match reference {
|
match reference {
|
||||||
Some(r) if !is_sha => {
|
Some(r) if !is_sha => {
|
||||||
run_git(vec![
|
let mut args = git_content_config();
|
||||||
|
args.extend([
|
||||||
"clone".into(),
|
"clone".into(),
|
||||||
"--depth".into(),
|
"--depth".into(),
|
||||||
"1".into(),
|
"1".into(),
|
||||||
@@ -1094,26 +1143,28 @@ fn clone_to_temp(url: &str, reference: Option<&str>) -> Result<TempRepoDir> {
|
|||||||
"--".into(),
|
"--".into(),
|
||||||
url.into(),
|
url.into(),
|
||||||
dest_arg,
|
dest_arg,
|
||||||
])?;
|
]);
|
||||||
|
run_git(args)?;
|
||||||
}
|
}
|
||||||
Some(r) => {
|
Some(r) => {
|
||||||
run_git(vec![
|
let mut args = git_content_config();
|
||||||
"clone".into(),
|
args.extend(["clone".into(), "--".into(), url.into(), dest_arg.clone()]);
|
||||||
"--".into(),
|
run_git(args)?;
|
||||||
url.into(),
|
let mut args = git_content_config();
|
||||||
dest_arg.clone(),
|
args.extend(["-C".into(), dest_arg, "checkout".into(), r.into()]);
|
||||||
])?;
|
run_git(args)?;
|
||||||
run_git(vec!["-C".into(), dest_arg, "checkout".into(), r.into()])?;
|
|
||||||
}
|
}
|
||||||
None => {
|
None => {
|
||||||
run_git(vec![
|
let mut args = git_content_config();
|
||||||
|
args.extend([
|
||||||
"clone".into(),
|
"clone".into(),
|
||||||
"--depth".into(),
|
"--depth".into(),
|
||||||
"1".into(),
|
"1".into(),
|
||||||
"--".into(),
|
"--".into(),
|
||||||
url.into(),
|
url.into(),
|
||||||
dest_arg,
|
dest_arg,
|
||||||
])?;
|
]);
|
||||||
|
run_git(args)?;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1123,11 +1174,7 @@ fn clone_to_temp(url: &str, reference: Option<&str>) -> Result<TempRepoDir> {
|
|||||||
"rev-parse".into(),
|
"rev-parse".into(),
|
||||||
"HEAD".into(),
|
"HEAD".into(),
|
||||||
])?;
|
])?;
|
||||||
|
Ok(head_sha)
|
||||||
Ok(TempRepoDir {
|
|
||||||
path: dest,
|
|
||||||
head_sha,
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn run_git(args: Vec<OsString>) -> Result<()> {
|
fn run_git(args: Vec<OsString>) -> Result<()> {
|
||||||
@@ -1816,7 +1863,17 @@ fn record_written_file(
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn provenance_path(dst: &Path) -> String {
|
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('\\', "/")
|
rel.to_string_lossy().replace('\\', "/")
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -2242,6 +2299,53 @@ mod tests {
|
|||||||
use std::env;
|
use std::env;
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
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 {
|
struct TestVaultConfigGuard {
|
||||||
dir_key: String,
|
dir_key: String,
|
||||||
file_key: String,
|
file_key: String,
|
||||||
@@ -4286,6 +4390,12 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
#[serial]
|
#[serial]
|
||||||
fn uninstall_shorthand_with_multiple_matches_bails_non_interactively() {
|
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 _guard = TestVaultConfigGuard::new("uninst-short-multi");
|
||||||
let mut store = BundleStore::load().unwrap();
|
let mut store = BundleStore::load().unwrap();
|
||||||
store
|
store
|
||||||
|
|||||||
@@ -287,6 +287,7 @@ fn discover_macros_in(dirs: &[(MacroSource, PathBuf)]) -> Vec<DiscoveredMacro> {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::repl;
|
||||||
use crate::utils::get_env_name;
|
use crate::utils::get_env_name;
|
||||||
use serial_test::serial;
|
use serial_test::serial;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
@@ -565,6 +566,34 @@ mod tests {
|
|||||||
assert_eq!(state_of(&policy, "a"), &MacroState::Enabled);
|
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]
|
#[test]
|
||||||
fn locked_wins_over_shadowed_builtin() {
|
fn locked_wins_over_shadowed_builtin() {
|
||||||
let l = list(&["a"]);
|
let l = list(&["a"]);
|
||||||
|
|||||||
@@ -22,6 +22,8 @@ pub(crate) mod todo;
|
|||||||
mod tool_scope;
|
mod tool_scope;
|
||||||
mod update;
|
mod update;
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
pub(crate) use self::agent::AgentConfig;
|
||||||
pub use self::agent::{
|
pub use self::agent::{
|
||||||
Agent, AgentVariable, AgentVariables, complete_agent_variables, list_agents,
|
Agent, AgentVariable, AgentVariables, complete_agent_variables, list_agents,
|
||||||
list_agents_with_descriptions,
|
list_agents_with_descriptions,
|
||||||
@@ -51,6 +53,11 @@ pub use self::skill::Skill;
|
|||||||
pub use self::skill_policy::SkillPolicy;
|
pub use self::skill_policy::SkillPolicy;
|
||||||
#[allow(unused_imports)]
|
#[allow(unused_imports)]
|
||||||
pub use self::skill_registry::SkillRegistry;
|
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;
|
pub use self::update::run_self_update;
|
||||||
use crate::client::{
|
use crate::client::{
|
||||||
self, ClientConfig, MessageContentToolCalls, Model, ModelType, OPENAI_COMPATIBLE_PROVIDERS,
|
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_policy::SkillPolicy;
|
||||||
use super::skill_registry::SkillRegistry;
|
use super::skill_registry::SkillRegistry;
|
||||||
use super::todo::TodoList;
|
use super::todo::TodoList;
|
||||||
use super::tool_scope::{McpRuntime, ToolScope};
|
use super::tool_scope::{McpPromptCompletion, McpRuntime, ToolScope, format_prompt_arguments};
|
||||||
use super::{
|
use super::{
|
||||||
AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, AssetCategory, CREATE_TITLE_ROLE,
|
AGENTS_DIR_NAME, Agent, AgentVariables, AppConfig, AppState, AssetCategory, CREATE_TITLE_ROLE,
|
||||||
Input, InstallFilter, LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, MacroAllowlistLevel,
|
Input, InstallFilter, LEFT_PROMPT, LastMessage, MESSAGES_FILE_NAME, MacroAllowlistLevel,
|
||||||
@@ -23,8 +23,8 @@ use crate::function::{
|
|||||||
user_interaction::USER_FUNCTION_PREFIX,
|
user_interaction::USER_FUNCTION_PREFIX,
|
||||||
};
|
};
|
||||||
use crate::mcp::{
|
use crate::mcp::{
|
||||||
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
CatalogItem, MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired,
|
||||||
MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpAuthReason, McpAuthRequired, is_auth_required_error,
|
McpServersConfig, is_auth_required_error, is_mcp_meta_function, mcp_meta_function_names,
|
||||||
};
|
};
|
||||||
use crate::rag::Rag;
|
use crate::rag::Rag;
|
||||||
use crate::supervisor::Supervisor;
|
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 struct AutoContinueConfig {
|
||||||
pub enabled: bool,
|
pub enabled: bool,
|
||||||
pub max_continues: usize,
|
pub max_continues: usize,
|
||||||
@@ -139,6 +164,20 @@ pub(crate) fn asset_table(header: &[&str]) -> Table {
|
|||||||
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>)> {
|
fn complete_skills_with_descriptions(names: Vec<String>) -> Vec<(String, Option<String>)> {
|
||||||
names
|
names
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@@ -2011,11 +2050,7 @@ impl RequestContext {
|
|||||||
.functions
|
.functions
|
||||||
.declarations()
|
.declarations()
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|v| {
|
.filter(|v| !is_mcp_meta_function(&v.name))
|
||||||
!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)
|
|
||||||
})
|
|
||||||
.map(|v| v.name.to_string())
|
.map(|v| v.name.to_string())
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
@@ -2025,11 +2060,7 @@ impl RequestContext {
|
|||||||
.functions()
|
.functions()
|
||||||
.declarations()
|
.declarations()
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|v| {
|
.filter(|v| !is_mcp_meta_function(&v.name))
|
||||||
!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)
|
|
||||||
})
|
|
||||||
.map(|v| v.name.to_string()),
|
.map(|v| v.name.to_string()),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -2102,11 +2133,7 @@ impl RequestContext {
|
|||||||
.declarations()
|
.declarations()
|
||||||
.to_vec()
|
.to_vec()
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.filter(|v| {
|
.filter(|v| !is_mcp_meta_function(&v.name))
|
||||||
!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)
|
|
||||||
})
|
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
if let Some(ref tool_names) = role_filter {
|
if let Some(ref tool_names) = role_filter {
|
||||||
@@ -2155,11 +2182,7 @@ impl RequestContext {
|
|||||||
.functions
|
.functions
|
||||||
.declarations()
|
.declarations()
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|v| {
|
.filter(|v| is_mcp_meta_function(&v.name))
|
||||||
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)
|
|
||||||
})
|
|
||||||
.map(|v| v.name.to_string())
|
.map(|v| v.name.to_string())
|
||||||
.collect();
|
.collect();
|
||||||
if let Some(agent) = &self.agent {
|
if let Some(agent) = &self.agent {
|
||||||
@@ -2168,12 +2191,7 @@ impl RequestContext {
|
|||||||
.functions()
|
.functions()
|
||||||
.declarations()
|
.declarations()
|
||||||
.iter()
|
.iter()
|
||||||
.filter(|v| {
|
.filter(|v| is_mcp_meta_function(&v.name))
|
||||||
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)
|
|
||||||
})
|
|
||||||
.map(|v| v.name.to_string()),
|
.map(|v| v.name.to_string()),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -2188,41 +2206,17 @@ impl RequestContext {
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
let item_invoke_name =
|
|
||||||
format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX);
|
|
||||||
let item_search_name =
|
let item_search_name =
|
||||||
format!("{}_{item}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX);
|
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) {
|
if let Some(values) = app.mapping_mcp_servers.get(item) {
|
||||||
server_names.extend(
|
server_names.extend(
|
||||||
values
|
values
|
||||||
.split(',')
|
.split(',')
|
||||||
.flat_map(|v| {
|
.flat_map(mcp_meta_function_names)
|
||||||
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()
|
|
||||||
),
|
|
||||||
]
|
|
||||||
})
|
|
||||||
.filter(|v| mcp_declaration_names.contains(v)),
|
.filter(|v| mcp_declaration_names.contains(v)),
|
||||||
)
|
)
|
||||||
} else if mcp_declaration_names.contains(&item_invoke_name) {
|
} else if mcp_declaration_names.contains(&item_search_name) {
|
||||||
server_names.insert(item_invoke_name);
|
server_names.extend(mcp_meta_function_names(item));
|
||||||
server_names.insert(item_search_name);
|
|
||||||
server_names.insert(item_describe_name);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -2251,11 +2245,7 @@ impl RequestContext {
|
|||||||
.declarations()
|
.declarations()
|
||||||
.to_vec()
|
.to_vec()
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.filter(|v| {
|
.filter(|v| is_mcp_meta_function(&v.name))
|
||||||
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)
|
|
||||||
})
|
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
if let Some(ref server_names) = role_filter {
|
if let Some(ref server_names) = role_filter {
|
||||||
@@ -2351,7 +2341,7 @@ impl RequestContext {
|
|||||||
Ok(())
|
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);
|
let mut role = Role::new(TEMP_ROLE_NAME, prompt);
|
||||||
role.set_model(self.current_model().clone());
|
role.set_model(self.current_model().clone());
|
||||||
self.use_role_obj(role)
|
self.use_role_obj(role)
|
||||||
@@ -2828,7 +2818,7 @@ impl RequestContext {
|
|||||||
}
|
}
|
||||||
"bundles" => bundles::list_installed_bundles(),
|
"bundles" => bundles::list_installed_bundles(),
|
||||||
_ => bail!(
|
_ => 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",
|
"rags",
|
||||||
"macros",
|
"macros",
|
||||||
"skills",
|
"skills",
|
||||||
|
"prompts",
|
||||||
"tools",
|
"tools",
|
||||||
"mcp-servers",
|
"mcp-servers",
|
||||||
"bundles",
|
"bundles",
|
||||||
@@ -3720,6 +3711,40 @@ impl RequestContext {
|
|||||||
fuzzy_filter(values, |v| v.0.as_str(), filter)
|
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(
|
async fn rebuild_tool_scope(
|
||||||
&mut self,
|
&mut self,
|
||||||
app: &AppConfig,
|
app: &AppConfig,
|
||||||
@@ -3763,24 +3788,7 @@ impl RequestContext {
|
|||||||
&& let Some(mcp_config) = &self.app.mcp_config
|
&& let Some(mcp_config) = &self.app.mcp_config
|
||||||
{
|
{
|
||||||
let server_ids: Vec<String> = match &enabled_mcp_servers {
|
let server_ids: Vec<String> = match &enabled_mcp_servers {
|
||||||
Some(servers) if servers.iter().any(|s| s.trim() == "all") => {
|
Some(servers) => expand_enabled_mcp_server_ids(app, mcp_config, servers),
|
||||||
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
|
|
||||||
}
|
|
||||||
None => vec![],
|
None => vec![],
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -3836,7 +3844,7 @@ impl RequestContext {
|
|||||||
functions.append_todo_functions();
|
functions.append_todo_functions();
|
||||||
}
|
}
|
||||||
if !mcp_runtime.is_empty() {
|
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 {
|
if app.function_calling_support && policy.skills_enabled {
|
||||||
functions.append_skill_functions();
|
functions.append_skill_functions();
|
||||||
@@ -4685,10 +4693,11 @@ mod tests {
|
|||||||
use crate::config::AppState;
|
use crate::config::AppState;
|
||||||
use crate::config::agent::AgentConfig;
|
use crate::config::agent::AgentConfig;
|
||||||
use crate::function::{ToolCall, skill};
|
use crate::function::{ToolCall, skill};
|
||||||
use crate::mcp::{McpServer, McpServersConfig, McpTransportType};
|
use crate::mcp::{McpServer, McpServerFeatures, McpServersConfig, McpTransportType};
|
||||||
use crate::utils;
|
use crate::utils;
|
||||||
use crate::utils::get_env_name;
|
use crate::utils::get_env_name;
|
||||||
use crate::vault::Vault;
|
use crate::vault::Vault;
|
||||||
|
use rmcp::model::PromptArgument;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use serial_test::serial;
|
use serial_test::serial;
|
||||||
use std::env;
|
use std::env;
|
||||||
@@ -4746,6 +4755,15 @@ mod tests {
|
|||||||
RequestContext::new(default_app_state(), WorkingMode::Cmd)
|
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]
|
#[test]
|
||||||
fn new_creates_clean_state() {
|
fn new_creates_clean_state() {
|
||||||
let ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd);
|
let ctx = RequestContext::new(default_app_state(), WorkingMode::Cmd);
|
||||||
@@ -4834,10 +4852,10 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn use_prompt_creates_temp_role() {
|
fn use_temp_role_creates_temp_role() {
|
||||||
let mut ctx = create_test_ctx();
|
let mut ctx = create_test_ctx();
|
||||||
let app = ctx.app.config.clone();
|
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!(ctx.role.is_some());
|
||||||
assert_eq!(ctx.role.as_ref().unwrap().name(), "temp");
|
assert_eq!(ctx.role.as_ref().unwrap().name(), "temp");
|
||||||
assert!(
|
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]
|
#[test]
|
||||||
fn extract_role_returns_standalone_role() {
|
fn extract_role_returns_standalone_role() {
|
||||||
let mut ctx = create_test_ctx();
|
let mut ctx = create_test_ctx();
|
||||||
@@ -5753,9 +5806,10 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn select_enabled_mcp_servers_all_returns_all_mcp_functions() {
|
fn select_enabled_mcp_servers_all_returns_all_mcp_functions() {
|
||||||
let mut ctx = create_test_ctx();
|
let mut ctx = create_test_ctx();
|
||||||
ctx.tool_scope
|
ctx.tool_scope.functions.append_mcp_meta_functions(vec![
|
||||||
.functions
|
tools_only_features("github"),
|
||||||
.append_mcp_meta_functions(vec!["github".into(), "slack".into()]);
|
tools_only_features("slack"),
|
||||||
|
]);
|
||||||
|
|
||||||
let mut role = Role::new("r", "p");
|
let mut role = Role::new("r", "p");
|
||||||
role.set_enabled_mcp_servers(Some(vec!["all".to_string()]));
|
role.set_enabled_mcp_servers(Some(vec!["all".to_string()]));
|
||||||
@@ -5770,9 +5824,10 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn select_enabled_mcp_servers_comma_filters() {
|
fn select_enabled_mcp_servers_comma_filters() {
|
||||||
let mut ctx = create_test_ctx();
|
let mut ctx = create_test_ctx();
|
||||||
ctx.tool_scope
|
ctx.tool_scope.functions.append_mcp_meta_functions(vec![
|
||||||
.functions
|
tools_only_features("github"),
|
||||||
.append_mcp_meta_functions(vec!["github".into(), "slack".into()]);
|
tools_only_features("slack"),
|
||||||
|
]);
|
||||||
|
|
||||||
let mut role = Role::new("r", "p");
|
let mut role = Role::new("r", "p");
|
||||||
role.set_enabled_mcp_servers(Some(vec!["github".to_string()]));
|
role.set_enabled_mcp_servers(Some(vec!["github".to_string()]));
|
||||||
@@ -5783,6 +5838,28 @@ mod tests {
|
|||||||
assert!(!names.contains(&"mcp_invoke_slack"));
|
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]
|
#[test]
|
||||||
fn state_empty_context_has_no_context_flags() {
|
fn state_empty_context_has_no_context_flags() {
|
||||||
let ctx = create_test_ctx();
|
let ctx = create_test_ctx();
|
||||||
|
|||||||
+1488
-34
File diff suppressed because it is too large
Load Diff
+1568
-71
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<()> {
|
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 functions = {
|
||||||
let agent = ctx
|
let agent = ctx
|
||||||
.agent
|
.agent
|
||||||
.as_mut()
|
.as_mut()
|
||||||
.with_context(|| "Agent should be initialized")?;
|
.with_context(|| "Agent should be initialized")?;
|
||||||
if !server_names.is_empty() {
|
if !server_features.is_empty() {
|
||||||
agent.append_mcp_meta_functions(server_names);
|
agent.append_mcp_meta_functions(server_features);
|
||||||
}
|
}
|
||||||
agent.functions().clone()
|
agent.functions().clone()
|
||||||
};
|
};
|
||||||
@@ -1453,7 +1453,8 @@ async fn summarize_output(ctx: &RequestContext, agent_name: &str, output: &str)
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
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 crate::supervisor::escalation::{EscalationQueue, EscalationRequest};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use serial_test::serial;
|
use serial_test::serial;
|
||||||
@@ -1510,6 +1511,28 @@ mod tests {
|
|||||||
.block_on(f)
|
.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]
|
#[test]
|
||||||
fn handle_list_running_empty_supervisor() {
|
fn handle_list_running_empty_supervisor() {
|
||||||
let mut ctx = ctx_with_supervisor(4, 3);
|
let mut ctx = ctx_with_supervisor(4, 3);
|
||||||
|
|||||||
+2
-2
@@ -380,8 +380,8 @@ async fn run(
|
|||||||
.await?;
|
.await?;
|
||||||
} else {
|
} else {
|
||||||
let app: Arc<AppConfig> = Arc::clone(&ctx.app.config);
|
let app: Arc<AppConfig> = Arc::clone(&ctx.app.config);
|
||||||
if let Some(prompt) = &cli.prompt {
|
if let Some(prompt) = &cli.temp_role {
|
||||||
ctx.use_prompt(app.as_ref(), prompt)?;
|
ctx.use_temp_role(app.as_ref(), prompt)?;
|
||||||
} else if let Some(name) = &cli.role {
|
} else if let Some(name) = &cli.role {
|
||||||
ctx.use_role(app.as_ref(), name, abort_signal.clone())
|
ctx.use_role(app.as_ref(), name, abort_signal.clone())
|
||||||
.await?;
|
.await?;
|
||||||
|
|||||||
+151
-47
@@ -1,6 +1,7 @@
|
|||||||
mod auth_client;
|
mod auth_client;
|
||||||
pub(crate) mod manage;
|
pub(crate) mod manage;
|
||||||
pub(crate) mod oauth;
|
pub(crate) mod oauth;
|
||||||
|
pub(crate) mod render;
|
||||||
mod sse_transport;
|
mod sse_transport;
|
||||||
|
|
||||||
use crate::config::AppConfig;
|
use crate::config::AppConfig;
|
||||||
@@ -15,6 +16,7 @@ use futures_util::{StreamExt, TryStreamExt, stream};
|
|||||||
use http::{HeaderName, HeaderValue};
|
use http::{HeaderName, HeaderValue};
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
use indoc::formatdoc;
|
use indoc::formatdoc;
|
||||||
|
use rmcp::model::{PromptArgument, ServerCapabilities};
|
||||||
use rmcp::service::RunningService;
|
use rmcp::service::RunningService;
|
||||||
use rmcp::transport::StreamableHttpClientTransport;
|
use rmcp::transport::StreamableHttpClientTransport;
|
||||||
use rmcp::transport::TokioChildProcess;
|
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_INVOKE_META_FUNCTION_NAME_PREFIX: &str = "mcp_invoke";
|
||||||
pub const MCP_SEARCH_META_FUNCTION_NAME_PREFIX: &str = "mcp_search";
|
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_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, ()>;
|
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)]
|
#[derive(Clone, Debug, Default, Serialize)]
|
||||||
pub struct CatalogItem {
|
pub struct CatalogItem {
|
||||||
|
pub kind: CatalogItemKind,
|
||||||
pub name: String,
|
pub name: String,
|
||||||
pub server: String,
|
pub server: String,
|
||||||
pub description: String,
|
pub description: String,
|
||||||
}
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
|
pub uri: Option<String>,
|
||||||
#[derive(Debug)]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
struct ServerCatalog {
|
pub mime_type: Option<String>,
|
||||||
items: HashMap<String, CatalogItem>,
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
}
|
pub size: Option<u64>,
|
||||||
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
impl Clone for ServerCatalog {
|
pub arguments: Option<Vec<PromptArgument>>,
|
||||||
fn clone(&self) -> Self {
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
Self {
|
pub audience: Option<Vec<String>>,
|
||||||
items: self.items.clone(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||||
@@ -160,7 +232,6 @@ pub struct McpRegistry {
|
|||||||
log_path: Option<PathBuf>,
|
log_path: Option<PathBuf>,
|
||||||
config: Option<McpServersConfig>,
|
config: Option<McpServersConfig>,
|
||||||
servers: HashMap<String, Arc<ConnectedServer>>,
|
servers: HashMap<String, Arc<ConnectedServer>>,
|
||||||
catalogs: HashMap<String, ServerCatalog>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl McpRegistry {
|
impl McpRegistry {
|
||||||
@@ -303,7 +374,7 @@ impl McpRegistry {
|
|||||||
|
|
||||||
debug!("Starting selected MCP servers: {:?}", ids_to_start);
|
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
|
ids_to_start
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|id| async { self.start_server(id).await }),
|
.map(|id| async { self.start_server(id).await }),
|
||||||
@@ -312,18 +383,14 @@ impl McpRegistry {
|
|||||||
.try_collect()
|
.try_collect()
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
for (id, server, catalog) in results.into_iter().flatten() {
|
for (id, server) in results.into_iter().flatten() {
|
||||||
self.servers.insert(id.clone(), server);
|
self.servers.insert(id, server);
|
||||||
self.catalogs.insert(id, catalog);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn start_server(
|
async fn start_server(&self, id: String) -> Result<Option<(String, Arc<ConnectedServer>)>> {
|
||||||
&self,
|
|
||||||
id: String,
|
|
||||||
) -> Result<Option<(String, Arc<ConnectedServer>, ServerCatalog)>> {
|
|
||||||
let spec = self
|
let spec = self
|
||||||
.config
|
.config
|
||||||
.as_ref()
|
.as_ref()
|
||||||
@@ -347,30 +414,9 @@ impl McpRegistry {
|
|||||||
Err(e) => return Err(e),
|
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}");
|
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> {
|
fn resolve_server_ids(&self, enabled_mcp_servers: Option<Vec<String>>) -> Vec<String> {
|
||||||
@@ -398,8 +444,21 @@ impl McpRegistry {
|
|||||||
&self.servers
|
&self.servers
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn list_started_servers(&self) -> Vec<String> {
|
pub fn server_features(&self) -> Vec<McpServerFeatures> {
|
||||||
self.servers.keys().cloned().collect()
|
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 {
|
pub fn is_empty(&self) -> bool {
|
||||||
@@ -1161,7 +1220,7 @@ mod tests {
|
|||||||
let registry = McpRegistry::default();
|
let registry = McpRegistry::default();
|
||||||
|
|
||||||
assert!(registry.is_empty());
|
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.mcp_config().is_none());
|
||||||
assert!(registry.log_path().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_INVOKE_META_FUNCTION_NAME_PREFIX, "mcp_invoke");
|
||||||
assert_eq!(MCP_SEARCH_META_FUNCTION_NAME_PREFIX, "mcp_search");
|
assert_eq!(MCP_SEARCH_META_FUNCTION_NAME_PREFIX, "mcp_search");
|
||||||
assert_eq!(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, "mcp_describe");
|
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]
|
#[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 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 parking_lot::RwLock;
|
||||||
use reedline::{Completer, Span, Suggestion};
|
use reedline::{Completer, Span, Suggestion};
|
||||||
|
use rmcp::model::Prompt;
|
||||||
use std::collections::HashMap;
|
use std::collections::HashMap;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
|
const PROMPT_COMPLETION_RPC_TIMEOUT: Duration = Duration::from_secs(2);
|
||||||
|
|
||||||
impl Completer for ReplCompleter {
|
impl Completer for ReplCompleter {
|
||||||
fn complete(&mut self, line: &str, pos: usize) -> Vec<Suggestion> {
|
fn complete(&mut self, line: &str, pos: usize) -> Vec<Suggestion> {
|
||||||
@@ -29,6 +35,22 @@ impl Completer for ReplCompleter {
|
|||||||
return suggestions;
|
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 ctx = self.ctx.read();
|
||||||
let state = ctx.state();
|
let state = ctx.state();
|
||||||
let model_has_reasoning = !ctx.current_model().reasoning_levels().is_empty();
|
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)> {
|
fn split_line(line: &str) -> Vec<(&str, usize)> {
|
||||||
let mut parts = vec![];
|
let mut parts = vec![];
|
||||||
let mut part_start = None;
|
let mut part_start = None;
|
||||||
@@ -178,3 +265,208 @@ fn test_split_line() {
|
|||||||
vec![(".set", 0), ("highlight", 5), ("t", 15)],
|
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::{
|
use crate::config::{
|
||||||
AgentVariables, AppConfig, AssertState, Input, LastMessage, MacroState, RequestContext,
|
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::config::{AssetCategory, paths};
|
||||||
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
|
use crate::function::supervisor::{GuardrailAction, check_pending_agents_guardrail};
|
||||||
@@ -29,6 +29,7 @@ use anyhow::{Context, Result, bail};
|
|||||||
use crossterm::cursor::SetCursorStyle;
|
use crossterm::cursor::SetCursorStyle;
|
||||||
use fancy_regex::Regex;
|
use fancy_regex::Regex;
|
||||||
use indoc::indoc;
|
use indoc::indoc;
|
||||||
|
use inquire::Text;
|
||||||
use log::warn;
|
use log::warn;
|
||||||
use parking_lot::RwLock;
|
use parking_lot::RwLock;
|
||||||
use reedline::CursorConfig;
|
use reedline::CursorConfig;
|
||||||
@@ -38,6 +39,7 @@ use reedline::{
|
|||||||
default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
|
default_emacs_keybindings, default_vi_insert_keybindings, default_vi_normal_keybindings,
|
||||||
};
|
};
|
||||||
use reedline::{MenuBuilder, Signal};
|
use reedline::{MenuBuilder, Signal};
|
||||||
|
use std::collections::HashMap;
|
||||||
use std::sync::LazyLock;
|
use std::sync::LazyLock;
|
||||||
use std::{env, process, sync::Arc};
|
use std::{env, process, sync::Arc};
|
||||||
use tokio::task;
|
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."
|
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(".help", "Show this help guide", AssertState::pass()),
|
||||||
ReplCommand::new(".info", "Show system info", 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(".model", "Switch LLM model", AssertState::pass()),
|
||||||
ReplCommand::new(
|
ReplCommand::new(
|
||||||
".prompt",
|
".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",
|
"Set a temporary role using a prompt",
|
||||||
AssertState::False(StateFlags::SESSION | StateFlags::AGENT),
|
AssertState::False(StateFlags::SESSION | StateFlags::AGENT),
|
||||||
),
|
),
|
||||||
@@ -307,7 +314,7 @@ static REPL_COMMANDS: LazyLock<[ReplCommand; 61]> = LazyLock::new(|| {
|
|||||||
),
|
),
|
||||||
ReplCommand::new(
|
ReplCommand::new(
|
||||||
".list",
|
".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(),
|
AssertState::pass(),
|
||||||
),
|
),
|
||||||
ReplCommand::new(
|
ReplCommand::new(
|
||||||
@@ -774,12 +781,47 @@ pub async fn run_repl_command(
|
|||||||
.tool disable <name> # Disable a single tool in the current context"#
|
.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) => {
|
Some(text) => {
|
||||||
let app = Arc::clone(&ctx.app.config);
|
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 {
|
".role" => match args {
|
||||||
Some(args) => match args.split_once(['\n', ' ']) {
|
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`)")
|
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) => {
|
Some(args) => {
|
||||||
ctx.list_assets(args.trim())?;
|
ctx.list_assets(args)?;
|
||||||
}
|
}
|
||||||
_ => {
|
_ => {
|
||||||
println!(
|
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) {
|
pub fn split_args_text(line: &str, is_win: bool) -> (Vec<String>, &str) {
|
||||||
let mut words = Vec::new();
|
let mut words = Vec::new();
|
||||||
let mut word = String::new();
|
let mut word = String::new();
|
||||||
@@ -1888,8 +1964,47 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn repl_commands_has_61_entries() {
|
fn repl_commands_has_62_entries() {
|
||||||
assert_eq!(REPL_COMMANDS.len(), 61);
|
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]
|
#[test]
|
||||||
@@ -2105,10 +2220,22 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[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();
|
let cmd = REPL_COMMANDS.iter().find(|c| c.name == ".prompt").unwrap();
|
||||||
assert!(cmd.is_valid(StateFlags::empty()));
|
assert!(cmd.is_valid(StateFlags::empty()));
|
||||||
assert!(cmd.is_valid(StateFlags::ROLE));
|
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::SESSION));
|
||||||
assert!(!cmd.is_valid(StateFlags::AGENT));
|
assert!(!cmd.is_valid(StateFlags::AGENT));
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user