Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 32 additions & 6 deletions src/commands.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,13 +69,10 @@ pub enum Commands {
#[command(subcommand)]
command: SecretCommands,
},
/// Inspect a WebAssembly component and display its JSON schema (for debugging).
/// Inspect a WebAssembly component (display schema or call component).
Inspect {
/// Component ID to inspect
component_id: String,
/// Directory where components are stored. Defaults to $XDG_DATA_HOME/wassette/components
#[arg(long)]
component_dir: Option<PathBuf>,
#[command(subcommand)]
command: InspectCommands,
},
/// Manage tools (list, read, invoke).
Tool {
Expand Down Expand Up @@ -435,6 +432,35 @@ pub enum ToolCommands {
},
}

#[derive(Subcommand, Debug)]
pub enum InspectCommands {
/// Display JSON schema of a component.
Schema {
/// Component ID to inspect
component_id: String,
/// Directory where components are stored. Defaults to $XDG_DATA_HOME/wassette/components
#[arg(long)]
component_dir: Option<PathBuf>,
},
/// Call a component with arguments.
Call {
/// Component ID to call
component_id: String,
/// Arguments in KEY=VALUE format. Can be specified multiple times.
#[arg(long = "arg", value_parser = crate::parse_env_var)]
args: Vec<(String, String)>,
/// Path to policy file for this component
#[arg(long)]
policy_file: Option<PathBuf>,
/// Directory where components are stored. Defaults to $XDG_DATA_HOME/wassette/components
#[arg(long)]
component_dir: Option<PathBuf>,
/// Output format
#[arg(short = 'o', long = "output-format", default_value = "json")]
output_format: OutputFormat,
},
}

