MythEclipse/zesdex
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: implement subagent tool gating and timeout mechanisms for enhanced security and performance

Browse Source
This commit is contained in:
asepharyana
2026-07-13 04:41:26 +07:00
parent d09e440e7e
commit 3b711bbf3b
7 changed files with 379 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/app/subagent/engine.rs
+240 -7
View File
@@ -1,6 +1,11 @@
//! Subagent execution loop: drive an LLM conversation, gate tool calls
//! against the context's allowlist, run tools, and stream progress events
//! to the parent via an mpsc channel.
//!
//! Security: subagent tool gating mirrors the main agent's `Harness` checks
//! (path traversal, reason validation, stub/denial/assumption scanning,
//! bash exfiltration and destructive-pattern detection) so that subagents
//! are not a weaker link than the main agent.
use tokio::sync::mpsc;
use crate::dto::chat::message::ChatMessage;
@@ -73,24 +78,236 @@ fn resolve_provider_config() -> (String, String, Option<String>) {
(api_key, model, base_url)
}
// ─── Subagent-level tool gating (mirrors Harness checks) ───
const STUB_PATTERNS: &[&str] = &[
"todo!()", "todo!(",
"unimplemented!()", "unimplemented!(",
"FIXME", "fixme:", "XXX:", "PLACEHOLDER",
"REPLACE_ME", "stub_value", "stub_function",
"fake_response", "fake_data",
"not implemented", "not yet implemented",
"to be implemented", "to be done",
];
const DENIAL_PATTERNS: &[&str] = &[
"// skip", "// skipping", "// skipping for now",
"// for now just", "// punt", "// hack:",
"// workaround:", "// cba", "// later",
"// do later", "// ignore for now", "// disable",
"// bypass", "// quick fix", "// temp fix",
"// temporary fix", "// temp:", "// temporary:",
"// noop",
];
const ASSUMPTION_PATTERNS: &[&str] = &[
"// assume", "// probably", "// guess",
"// should work", "// hopefully", "// i think",
"// should be fine", "// likely",
];
const EXFIL_PATTERNS: &[&str] = &[
"curl ", "wget ", "nc -e ", "ncat ", "/dev/tcp/",
"base64 -d |", "base64 --decode |",
"openssl s_client", "ssh -R ",
"scp /", "rsync /",
];
const SENSITIVE_PATH_PATTERNS: &[&str] = &[
".ssh/id_rsa", ".ssh/id_ed25519",
".aws/credentials", ".aws/config",
".kube/config", ".docker/config.json",
"/etc/shadow", "/etc/passwd", "/proc/self/environ",
];
const MIN_REASON_LEN: usize = 8;
/// Gate a tool call in the subagent context. Returns `Some(block_reason)` if
/// the call should be blocked, `None` to allow.
///
/// Flow: always blocks dangerous patterns — path traversal, stub/denial/
/// assumption language, bash exfiltration, destructive commands, sensitive
/// path reads — regardless of the allowed-tools list. Tools that are not
/// risky only get the basic allowlist check.
fn gate_subagent_tool_call(
tool_name: &str,
args: &serde_json::Value,
) -> Option<String> {
// File-mutating tools: write / edit / delete
if matches!(tool_name, "write" | "edit" | "delete") {
if let Some(path) = args.get("path").and_then(|v| v.as_str()) {
if path.contains("..") {
return Some("path traversal detected in 'path' argument".to_string());
}
}
}
// write / edit require a non-trivial `reason`
if matches!(tool_name, "write" | "edit" | "delete") {
let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("");
if reason.trim().len() < MIN_REASON_LEN {
return Some(format!(
"{} requires a non-trivial 'reason' (>= {} chars) explaining why",
tool_name, MIN_REASON_LEN,
));
}
}
// write / edit content must not contain stubs, denial, or assumption language
if matches!(tool_name, "write" | "edit") {
let content = match tool_name {
"write" => args.get("content").and_then(|v| v.as_str()).unwrap_or(""),
"edit" => {
let old = args.get("old").and_then(|v| v.as_str()).unwrap_or("");
let new = args.get("new").and_then(|v| v.as_str()).unwrap_or("");
// For edits, scanning old+new together catches stubs in both
return if contains_any(old, STUB_PATTERNS) {
Some("content contains stub/placeholder pattern; production code must be fully implemented".to_string())
} else if contains_any(new, STUB_PATTERNS) {
Some("content contains stub/placeholder pattern; production code must be fully implemented".to_string())
} else if contains_any(new, DENIAL_PATTERNS) {
Some("content contains denial/punt pattern; implement properly instead of skipping".to_string())
} else if contains_any(new, ASSUMPTION_PATTERNS) {
Some("content contains assumption pattern; verify against data instead of guessing".to_string())
} else {
return None;
};
}
_ => "",
};
if contains_any(content, STUB_PATTERNS) {
return Some("content contains stub/placeholder pattern; production code must be fully implemented".to_string());
}
if contains_any(content, DENIAL_PATTERNS) {
return Some("content contains denial/punt pattern; implement properly instead of skipping".to_string());
}
if contains_any(content, ASSUMPTION_PATTERNS) {
return Some("content contains assumption pattern; verify against data instead of guessing".to_string());
}
}
// Bash: exfiltration, sensitive paths, destructive commands
if tool_name == "bash" {
let cmd = args.get("command").and_then(|v| v.as_str()).unwrap_or("");
if cmd.contains("..") {
return Some("path traversal detected in bash command".to_string());
}
// Only check exfiltration for non-standard commands
let is_standard = cmd.trim_start().starts_with("cargo")
|| cmd.trim_start().starts_with("rustc")
|| cmd.trim_start().starts_with("git ")
|| cmd.trim_start().starts_with("ls")
|| cmd.trim_start().starts_with("pwd")
|| cmd.trim_start().starts_with("echo")
|| cmd.trim_start().starts_with("cat")
|| cmd.trim_start().starts_with("find")
|| cmd.trim_start().starts_with("grep")
|| cmd.trim_start().starts_with("test");
if !is_standard {
for pat in EXFIL_PATTERNS {
if cmd.contains(pat) {
return Some(format!("potential data-exfiltration command blocked (matched '{}')", pat));
}
}
}
for pat in SENSITIVE_PATH_PATTERNS {
if cmd.contains(pat) {
return Some(format!("refused to read/write sensitive path '{}'", pat));
}
}
let dangerous = ["rm -rf /", "rm -rf --no-preserve-root", "rm -rf ~",
"rm -fr /", "mkfs.", "dd if=", ":(){", "> /dev/sda",
"chmod -R 000 /", "shutdown ", "poweroff ", "reboot ", "halt "];
for pat in &dangerous {
if cmd.contains(pat) {
return Some(format!("destructive command pattern blocked: {}", pat));
}
}
if contains_any(cmd, STUB_PATTERNS) {
return Some("bash command contains stub pattern".to_string());
}
}
// git_operator: require reason
if tool_name == "git_operator" {
let reason = args.get("reason").and_then(|v| v.as_str()).unwrap_or("");
if reason.trim().len() < MIN_REASON_LEN {
return Some("git_operator requires a non-trivial 'reason' (>= 8 chars)".to_string());
}
}
None
}
/// Check if `text` matches any pattern (case-insensitive substring).
fn contains_any(text: &str, patterns: &[&str]) -> bool {
let lower = text.to_lowercase();
patterns.iter().any(|p| lower.contains(&p.to_lowercase()))
}
/// Build an ASCII tree of the workspace directory structure for the
/// system prompt, so the LLM can see the file layout.
///
/// Flow: for each root, walk using `ignore::WalkBuilder` (respecting
/// `.gitignore` and hidden files) → prefix `[DIR]` for directories →
/// truncate after 1000 entries.
fn generate_workspace_tree(roots: &[std::path::PathBuf]) -> String {
let mut out = String::new();
out.push_str("Current Workspace Directory Structure:\n");
for root in roots {
out.push_str(&format!("Root: {}\n", root.display()));
let walker = ignore::WalkBuilder::new(root)
.hidden(true)
.git_ignore(true)
.build();
let mut count = 0;
for entry in walker.flatten() {
let path = entry.path();
if let Ok(rel) = path.strip_prefix(root) {
if rel.as_os_str().is_empty() { continue; }
let is_dir = entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false);
let prefix = if is_dir { "[DIR] " } else { " " };
out.push_str(&format!(" {}{}\n", prefix, rel.display()));
count += 1;
if count > 1000 {
out.push_str(" ... (truncated)\n");
break;
}
}
}
}
out
}
/// Synchronous subagent entry point: run up to `ctx.max_steps` iterations
/// of the LLM tool loop.
///
/// Flow: inject system prompt → for each step: resolve provider config,
/// build an LLM client, call `chat_with_tools_non_streaming`, process tool
/// calls or collect text output → send `SubagentEvent`s on `tx` → break on
/// first text-only (non-empty) response.
/// Flow: inject system prompt (with workspace tree if available) → for each
/// step: resolve provider config, build an LLM client, call
/// `chat_with_tools_non_streaming`, process tool calls (gated against both
/// the allowlist and Harness-style content safety checks) or collect text
/// output → send `SubagentEvent`s on `tx` → break on first text-only
/// (non-empty) response.
///
/// Why: runs synchronously on a dedicated thread so the main async event
/// loop is not blocked. Tool gating prevents restricted or risky tools from
/// executing unless explicitly allowed.
/// loop is not blocked. Tool gating prevents restricted, risky, or
/// malicious/poor-quality tool calls from executing.
///
/// Return: the concatenated text output, or an `anyhow::Error` if the LLM
/// call fails at any step.
pub fn run_subagent(ctx: SubagentContext, tx: mpsc::Sender<SubagentEvent>) -> anyhow::Result<String> {
let mut output = String::new();
let mut messages: Vec<ChatMessage> = Vec::new();
messages.push(ChatMessage::system(ctx.system_prompt.clone()));
// Build system prompt with workspace tree context if we have workspaces,
// giving subagents the same project-awareness as the main agent.
let system_with_context = if ctx.workspaces.is_empty() {
ctx.system_prompt.clone()
} else {
let tree_info = generate_workspace_tree(&ctx.workspaces);
format!("{}\n\n{}", ctx.system_prompt, tree_info)
};
messages.push(ChatMessage::system(system_with_context));
let tool_ctx = crate::tool::ToolCtx::builder()
.session_dir(ctx.session_dir.clone())
@@ -145,6 +362,7 @@ pub fn run_subagent(ctx: SubagentContext, tx: mpsc::Sender<SubagentEvent>) -> an
_args: args.clone(),
});
// Level 1: allowlist check — is this tool even permitted?
if !generally_allowed {
let msg = format!("tool '{}' not allowed for this subagent", tool_name);
messages.push(ChatMessage::tool_result(tool_call.id.clone(), msg.clone()));
@@ -155,6 +373,7 @@ pub fn run_subagent(ctx: SubagentContext, tx: mpsc::Sender<SubagentEvent>) -> an
continue;
}
// Level 2: risky tool check — risky tools require explicit permission
if tool_is_risky(tool_name) && !explicitly_allowed {
let msg = format!("risky tool '{}' requires explicit permission; not allowed for this subagent", tool_name);
messages.push(ChatMessage::tool_result(tool_call.id.clone(), msg.clone()));
@@ -165,6 +384,20 @@ pub fn run_subagent(ctx: SubagentContext, tx: mpsc::Sender<SubagentEvent>) -> an
continue;
}
// Level 3: Harness-style content safety gating — mirrors the main
// agent's gate_tool_call checks (path traversal, reason validation,
// stub/denial/assumption scanning, bash exfiltration, destructive
// commands, sensitive path reads).
if let Some(block_reason) = gate_subagent_tool_call(tool_name, &args) {
let msg = format!("Blocked by subagent gate: {}", block_reason);
messages.push(ChatMessage::tool_result(tool_call.id.clone(), msg.clone()));
let _ = tx.blocking_send(SubagentEvent::ToolResult {
_tool: tool_name.clone(),
_output: msg,
});
continue;
}
let result = match tools.iter().find(|t| t.name() == tool_name.as_str()) {
Some(tool) => tool.run(&tool_ctx, &args),
None => Err(anyhow::anyhow!("tool '{}' not found", tool_name)),