use super::super::args::{McpArgs, McpSub, McpWrapArgs};
use super::coverage;
use super::session_state_window;
use anyhow::{Context, Result};
use assay_core::mcp::policy::McpPolicy;
use assay_core::mcp::proxy::{McpProxy, ProxyConfig, ProxyConfigRaw};
use serde_json::{json, Value};
use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};
pub async fn run(args: McpArgs) -> anyhow::Result<i32> {
match args.cmd {
McpSub::Wrap(wrap_args) => cmd_wrap(wrap_args).await,
McpSub::ConfigPath(config_args) => {
super::config_path::run(config_args);
Ok(0)
}
McpSub::Discover(discover_args) => super::discover::run(discover_args).await,
McpSub::Kill(kill_args) => super::kill::run(kill_args).await,
McpSub::Tool(tool_args) => Ok(super::tool::cmd_tool(tool_args.cmd)),
}
}
fn is_explicit_tool_name(pattern: &str) -> bool {
let trimmed = pattern.trim();
!trimmed.is_empty() && !trimmed.contains('*')
}
fn collect_declared_tools(policy: &McpPolicy) -> Vec<String> {
let mut declared = BTreeSet::new();
if let Some(allow) = &policy.tools.allow {
for tool in allow {
if is_explicit_tool_name(tool) {
declared.insert(tool.trim().to_string());
}
}
}
if let Some(deny) = &policy.tools.deny {
for tool in deny {
if is_explicit_tool_name(tool) {
declared.insert(tool.trim().to_string());
}
}
}
for tool in policy.schemas.keys() {
if tool != "$defs" && is_explicit_tool_name(tool) {
declared.insert(tool.trim().to_string());
}
}
for constraint in &policy.constraints {
if is_explicit_tool_name(&constraint.tool) {
declared.insert(constraint.tool.trim().to_string());
}
}
for tool in policy.tool_pins.keys() {
if is_explicit_tool_name(tool) {
declared.insert(tool.trim().to_string());
}
}
declared.into_iter().collect()
}
fn extract_tool_name_from_decision_event(v: &Value) -> Result<String> {
for (pointer, label) in [
(Some("/data/tool"), "data.tool"),
(Some("/data/tool_name"), "data.tool_name"),
(None, "tool"),
(None, "tool_name"),
] {
let value = match pointer {
Some(path) => v.pointer(path),
None => v.get(label),
};
if let Some(s) = value.and_then(Value::as_str) {
let tool = s.trim();
if !tool.is_empty() {
return Ok(tool.to_string());
}
anyhow::bail!("field '{label}' must be non-empty string");
}
}
anyhow::bail!("missing required field: 'data.tool', 'data.tool_name', 'tool', or 'tool_name'")
}
fn extract_tool_classes_from_decision_event(v: &Value) -> BTreeSet<String> {
match v.pointer("/data/tool_classes") {
Some(Value::Array(arr)) => arr
.iter()
.filter_map(Value::as_str)
.map(str::trim)
.filter(|s| !s.is_empty())
.map(ToOwned::to_owned)
.collect(),
_ => BTreeSet::new(),
}
}
async fn normalize_decision_jsonl_to_coverage_jsonl(input: &Path, output: &Path) -> Result<()> {
let content = tokio::fs::read_to_string(input)
.await
.with_context(|| format!("failed to read decision log {}", input.display()))?;
let mut lines = Vec::new();
for (lineno, raw) in content.lines().enumerate() {
let line = raw.trim();
if line.is_empty() {
continue;
}
let value: Value = serde_json::from_str(line)
.with_context(|| format!("invalid json at line {}", lineno + 1))?;
let tool = extract_tool_name_from_decision_event(&value)
.with_context(|| format!("measurement error at line {}", lineno + 1))?;
let tool_classes = extract_tool_classes_from_decision_event(&value)
.into_iter()
.collect::<Vec<_>>();
lines.push(serde_json::to_string(&json!({
"tool": tool,
"tool_classes": tool_classes,
}))?);
}
let mut normalized = lines.join("\n");
if !normalized.is_empty() {
normalized.push('\n');
}
tokio::fs::write(output, normalized)
.await
.with_context(|| {
format!(
"failed to write normalized coverage input {}",
output.display()
)
})?;
Ok(())
}
fn unique_temp_path(stem: &str, extension: &str) -> PathBuf {
let stamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_nanos();
std::env::temp_dir().join(format!(
"assay-{stem}-{}-{stamp}.{extension}",
std::process::id()
))
}
fn generate_session_id() -> String {
let stamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_nanos();
format!("mcpwrap-{}-{stamp}", std::process::id())
}
struct TempPathGuard {
path: PathBuf,
}
impl TempPathGuard {
fn new(path: PathBuf) -> Self {
Self { path }
}
fn path(&self) -> &Path {
&self.path
}
}
impl Drop for TempPathGuard {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.path);
}
}
async fn cmd_wrap(args: McpWrapArgs) -> anyhow::Result<i32> {
if args.deny_deprecations {
std::env::set_var("ASSAY_STRICT_DEPRECATIONS", "1");
}
if args.command.is_empty() {
anyhow::bail!("No command specified. Usage: assay mcp wrap -- <cmd> [args]");
}
let cmd = &args.command[0];
let cmd_args = &args.command[1..];
let session_id = generate_session_id();
let policy = if args.policy.exists() {
eprintln!("[assay] loading policy from {}", args.policy.display());
McpPolicy::from_file(&args.policy)?
} else {
anyhow::bail!(
"Policy file not found at {}. Aborting instead of using insecure default (allow-all) policy.",
args.policy.display()
);
};
let declared_tools = if args.coverage_out.is_some() {
collect_declared_tools(&policy)
} else {
Vec::new()
};
let temp_decision_log = if args.coverage_out.is_some() && args.decision_log.is_none() {
Some(TempPathGuard::new(unique_temp_path(
"wrap-decision-log",
"ndjson",
)))
} else {
None
};
let effective_decision_log = args.decision_log.clone().or_else(|| {
temp_decision_log
.as_ref()
.map(|guard| guard.path().to_path_buf())
});
let raw = ProxyConfigRaw {
dry_run: args.dry_run,
verbose: args.verbose,
audit_log_path: args.audit_log.clone(),
decision_log_path: effective_decision_log.clone(),
event_source: args.event_source.clone(),
server_id: args
.label
.clone()
.unwrap_or_else(|| "default-mcp-server".into()),
};
let config = ProxyConfig::try_from_raw(raw)?;
let state_window_event_source = config.event_source.clone();
let state_window_server_id = config.server_id.clone();
if config.dry_run {
eprintln!("[assay] DRY RUN MODE: No actions will be blocked.");
}
eprintln!("[assay] wrapping command: {} {:?}", cmd, cmd_args);
let proxy = McpProxy::spawn(cmd, cmd_args, policy, config)?;
let wrapped_code = tokio::task::spawn_blocking(move || proxy.run()).await??;
let coverage_status = if let Some(coverage_out) = args.coverage_out.as_ref() {
let normalized_log = TempPathGuard::new(unique_temp_path("wrap-coverage-input", "jsonl"));
let decision_log_path = effective_decision_log
.as_ref()
.context("coverage generation requires a decision log path")?;
match normalize_decision_jsonl_to_coverage_jsonl(decision_log_path, normalized_log.path())
.await
{
Ok(()) => {
coverage::write_generated_coverage_report(
normalized_log.path(),
coverage_out,
&declared_tools,
"decision_jsonl",
)
.await?
}
Err(e) => {
eprintln!(
"Measurement error: failed to normalize decision log {}: {e}",
decision_log_path.display()
);
crate::exit_codes::EXIT_CONFIG_ERROR
}
}
} else {
crate::exit_codes::EXIT_SUCCESS
};
let state_status = if let Some(state_window_out) = args.state_window_out.as_ref() {
let event_source = state_window_event_source
.as_deref()
.context("state window export requires event_source")?;
session_state_window::write_state_window_out(
state_window_out,
event_source,
&state_window_server_id,
&session_id,
)
.await?
} else {
crate::exit_codes::EXIT_SUCCESS
};
if wrapped_code != crate::exit_codes::EXIT_SUCCESS {
if coverage_status != crate::exit_codes::EXIT_SUCCESS
|| state_status != crate::exit_codes::EXIT_SUCCESS
{
eprintln!(
"Coverage/state-window generation failed after wrapped command exited with code {}; preserving wrapped exit code",
wrapped_code
);
}
return Ok(wrapped_code);
}
if coverage_status != crate::exit_codes::EXIT_SUCCESS {
if state_status != crate::exit_codes::EXIT_SUCCESS {
eprintln!(
"State window export failed after coverage generation failed; preserving coverage exit code"
);
}
return Ok(coverage_status);
}
Ok(state_status)
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
fn fixture_path(name: &str) -> PathBuf {
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("../../scripts/ci/fixtures/coverage")
.join(name)
}
#[tokio::test]
async fn mcp_wrap_coverage_normalizer_extracts_nested_tool_and_classes() {
let dir = tempdir().unwrap();
let out = dir.path().join("normalized.jsonl");
normalize_decision_jsonl_to_coverage_jsonl(
&fixture_path("decision_event_basic.jsonl"),
&out,
)
.await
.expect("normalization should succeed");
let content = std::fs::read_to_string(&out).unwrap();
let line = content.lines().next().unwrap();
let value: Value = serde_json::from_str(line).unwrap();
assert_eq!(value["tool"], "assay_policy_decide");
assert_eq!(value["tool_classes"], json!(["sink:network"]));
}
#[tokio::test]
async fn mcp_wrap_coverage_normalizer_accepts_data_tool_name_fallback() {
let dir = tempdir().unwrap();
let input = dir.path().join("decision.jsonl");
let out = dir.path().join("normalized.jsonl");
std::fs::write(
&input,
r#"{"data":{"tool_name":"web_search_alt"}}
"#,
)
.unwrap();
normalize_decision_jsonl_to_coverage_jsonl(&input, &out)
.await
.expect("data.tool_name fallback should succeed");
let content = std::fs::read_to_string(&out).unwrap();
let line = content.lines().next().unwrap();
let value: Value = serde_json::from_str(line).unwrap();
assert_eq!(value["tool"], "web_search_alt");
}
#[tokio::test]
async fn mcp_wrap_coverage_normalizer_accepts_top_level_tool_name_fallback() {
let dir = tempdir().unwrap();
let input = dir.path().join("decision.jsonl");
let out = dir.path().join("normalized.jsonl");
std::fs::write(
&input,
r#"{"tool_name":"web_search"}
"#,
)
.unwrap();
normalize_decision_jsonl_to_coverage_jsonl(&input, &out)
.await
.expect("top-level tool_name fallback should succeed");
let content = std::fs::read_to_string(&out).unwrap();
let line = content.lines().next().unwrap();
let value: Value = serde_json::from_str(line).unwrap();
assert_eq!(value["tool"], "web_search");
assert_eq!(value["tool_classes"], json!([]));
}
#[tokio::test]
async fn mcp_wrap_coverage_normalizer_rejects_missing_tool_fields() {
let dir = tempdir().unwrap();
let input = dir.path().join("decision.jsonl");
let out = dir.path().join("normalized.jsonl");
std::fs::write(
&input,
r#"{"decision":"deny"}
"#,
)
.unwrap();
let err = normalize_decision_jsonl_to_coverage_jsonl(&input, &out)
.await
.expect_err("missing tool fields must fail");
let msg = format!("{err:#}");
assert!(msg.contains("missing required field"));
}
#[tokio::test]
async fn mcp_wrap_coverage_normalizer_and_report_writer_emit_v1_report() {
let dir = tempdir().unwrap();
let normalized = dir.path().join("normalized.jsonl");
let coverage_out = dir.path().join("coverage.json");
normalize_decision_jsonl_to_coverage_jsonl(
&fixture_path("decision_event_basic.jsonl"),
&normalized,
)
.await
.expect("normalization should succeed");
let exit = coverage::write_generated_coverage_report(
&normalized,
&coverage_out,
&["assay_policy_decide".to_string()],
"decision_jsonl",
)
.await
.expect("coverage generation should complete");
assert_eq!(exit, crate::exit_codes::EXIT_SUCCESS);
let report: Value =
serde_json::from_str(&std::fs::read_to_string(&coverage_out).unwrap()).unwrap();
assert_eq!(report["schema_version"], "coverage_report_v1");
assert_eq!(report["run"]["source"], "decision_jsonl");
assert_eq!(
report["tools"]["tools_seen"],
json!(["assay_policy_decide"])
);
assert_eq!(
report["taxonomy"]["tool_classes_seen"],
json!(["sink:network"])
);
}
#[test]
fn mcp_wrap_coverage_collect_declared_tools_ignores_wildcards() {
let mut policy = McpPolicy::default();
policy.tools.allow = Some(vec!["assay_policy_decide".into(), "*".into()]);
policy.tools.deny = Some(vec!["assay_check_args".into(), "exec*".into()]);
policy.schemas.insert("$defs".into(), json!({}));
policy
.schemas
.insert("assay_check_sequence".into(), json!({"type": "object"}));
let declared = collect_declared_tools(&policy);
assert_eq!(
declared,
vec![
"assay_check_args".to_string(),
"assay_check_sequence".to_string(),
"assay_policy_decide".to_string(),
]
);
}
}