#[derive(Subcommand, Debug)]
pub enum RegistryCommands {
/// Search for components in the registry.
Expand Down
312 changes: 268 additions & 44 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ use rmcp::service::serve_server;
use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
use rmcp::transport::streamable_http_server::StreamableHttpService;
use rmcp::transport::{stdio as stdio_transport, SseServer};
use serde_json::{json, Map};
use serde_json::{json, Map, Value};
use tracing_subscriber::layer::SubscriberExt as _;
use tracing_subscriber::util::SubscriberInitExt as _;

Expand All @@ -31,8 +31,9 @@ mod utils;

use cli_handlers::{create_lifecycle_manager, handle_tool_cli_command};
use commands::{
Cli, Commands, ComponentCommands, GrantPermissionCommands, PermissionCommands, PolicyCommands,
RegistryCommands, RevokePermissionCommands, SecretCommands, Shell, ToolCommands, Transport,
Cli, Commands, ComponentCommands, GrantPermissionCommands, InspectCommands, PermissionCommands,
PolicyCommands, RegistryCommands, RevokePermissionCommands, SecretCommands, Shell,
ToolCommands, Transport,
};
use format::{print_result, OutputFormat};
use server::McpServer;
Expand Down Expand Up @@ -716,50 +717,273 @@ async fn main() -> Result<()> {
}
}
},
Commands::Inspect {
component_id,
component_dir,
} => {
let component_dir = component_dir.clone().or_else(|| cli.component_dir.clone());
let lifecycle_manager = create_lifecycle_manager(component_dir).await?;

// Get the component schema from the lifecycle manager
let schema = lifecycle_manager
.get_component_schema(component_id)
.await
.context(format!(
"Component '{}' not found. Use 'component load' to load the component first.",
component_id
))?;

// Display tools information
if let Some(arr) = schema["tools"].as_array() {
for t in arr {
// The tool info is nested in properties.result
let tool_info = &t["properties"]["result"];
let name = tool_info["name"]
Commands::Inspect { command } => match command {
InspectCommands::Schema {
component_id,
component_dir,
} => {
let component_dir = component_dir.clone().or_else(|| cli.component_dir.clone());
let lifecycle_manager = create_lifecycle_manager(component_dir).await?;

// Get the component schema from the lifecycle manager
let schema = lifecycle_manager
.get_component_schema(component_id)
.await
.context(format!(
"Component '{}' not found. Use 'component load' to load the component first.",
component_id
))?;

// Display tools information
if let Some(arr) = schema["tools"].as_array() {
for t in arr {
// The tool info is nested in properties.result
let tool_info = &t["properties"]["result"];
let name = tool_info["name"]
.as_str()
.unwrap_or("<unnamed>")
.to_string();
let description: Option<String> =
tool_info["description"].as_str().map(|s| s.to_string());
let input_schema = tool_info["inputSchema"].clone();
let output_schema = tool_info["outputSchema"].clone();

println!("{name}, {description:?}");
println!(
"input schema: {}",
serde_json::to_string_pretty(&input_schema)?
);
println!(
"output schema: {}",
serde_json::to_string_pretty(&output_schema)?
);
}
} else {
println!("No tools found in component");
}
}
InspectCommands::Call {
component_id,
args,
policy_file,
component_dir,
output_format,
} => {
let component_dir = component_dir.clone().or_else(|| cli.component_dir.clone());
let lifecycle_manager = create_lifecycle_manager(component_dir).await?;

// Ensure the component is loaded in memory first
lifecycle_manager
.ensure_component_loaded(component_id)
.await
.context("Failed to load component into memory")?;

// Apply policy file if provided
if let Some(policy_path) = policy_file {
tracing::info!("Applying policy from file to component {}", component_id);

// Use PolicyParser from the policy crate to parse the policy file
let policy_doc = policy::PolicyParser::parse_file(policy_path)
.context("Failed to parse policy file")?;

// Apply storage permissions from the allow list
if let Some(storage) = &policy_doc.permissions.storage {
if let Some(allow_list) = &storage.allow {
for perm in allow_list {
lifecycle_manager
.grant_permission(
component_id,
"storage",
&json!({
"uri": perm.uri,
"access": perm.access.iter().map(|a| match a {
policy::AccessType::Read => "read",
policy::AccessType::Write => "write",
}).collect::<Vec<_>>()
}),
)
.await?;
}
}
}

// Apply network permissions from the allow list
if let Some(network) = &policy_doc.permissions.network {
if let Some(allow_list) = &network.allow {
for perm in allow_list {
match perm {
policy::NetworkPermission::Host(host) => {
lifecycle_manager
.grant_permission(
component_id,
"network",
&json!({
"host": host.host
}),
)
.await?;
}
policy::NetworkPermission::Cidr(cidr) => {
lifecycle_manager
.grant_permission(
component_id,
"network",
&json!({
"cidr": cidr.cidr
}),
)
.await?;
}
}
}
}
}

// Apply environment variable permissions from the allow list
if let Some(environment) = &policy_doc.permissions.environment {
if let Some(allow_list) = &environment.allow {
for perm in allow_list {
lifecycle_manager
.grant_permission(
component_id,
"environment-variable",
&json!({
"key": perm.key
}),
)
.await?;
}
}
}

// Apply resource limits if specified
if let Some(resources) = &policy_doc.permissions.resources {
if let Some(limits) = &resources.limits {
let mut limits_json = json!({});

if let Some(cpu) = &limits.cpu {
limits_json["cpu"] = match cpu {
policy::CpuLimit::String(s) => json!(s),
policy::CpuLimit::Number(n) => json!(n),
};
}

if let Some(memory) = &limits.memory {
limits_json["memory"] = match memory {
policy::MemoryLimit::String(s) => json!(s),
policy::MemoryLimit::Number(n) => json!(n),
};
}

if !limits_json.as_object().unwrap().is_empty() {
lifecycle_manager
.grant_permission(
component_id,
"resources",
&json!({
"limits": limits_json
}),
)
.await?;
}
}
}
}

// Get the component schema to find the tool name
let schema = lifecycle_manager
.get_component_schema(component_id)
.await
.context(format!(
"Component '{}' not found. Use 'component load' to load the component first.",
component_id
))?;

// Find the tool to call in the component
let tool_name = if let Some(arr) = schema["tools"].as_array() {
if arr.is_empty() {
bail!("No tools found in component");
}

// Check if component has multiple tools
if arr.len() > 1 {
// Collect tool names for error message
let tool_names: Vec<String> = arr
.iter()
.map(|tool| {
let tool_info = if tool.get("properties").is_some() {
&tool["properties"]["result"]
} else {
tool
};
tool_info["name"]
.as_str()
.unwrap_or("<unknown>")
.to_string()
})
.collect();
bail!(
"Component '{}' has multiple tools: [{}]. This command currently only supports components with a single tool.",
component_id,
tool_names.join(", ")
);
}

// Only one tool present, proceed
let first_tool = &arr[0];
let tool_info = if first_tool.get("properties").is_some() {
&first_tool["properties"]["result"]
} else {
// Fallback to flat structure (metadata schema)
first_tool
};

tool_info["name"]
.as_str()
.unwrap_or("<unnamed>")
.to_string();
let description: Option<String> =
tool_info["description"].as_str().map(|s| s.to_string());
let input_schema = tool_info["inputSchema"].clone();
let output_schema = tool_info["outputSchema"].clone();

println!("{name}, {description:?}");
println!(
"input schema: {}",
serde_json::to_string_pretty(&input_schema)?
);
println!(
"output schema: {}",
serde_json::to_string_pretty(&output_schema)?
);
.ok_or_else(|| anyhow::anyhow!("Tool name not found in schema"))?
.to_string()
} else {
bail!("No tools found in component");
};

// Convert args to JSON
let mut arguments = Map::new();
for (key, value) in args {
// Try to parse value as JSON, otherwise treat as string
let json_value: Value = serde_json::from_str(value)
.unwrap_or_else(|_| Value::String(value.clone()));
arguments.insert(key.clone(), json_value);
}

// Call the component
let parameters_json = serde_json::to_string(&arguments)?;
let result = lifecycle_manager
.execute_component_call(component_id, &tool_name, &parameters_json)
.await;

match result {
Ok(result_str) => {
let result_value: Value = serde_json::from_str(&result_str)
.unwrap_or_else(|_| Value::String(result_str.clone()));

print_result(
&rmcp::model::CallToolResult {
content: Some(vec![rmcp::model::Content::text(
serde_json::to_string_pretty(&result_value)?,
)]),
structured_content: None,
is_error: Some(false),
},
*output_format,
)?;
}
Err(e) => {
eprintln!("Error calling component '{}': {}", component_id, e);
std::process::exit(1);
}
}
} else {
println!("No tools found in component");
}
}
},
Commands::Registry { command } => match command {
RegistryCommands::Search {
query,
Expand Down
Loading
Loading