use std::borrow::Cow;
use std::path::{Path, PathBuf};
use aither_core::llm::tool::{Tool, ToolResult};
use eyre::{Context as _, Result, bail};
use schemars::JsonSchema;
use serde::Deserialize;
use tracing::info;
use crate::preview::request::{
self, CliHydrolysisPreviewTheme, CliPreviewBackend, CliPreviewPlatform, DEFAULT_FRAME,
PreviewRequest, PreviewTarget,
};
use crate::preview::{
HydrolysisPreviewRequest, launch_preview_session, render_preview_with_hydrolysis,
};
use crate::project::read_project_crate_name;
#[derive(Debug, Deserialize, JsonSchema)]
pub struct PreviewArgs {
pub target: String,
#[serde(default)]
pub expr: bool,
#[serde(default)]
pub frame: Option<String>,
#[serde(default)]
pub backend: Option<CliPreviewBackend>,
#[serde(default)]
pub theme: Option<CliHydrolysisPreviewTheme>,
#[serde(default)]
pub platform: Option<CliPreviewPlatform>,
}
impl PreviewArgs {
pub fn resolve(&self, crate_name: &str) -> Result<PreviewRequest> {
let frame = self.frame.as_deref().unwrap_or(DEFAULT_FRAME);
let (width, height) = request::parse_frame(frame)?;
let platform = request::resolve_preview_platform(self.platform)?;
let backend = request::resolve_preview_backend(platform, self.backend)?;
let hydrolysis_theme = request::resolve_hydrolysis_preview_theme(backend, self.theme)?;
let target = request::resolve_preview_target(crate_name, &self.target, self.expr);
Ok(PreviewRequest {
platform,
backend,
hydrolysis_theme,
target,
width,
height,
})
}
}
fn sanitize_target_name(target: &str) -> String {
let mut name = String::with_capacity(target.len());
for ch in target.chars() {
if ch.is_ascii_alphanumeric() || matches!(ch, '_' | '.' | '-') {
name.push(ch);
} else if !name.ends_with('_') {
name.push('_');
}
}
name
}
#[derive(Debug)]
pub struct PreviewTool {
project_path: PathBuf,
sccache_path: Option<PathBuf>,
}
impl PreviewTool {
#[must_use]
pub const fn new(project_path: PathBuf, sccache_path: Option<PathBuf>) -> Self {
Self {
project_path,
sccache_path,
}
}
async fn render(&self, args: PreviewArgs) -> ToolResult {
match self.run(&args).await {
Ok((output_path, bytes)) => {
info!(path = %output_path.display(), "preview rendered");
ToolResult::image(bytes, "image/png")
}
Err(error) => ToolResult::error(format!("{error:#}")),
}
}
async fn output_path(&self, request: &PreviewRequest) -> Result<PathBuf> {
let dir = crate::water_dir::build_cache_container_for(&self.project_path)
.await?
.join("mcp")
.join("preview");
smol::fs::create_dir_all(&dir)
.await
.wrap_err_with(|| format!("failed to create {}", dir.display()))?;
Ok(dir.join(format!(
"{}-{}x{}.png",
sanitize_target_name(request.target.display_name()),
request.width,
request.height
)))
}
async fn run(&self, args: &PreviewArgs) -> Result<(PathBuf, Vec<u8>)> {
let crate_name = read_project_crate_name(&self.project_path).await?;
let request = args.resolve(&crate_name)?;
request::check_toolchain_for_backend(request.platform, request.backend).await?;
let output_path = self.output_path(&request).await?;
match request.backend {
CliPreviewBackend::Hydrolysis => {
render_preview_with_hydrolysis(
HydrolysisPreviewRequest {
project_path: &self.project_path,
source: request.target.hydrolysis_source(),
theme: request
.hydrolysis_theme
.expect("resolve guarantees a theme for hydrolysis"),
width: request.width,
height: request.height,
sccache_path: self.sccache_path.clone(),
progress: None,
},
&output_path,
None,
)
.await?;
}
CliPreviewBackend::Apple | CliPreviewBackend::Android => {
let PreviewTarget::Function {
function_path,
symbol,
} = &request.target
else {
bail!(
"Expression preview is currently supported only with the `hydrolysis` backend."
);
};
self.render_support_app(&request, function_path, symbol, &output_path)
.await?;
}
}
let bytes = smol::fs::read(&output_path).await?;
Ok((output_path, bytes))
}
async fn render_support_app(
&self,
request: &PreviewRequest,
function_path: &str,
symbol: &str,
output_path: &Path,
) -> Result<()> {
let mut session = launch_preview_session(
&self.project_path,
request.platform.into(),
self.sccache_path.clone(),
None,
)
.await?;
let result = async {
let dylib = session.build_dylib(&self.project_path).await?;
let png_data = request::render_with_symbol(
&mut session,
function_path,
symbol,
dylib.id,
&dylib.path,
request.width,
request.height,
)
.await?;
if png_data.is_empty() {
bail!("Preview returned empty PNG data");
}
smol::fs::write(output_path, &png_data).await?;
Ok(())
}
.await;
match result {
Ok(()) => {
session.detach();
Ok(())
}
Err(err) => match session.shutdown().await {
Ok(()) => Err(err),
Err(shutdown_error) => Err(err.wrap_err(format!(
"preview support app shutdown also failed: {shutdown_error}"
))),
},
}
}
}
impl Tool for PreviewTool {
type Arguments = PreviewArgs;
type Res = ToolResult;
fn name(&self) -> Cow<'static, str> {
"preview".into()
}
async fn call(&self, args: Self::Arguments) -> aither_core::Result<Self::Res> {
Ok(self.render(args).await)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn args_default_to_water_preview_defaults() {
let args: PreviewArgs =
serde_json::from_str(r#"{"target": "views::home"}"#).expect("minimal args parse");
assert_eq!(args.target, "views::home");
assert!(!args.expr);
assert_eq!(args.frame, None);
assert_eq!(args.backend, None);
assert_eq!(args.theme, None);
assert_eq!(args.platform, None);
}
#[test]
fn args_parse_all_fields() {
let args: PreviewArgs = serde_json::from_str(
r#"{
"target": "text(\"hi\")",
"expr": true,
"frame": "800x600",
"backend": "hydrolysis",
"theme": "material3",
"platform": "macos"
}"#,
)
.expect("full args parse");
assert!(args.expr);
assert_eq!(args.frame.as_deref(), Some("800x600"));
assert_eq!(args.backend, Some(CliPreviewBackend::Hydrolysis));
assert_eq!(args.theme, Some(CliHydrolysisPreviewTheme::Material3));
assert_eq!(args.platform, Some(CliPreviewPlatform::Macos));
}
#[test]
fn sanitize_collapses_unsafe_runs() {
assert_eq!(sanitize_target_name("views::home"), "views_home");
assert_eq!(sanitize_target_name("text(\"hello\")"), "text_hello_");
assert_eq!(sanitize_target_name("a.b-c_d"), "a.b-c_d");
assert_eq!(sanitize_target_name("**"), "_");
}
#[test]
fn resolves_to_the_same_request_as_water_preview() {
let args: PreviewArgs = serde_json::from_str(
r#"{
"target": "text(\"hi\")",
"expr": true,
"frame": "800x600",
"backend": "hydrolysis",
"theme": "material3",
"platform": "macos"
}"#,
)
.expect("args parse");
let request = args.resolve("demo_app").expect("resolve");
assert_eq!(
request,
PreviewRequest {
platform: CliPreviewPlatform::Macos,
backend: CliPreviewBackend::Hydrolysis,
hydrolysis_theme: Some(crate::preview::HydrolysisPreviewTheme::Material3),
target: PreviewTarget::Expression {
expression: "text(\"hi\")".to_string(),
},
width: 800.0,
height: 600.0,
}
);
}
}