1240 lines
40 KiB
Rust

mod window_capture;
use screenshots::Screen;
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::PathBuf;
use std::process::Command;
use std::sync::Mutex;
use std::time::Instant;
use tauri::webview::PageLoadEvent;
use tauri::{Emitter, Manager, WebviewUrl, WebviewWindowBuilder};
use tauri_plugin_shell::{
process::{CommandChild, CommandEvent},
ShellExt,
};
use xcap::{Monitor, Window};
struct AgentProcess(Mutex<Option<CommandChild>>);
struct VisionStreamProcess(Mutex<Option<CommandChild>>);
impl Drop for AgentProcess {
fn drop(&mut self) {
if let Ok(mut process) = self.0.lock() {
if let Some(child) = process.take() {
let _ = child.kill();
}
}
}
}
impl Drop for VisionStreamProcess {
fn drop(&mut self) {
if let Ok(mut process) = self.0.lock() {
if let Some(child) = process.take() {
let _ = child.kill();
}
}
}
}
#[derive(Clone, Deserialize, Serialize)]
struct AgentLog {
level: String,
message: String,
timestamp: Option<String>,
}
#[derive(Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct CaptureResult {
pub(crate) screenshot_path: String,
pub(crate) screenshot_width: u32,
pub(crate) screenshot_height: u32,
pub(crate) scale_factor: f64,
pub(crate) source: Option<CaptureSource>,
pub(crate) screen_list_ms: u128,
pub(crate) capture_ms: u128,
pub(crate) save_ms: u128,
pub(crate) total_ms: u128,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct CaptureSource {
pub(crate) id: String,
pub(crate) kind: String,
pub(crate) label: String,
pub(crate) app_name: Option<String>,
pub(crate) title: Option<String>,
pub(crate) pid: Option<u32>,
pub(crate) x: i32,
pub(crate) y: i32,
pub(crate) width: u32,
pub(crate) height: u32,
pub(crate) scale_factor: f64,
}
#[derive(Clone, Deserialize, Serialize)]
pub(crate) struct Region {
pub(crate) id: String,
pub(crate) name: String,
pub(crate) description: Option<String>,
#[serde(rename = "type")]
pub(crate) region_type: String,
pub(crate) bbox_image: [f64; 4],
pub(crate) bbox_source: Option<[f64; 4]>,
pub(crate) bbox_screen: [f64; 4],
#[serde(rename = "scaleFactor")]
pub(crate) scale_factor: f64,
}
#[derive(Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct AnnotationFile {
pub(crate) app: String,
pub(crate) screenshot_path: String,
pub(crate) screenshot_width: u32,
pub(crate) screenshot_height: u32,
pub(crate) scale_factor: f64,
pub(crate) source: Option<CaptureSource>,
pub(crate) regions: Vec<Region>,
pub(crate) created_at: String,
pub(crate) updated_at: String,
}
pub(crate) fn data_dir(app: &tauri::AppHandle) -> Result<PathBuf, String> {
let dir = app
.path()
.app_data_dir()
.map_err(|error| error.to_string())?
.join("data");
fs::create_dir_all(&dir).map_err(|error| error.to_string())?;
Ok(dir)
}
pub(crate) fn regions_path(app: &tauri::AppHandle) -> Result<PathBuf, String> {
let dir = data_dir(app)?.join("regions");
fs::create_dir_all(&dir).map_err(|error| error.to_string())?;
Ok(dir.join("wechat.json"))
}
fn screenshot_path(app: &tauri::AppHandle) -> Result<PathBuf, String> {
let screenshot_dir = data_dir(app)?.join("screenshots");
fs::create_dir_all(&screenshot_dir).map_err(|error| error.to_string())?;
Ok(screenshot_dir.join("WeChat.jpg"))
}
fn save_capture_image(
width: u32,
height: u32,
rgba_bytes: Vec<u8>,
screenshot_path: PathBuf,
source: Option<CaptureSource>,
default_scale_factor: f64,
screen_list_ms: u128,
capture_ms: u128,
total_started_at: Instant,
) -> Result<CaptureResult, String> {
let save_started_at = Instant::now();
let rgba_image = image::RgbaImage::from_raw(width, height, rgba_bytes)
.ok_or_else(|| "invalid screenshot buffer".to_string())?;
let rgb_image = image::DynamicImage::ImageRgba8(rgba_image).to_rgb8();
let mut file = fs::File::create(&screenshot_path).map_err(|error| error.to_string())?;
let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut file, 88);
encoder
.encode_image(&rgb_image)
.map_err(|error| error.to_string())?;
let save_ms = save_started_at.elapsed().as_millis();
let total_ms = total_started_at.elapsed().as_millis();
let scale_factor = source
.as_ref()
.map(|item| item.scale_factor)
.unwrap_or(default_scale_factor);
println!(
"capture_source timing: screen_list={}ms capture={}ms save={}ms total={}ms size={}x{} scale={}",
screen_list_ms, capture_ms, save_ms, total_ms, width, height, scale_factor
);
Ok(CaptureResult {
screenshot_path: screenshot_path.to_string_lossy().to_string(),
screenshot_width: width,
screenshot_height: height,
scale_factor,
source,
screen_list_ms,
capture_ms,
save_ms,
total_ms,
})
}
#[cfg(target_os = "windows")]
#[derive(Clone, Copy)]
pub(crate) struct WinPhysicalRect {
pub(crate) left: i32,
pub(crate) top: i32,
pub(crate) right: i32,
pub(crate) bottom: i32,
}
#[cfg(target_os = "windows")]
impl WinPhysicalRect {
pub(crate) fn width(self) -> u32 {
(self.right - self.left).max(0) as u32
}
pub(crate) fn height(self) -> u32 {
(self.bottom - self.top).max(0) as u32
}
}
#[cfg(target_os = "windows")]
pub(crate) struct WinScreenMatch {
pub(crate) screen: Screen,
pub(crate) physical_x: i32,
pub(crate) physical_y: i32,
pub(crate) physical_width: u32,
pub(crate) physical_height: u32,
pub(crate) logical_x: i32,
pub(crate) logical_y: i32,
pub(crate) logical_width: u32,
pub(crate) logical_height: u32,
pub(crate) scale_factor: f64,
}
#[cfg(target_os = "windows")]
fn windows_capture_sources() -> Vec<CaptureSource> {
use windows_sys::Win32::Foundation::{HWND, LPARAM};
use windows_sys::Win32::UI::WindowsAndMessaging::EnumWindows;
unsafe extern "system" fn enum_window(hwnd: HWND, lparam: LPARAM) -> i32 {
let sources = &mut *(lparam as *mut Vec<CaptureSource>);
if let Some(source) = capture_source_from_hwnd(hwnd) {
let duplicate = sources.iter().any(|item| {
item.pid == source.pid
&& item.title == source.title
&& (item.x - source.x).abs() <= 2
&& (item.y - source.y).abs() <= 2
&& item.width.abs_diff(source.width) <= 2
&& item.height.abs_diff(source.height) <= 2
});
if !duplicate {
sources.push(source);
}
}
1
}
let mut sources = Vec::new();
unsafe {
EnumWindows(Some(enum_window), &mut sources as *mut _ as isize);
}
sources
}
#[cfg(not(target_os = "windows"))]
fn windows_capture_sources() -> Vec<CaptureSource> {
Vec::new()
}
#[cfg(target_os = "windows")]
pub(crate) fn capture_source_from_hwnd(
hwnd: windows_sys::Win32::Foundation::HWND,
) -> Option<CaptureSource> {
use windows_sys::Win32::UI::WindowsAndMessaging::{
GetWindowThreadProcessId, IsIconic, IsWindowVisible,
};
unsafe {
if IsWindowVisible(hwnd) == 0 || IsIconic(hwnd) != 0 {
return None;
}
}
let rect = window_physical_rect(hwnd).ok()?;
if rect.width() < 80 || rect.height() < 80 {
return None;
}
let title = window_title(hwnd);
let mut pid = 0u32;
unsafe {
GetWindowThreadProcessId(hwnd, &mut pid);
}
let app_name = process_name(pid).unwrap_or_default();
if app_name.is_empty() && title.is_empty() {
return None;
}
let screen_match = windows_screen_for_physical_rect(rect).ok()?;
let label_title = if title.is_empty() {
"无标题窗口"
} else {
&title
};
let label_app = if app_name.is_empty() {
"Windows 应用"
} else {
&app_name
};
Some(CaptureSource {
id: format!("win-window:{}", hwnd as usize),
kind: "window".to_string(),
label: format!("{label_app} · {label_title} · Windows"),
app_name: if app_name.is_empty() {
None
} else {
Some(app_name)
},
title: if title.is_empty() { None } else { Some(title) },
pid: if pid == 0 { None } else { Some(pid) },
x: screen_match.logical_x,
y: screen_match.logical_y,
width: screen_match.logical_width,
height: screen_match.logical_height,
scale_factor: screen_match.scale_factor,
})
}
#[cfg(target_os = "windows")]
pub(crate) fn window_physical_rect(
hwnd: windows_sys::Win32::Foundation::HWND,
) -> Result<WinPhysicalRect, String> {
use std::mem::size_of;
use windows_sys::Win32::Foundation::RECT;
use windows_sys::Win32::Graphics::Dwm::{DwmGetWindowAttribute, DWMWA_EXTENDED_FRAME_BOUNDS};
use windows_sys::Win32::UI::WindowsAndMessaging::GetWindowRect;
let mut rect = RECT::default();
let dwm_result = unsafe {
DwmGetWindowAttribute(
hwnd,
DWMWA_EXTENDED_FRAME_BOUNDS as u32,
&mut rect as *mut _ as *mut core::ffi::c_void,
size_of::<RECT>() as u32,
)
};
if dwm_result < 0 {
let ok = unsafe { GetWindowRect(hwnd, &mut rect) };
if ok == 0 {
return Err("GetWindowRect failed".to_string());
}
}
if rect.right <= rect.left || rect.bottom <= rect.top {
return Err("window rect is empty".to_string());
}
Ok(WinPhysicalRect {
left: rect.left,
top: rect.top,
right: rect.right,
bottom: rect.bottom,
})
}
#[cfg(target_os = "windows")]
pub(crate) fn window_title(hwnd: windows_sys::Win32::Foundation::HWND) -> String {
use windows_sys::Win32::UI::WindowsAndMessaging::{GetWindowTextLengthW, GetWindowTextW};
let length = unsafe { GetWindowTextLengthW(hwnd) };
if length <= 0 {
return String::new();
}
let mut buffer = vec![0u16; length as usize + 1];
let copied = unsafe { GetWindowTextW(hwnd, buffer.as_mut_ptr(), buffer.len() as i32) };
if copied <= 0 {
return String::new();
}
String::from_utf16_lossy(&buffer[..copied as usize])
.trim()
.to_string()
}
#[cfg(target_os = "windows")]
pub(crate) fn process_name(pid: u32) -> Option<String> {
use windows_sys::Win32::Foundation::CloseHandle;
use windows_sys::Win32::System::Threading::{
OpenProcess, QueryFullProcessImageNameW, PROCESS_QUERY_LIMITED_INFORMATION,
};
if pid == 0 {
return None;
}
let process = unsafe { OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid) };
if process.is_null() {
return None;
}
let mut buffer = vec![0u16; 32768];
let mut length = buffer.len() as u32;
let ok = unsafe { QueryFullProcessImageNameW(process, 0, buffer.as_mut_ptr(), &mut length) };
unsafe {
CloseHandle(process);
}
if ok == 0 || length == 0 {
return None;
}
let path = String::from_utf16_lossy(&buffer[..length as usize]);
PathBuf::from(path)
.file_name()
.and_then(|name| name.to_str())
.map(|name| name.to_string())
}
#[cfg(target_os = "windows")]
pub(crate) fn windows_screen_for_physical_rect(
rect: WinPhysicalRect,
) -> Result<WinScreenMatch, String> {
let screens = Screen::all().map_err(|error| error.to_string())?;
let mut best: Option<(Screen, i32, i32, u32, u32, f64, i64)> = None;
for screen in screens {
let scale_factor = screen.display_info.scale_factor as f64;
let physical_x = (screen.display_info.x as f64 * scale_factor).round() as i32;
let physical_y = (screen.display_info.y as f64 * scale_factor).round() as i32;
let physical_width = (screen.display_info.width as f64 * scale_factor)
.round()
.max(1.0) as u32;
let physical_height = (screen.display_info.height as f64 * scale_factor)
.round()
.max(1.0) as u32;
let physical_right = physical_x + physical_width as i32;
let physical_bottom = physical_y + physical_height as i32;
let intersection_width =
(rect.right.min(physical_right) - rect.left.max(physical_x)).max(0) as i64;
let intersection_height =
(rect.bottom.min(physical_bottom) - rect.top.max(physical_y)).max(0) as i64;
let intersection_area = intersection_width * intersection_height;
if intersection_area > best.map(|item| item.6).unwrap_or(-1) {
best = Some((
screen,
physical_x,
physical_y,
physical_width,
physical_height,
scale_factor,
intersection_area,
));
}
}
let (screen, physical_x, physical_y, physical_width, physical_height, scale_factor, area) =
best.ok_or_else(|| "no screen found".to_string())?;
if area <= 0 {
return Err("window is outside all screens".to_string());
}
let logical_x =
screen.display_info.x + ((rect.left - physical_x) as f64 / scale_factor).round() as i32;
let logical_y =
screen.display_info.y + ((rect.top - physical_y) as f64 / scale_factor).round() as i32;
let logical_width = (rect.width() as f64 / scale_factor).round().max(1.0) as u32;
let logical_height = (rect.height() as f64 / scale_factor).round().max(1.0) as u32;
Ok(WinScreenMatch {
screen,
physical_x,
physical_y,
physical_width,
physical_height,
logical_x,
logical_y,
logical_width,
logical_height,
scale_factor,
})
}
#[cfg(target_os = "windows")]
pub(crate) fn capture_windows_window(
hwnd_text: &str,
screenshot_path: PathBuf,
screen_list_ms: u128,
total_started_at: Instant,
) -> Result<CaptureResult, String> {
let hwnd_value = hwnd_text
.parse::<usize>()
.map_err(|error| error.to_string())?;
let hwnd = hwnd_value as windows_sys::Win32::Foundation::HWND;
let rect = window_physical_rect(hwnd)?;
let source =
capture_source_from_hwnd(hwnd).ok_or_else(|| "window source not found".to_string())?;
let screen_match = windows_screen_for_physical_rect(rect)?;
let physical_right = screen_match.physical_x + screen_match.physical_width as i32;
let physical_bottom = screen_match.physical_y + screen_match.physical_height as i32;
let x1 = rect.left.max(screen_match.physical_x);
let y1 = rect.top.max(screen_match.physical_y);
let x2 = rect.right.min(physical_right);
let y2 = rect.bottom.min(physical_bottom);
if x1 >= x2 || y1 >= y2 {
return Err("window source is outside selected screen".to_string());
}
let capture_started_at = Instant::now();
let image = screen_match
.screen
.capture_area_ignore_area_check(
x1 - screen_match.physical_x,
y1 - screen_match.physical_y,
(x2 - x1) as u32,
(y2 - y1) as u32,
)
.map_err(|error| error.to_string())?;
let capture_ms = capture_started_at.elapsed().as_millis();
save_capture_image(
image.width(),
image.height(),
image.into_raw(),
screenshot_path,
Some(source.clone()),
source.scale_factor,
screen_list_ms,
capture_ms,
total_started_at,
)
}
#[cfg(not(target_os = "windows"))]
fn capture_windows_window(
_hwnd_text: &str,
_screenshot_path: PathBuf,
_screen_list_ms: u128,
_total_started_at: Instant,
) -> Result<CaptureResult, String> {
Err("Windows window capture is only available on Windows".to_string())
}
fn emit_agent_log(app: &tauri::AppHandle, level: &str, message: impl Into<String>) {
let _ = app.emit(
"agent-log",
AgentLog {
level: level.to_string(),
message: message.into(),
timestamp: None,
},
);
}
fn emit_agent_line(app: &tauri::AppHandle, fallback_level: &str, line: String) {
if let Ok(log) = serde_json::from_str::<AgentLog>(&line) {
let _ = app.emit("agent-log", log);
return;
}
emit_agent_log(app, fallback_level, line);
}
#[tauri::command]
fn start_agent(app: tauri::AppHandle, process: tauri::State<AgentProcess>) -> Result<(), String> {
let mut process_guard = process.0.lock().map_err(|error| error.to_string())?;
if process_guard.is_some() {
emit_agent_log(&app, "warning", "agent already running");
return Ok(());
}
let (mut rx, child) = app
.shell()
.sidecar("agent")
.map_err(|error| error.to_string())?
.spawn()
.map_err(|error| error.to_string())?;
*process_guard = Some(child);
emit_agent_log(&app, "info", "agent started");
let app_handle = app.clone();
tauri::async_runtime::spawn(async move {
while let Some(event) = rx.recv().await {
match event {
CommandEvent::Stdout(line) => {
emit_agent_line(
&app_handle,
"info",
String::from_utf8_lossy(&line).trim().to_string(),
);
}
CommandEvent::Stderr(line) => {
emit_agent_line(
&app_handle,
"error",
String::from_utf8_lossy(&line).trim().to_string(),
);
}
CommandEvent::Error(error) => {
emit_agent_log(&app_handle, "error", error);
}
CommandEvent::Terminated(payload) => {
emit_agent_log(
&app_handle,
"warning",
format!("agent exited with code {:?}", payload.code),
);
if let Ok(mut process_guard) = app_handle.state::<AgentProcess>().0.lock() {
*process_guard = None;
}
break;
}
_ => {}
}
}
});
Ok(())
}
#[tauri::command]
fn stop_agent(app: tauri::AppHandle, process: tauri::State<AgentProcess>) -> Result<(), String> {
let mut process_guard = process.0.lock().map_err(|error| error.to_string())?;
if let Some(child) = process_guard.take() {
child.kill().map_err(|error| error.to_string())?;
emit_agent_log(&app, "warning", "agent stopped");
} else {
emit_agent_log(&app, "warning", "agent is not running");
}
Ok(())
}
#[tauri::command]
fn start_vision_stream(
app: tauri::AppHandle,
process: tauri::State<VisionStreamProcess>,
) -> Result<(), String> {
let mut process_guard = process.0.lock().map_err(|error| error.to_string())?;
if process_guard.is_some() {
return Ok(());
}
stop_legacy_python_vision_stream();
let (mut rx, child) = app
.shell()
.sidecar("agent")
.map_err(|error| error.to_string())?
.args(["vision-stream"])
.spawn()
.map_err(|error| error.to_string())?;
*process_guard = Some(child);
let app_handle = app.clone();
tauri::async_runtime::spawn(async move {
while let Some(event) = rx.recv().await {
match event {
CommandEvent::Stdout(line) => {
emit_agent_line(
&app_handle,
"info",
String::from_utf8_lossy(&line).trim().to_string(),
);
}
CommandEvent::Stderr(line) => {
emit_agent_line(
&app_handle,
"error",
String::from_utf8_lossy(&line).trim().to_string(),
);
}
CommandEvent::Error(error) => {
emit_agent_log(
&app_handle,
"error",
format!("vision stream error: {error}"),
);
}
CommandEvent::Terminated(payload) => {
emit_agent_log(
&app_handle,
"warning",
format!("vision stream exited with code {:?}", payload.code),
);
if let Ok(mut process_guard) =
app_handle.state::<VisionStreamProcess>().0.lock()
{
*process_guard = None;
}
break;
}
_ => {}
}
}
});
Ok(())
}
fn stop_legacy_python_vision_stream() {
let Ok(cwd) = std::env::current_dir() else {
return;
};
let candidates = [
cwd.join("wechat_vision"),
cwd.join("..").join("wechat_vision"),
cwd.join("..").join("..").join("wechat_vision"),
];
for dir in candidates {
let pid_file = dir.join("ouptsw").join("wechat_algorithm_live.pid");
let stop_script = dir.join("stop_wechat_algorithm_live.sh");
if pid_file.exists() && stop_script.exists() {
let _ = Command::new("bash")
.arg(stop_script)
.current_dir(dir)
.output();
return;
}
}
}
#[tauri::command]
fn stop_vision_stream(
_app: tauri::AppHandle,
process: tauri::State<VisionStreamProcess>,
) -> Result<(), String> {
let mut process_guard = process.0.lock().map_err(|error| error.to_string())?;
if let Some(child) = process_guard.take() {
child.kill().map_err(|error| error.to_string())?;
}
Ok(())
}
#[tauri::command]
async fn capture_screen(app: tauri::AppHandle) -> Result<CaptureResult, String> {
let screenshot_path = screenshot_path(&app)?;
tauri::async_runtime::spawn_blocking(move || {
let total_started_at = Instant::now();
let screen_list_started_at = Instant::now();
let screen = Screen::all()
.map_err(|error| error.to_string())?
.into_iter()
.next()
.ok_or_else(|| "no screen found".to_string())?;
let screen_list_ms = screen_list_started_at.elapsed().as_millis();
let capture_started_at = Instant::now();
let image = screen.capture().map_err(|error| error.to_string())?;
let capture_ms = capture_started_at.elapsed().as_millis();
let width = image.width();
let height = image.height();
save_capture_image(
width,
height,
image.into_raw(),
screenshot_path,
None,
screen.display_info.scale_factor as f64,
screen_list_ms,
capture_ms,
total_started_at,
)
})
.await
.map_err(|error| error.to_string())?
}
#[tauri::command]
async fn list_capture_sources() -> Result<Vec<CaptureSource>, String> {
tauri::async_runtime::spawn_blocking(move || {
let mut sources = Vec::new();
for (index, monitor) in Monitor::all()
.map_err(|error| error.to_string())?
.into_iter()
.enumerate()
{
let width = monitor.width().map_err(|error| error.to_string())?;
let height = monitor.height().map_err(|error| error.to_string())?;
let x = monitor.x().map_err(|error| error.to_string())?;
let y = monitor.y().map_err(|error| error.to_string())?;
let scale_factor = monitor.scale_factor().map_err(|error| error.to_string())? as f64;
sources.push(CaptureSource {
id: format!("display:{index}"),
kind: "display".to_string(),
label: format!("桌面 {} · {}x{}", index + 1, width, height),
app_name: None,
title: None,
pid: None,
x,
y,
width,
height,
scale_factor,
});
}
if let Ok(windows) = Window::all() {
for window in windows {
let width = match window.width() {
Ok(width) => width,
Err(_) => continue,
};
let height = match window.height() {
Ok(height) => height,
Err(_) => continue,
};
if window.is_minimized().unwrap_or(true) || width < 80 || height < 80 {
continue;
}
let app_name = window.app_name().unwrap_or_default();
let title = window.title().unwrap_or_default();
if app_name.is_empty() && title.is_empty() {
continue;
}
sources.push(CaptureSource {
id: format!("window:{}", window.id().map_err(|error| error.to_string())?),
kind: "window".to_string(),
label: format!("{} · {}", app_name, title),
app_name: Some(app_name),
title: Some(title),
pid: window.pid().ok(),
x: window.x().unwrap_or(0),
y: window.y().unwrap_or(0),
width,
height,
scale_factor: 1.0,
});
}
}
for source in windows_capture_sources() {
let duplicate = sources.iter().any(|item| {
item.pid == source.pid
&& item.title == source.title
&& (item.x - source.x).abs() <= 2
&& (item.y - source.y).abs() <= 2
&& item.width.abs_diff(source.width) <= 2
&& item.height.abs_diff(source.height) <= 2
});
if !duplicate {
sources.push(source);
}
}
Ok(sources)
})
.await
.map_err(|error| error.to_string())?
}
#[tauri::command]
async fn capture_source(app: tauri::AppHandle, source_id: String) -> Result<CaptureResult, String> {
let screenshot_path = screenshot_path(&app)?;
tauri::async_runtime::spawn_blocking(move || {
let total_started_at = Instant::now();
let screen_list_started_at = Instant::now();
if let Some(index_text) = source_id.strip_prefix("display:") {
let index = index_text
.parse::<usize>()
.map_err(|error| error.to_string())?;
let monitors = Monitor::all().map_err(|error| error.to_string())?;
let monitor = monitors
.into_iter()
.nth(index)
.ok_or_else(|| "display source not found".to_string())?;
let screen_list_ms = screen_list_started_at.elapsed().as_millis();
let width = monitor.width().map_err(|error| error.to_string())?;
let height = monitor.height().map_err(|error| error.to_string())?;
let scale_factor = monitor.scale_factor().map_err(|error| error.to_string())? as f64;
let source = CaptureSource {
id: source_id,
kind: "display".to_string(),
label: format!("桌面 {} · {}x{}", index + 1, width, height),
app_name: None,
title: None,
pid: None,
x: monitor.x().map_err(|error| error.to_string())?,
y: monitor.y().map_err(|error| error.to_string())?,
width,
height,
scale_factor,
};
let capture_started_at = Instant::now();
let image = monitor.capture_image().map_err(|error| error.to_string())?;
let capture_ms = capture_started_at.elapsed().as_millis();
let image_width = image.width();
let image_height = image.height();
return save_capture_image(
image_width,
image_height,
image.into_raw(),
screenshot_path,
Some(source.clone()),
source.scale_factor,
screen_list_ms,
capture_ms,
total_started_at,
);
}
if let Some(window_id_text) = source_id.strip_prefix("window:") {
let window_id = window_id_text
.parse::<u32>()
.map_err(|error| error.to_string())?;
let windows = Window::all().map_err(|error| error.to_string())?;
let window = windows
.into_iter()
.find(|item| item.id().ok() == Some(window_id))
.ok_or_else(|| "window source not found".to_string())?;
let screen_list_ms = screen_list_started_at.elapsed().as_millis();
let window_width = window.width().map_err(|error| error.to_string())?;
let window_height = window.height().map_err(|error| error.to_string())?;
let app_name = window.app_name().unwrap_or_default();
let title = window.title().unwrap_or_default();
let capture_started_at = Instant::now();
let image = window.capture_image().map_err(|error| error.to_string())?;
let capture_ms = capture_started_at.elapsed().as_millis();
let image_width = image.width();
let image_height = image.height();
let scale_factor = if window_width > 0 {
image_width as f64 / window_width as f64
} else {
1.0
};
let source = CaptureSource {
id: source_id,
kind: "window".to_string(),
label: format!("{} · {}", app_name, title),
app_name: Some(app_name),
title: Some(title),
pid: window.pid().ok(),
x: window.x().unwrap_or(0),
y: window.y().unwrap_or(0),
width: window_width,
height: window_height,
scale_factor,
};
return save_capture_image(
image_width,
image_height,
image.into_raw(),
screenshot_path,
Some(source.clone()),
source.scale_factor,
screen_list_ms,
capture_ms,
total_started_at,
);
}
if let Some(hwnd_text) = source_id.strip_prefix("win-window:") {
let screen_list_ms = screen_list_started_at.elapsed().as_millis();
return capture_windows_window(
hwnd_text,
screenshot_path,
screen_list_ms,
total_started_at,
);
}
Err("unsupported capture source".to_string())
})
.await
.map_err(|error| error.to_string())?
}
#[tauri::command]
fn read_screenshot_bytes(app: tauri::AppHandle, path: String) -> Result<Vec<u8>, String> {
let screenshots_dir = data_dir(&app)?.join("screenshots");
fs::create_dir_all(&screenshots_dir).map_err(|error| error.to_string())?;
let screenshots_dir = screenshots_dir
.canonicalize()
.map_err(|error| error.to_string())?;
let requested_path = PathBuf::from(path)
.canonicalize()
.map_err(|error| error.to_string())?;
if !requested_path.starts_with(&screenshots_dir) {
return Err("screenshot path is outside app screenshot directory".to_string());
}
fs::read(requested_path).map_err(|error| error.to_string())
}
pub(crate) fn write_annotation_file_atomic(
app: &tauri::AppHandle,
annotation: &AnnotationFile,
) -> Result<PathBuf, String> {
use std::io::Write;
let path = regions_path(app)?;
let bytes = serde_json::to_vec_pretty(annotation).map_err(|error| error.to_string())?;
let parent = path
.parent()
.ok_or_else(|| "regions path has no parent directory".to_string())?;
let nonce = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_err(|error| error.to_string())?
.as_nanos();
let temp_path = parent.join(format!(".wechat.json.{}.{}.tmp", std::process::id(), nonce));
let result = (|| -> Result<(), String> {
let mut file = fs::OpenOptions::new()
.create_new(true)
.write(true)
.open(&temp_path)
.map_err(|error| error.to_string())?;
file.write_all(&bytes).map_err(|error| error.to_string())?;
file.sync_all().map_err(|error| error.to_string())?;
drop(file);
#[cfg(target_os = "windows")]
{
use std::os::windows::ffi::OsStrExt;
use windows_sys::Win32::Storage::FileSystem::{
MoveFileExW, MOVEFILE_REPLACE_EXISTING, MOVEFILE_WRITE_THROUGH,
};
let source = temp_path
.as_os_str()
.encode_wide()
.chain(std::iter::once(0))
.collect::<Vec<_>>();
let destination = path
.as_os_str()
.encode_wide()
.chain(std::iter::once(0))
.collect::<Vec<_>>();
let moved = unsafe {
MoveFileExW(
source.as_ptr(),
destination.as_ptr(),
MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH,
)
};
if moved == 0 {
return Err(std::io::Error::last_os_error().to_string());
}
}
#[cfg(not(target_os = "windows"))]
fs::rename(&temp_path, &path).map_err(|error| error.to_string())?;
Ok(())
})();
if result.is_err() {
let _ = fs::remove_file(&temp_path);
}
result?;
Ok(path)
}
#[tauri::command]
fn save_regions(app: tauri::AppHandle, annotation: AnnotationFile) -> Result<String, String> {
write_annotation_file_atomic(&app, &annotation).map(|path| path.to_string_lossy().to_string())
}
#[tauri::command]
fn load_regions(app: tauri::AppHandle) -> Result<Option<AnnotationFile>, String> {
let path = regions_path(&app)?;
if !path.exists() {
return Ok(None);
}
let json = fs::read_to_string(&path).map_err(|error| error.to_string())?;
serde_json::from_str(&json)
.map(Some)
.map_err(|error| error.to_string())
}
#[tauri::command]
fn close_current_window(window: tauri::WebviewWindow) -> Result<(), String> {
window.close().map_err(|error| error.to_string())
}
#[tauri::command]
fn exit_application(app: tauri::AppHandle) {
app.exit(0);
}
#[cfg(target_os = "windows")]
fn apply_native_window_corner(window: &tauri::WebviewWindow) -> Result<(), String> {
use std::ffi::c_void;
use std::mem::size_of_val;
use windows_sys::Win32::Graphics::Dwm::{
DwmSetWindowAttribute, DWMWA_WINDOW_CORNER_PREFERENCE, DWMWCP_ROUND,
};
let hwnd = window.hwnd().map_err(|error| error.to_string())?;
let preference = DWMWCP_ROUND;
let result = unsafe {
DwmSetWindowAttribute(
hwnd.0 as _,
DWMWA_WINDOW_CORNER_PREFERENCE as u32,
&preference as *const _ as *const c_void,
size_of_val(&preference) as u32,
)
};
if result < 0 {
return Err(format!(
"DwmSetWindowAttribute(DWMWA_WINDOW_CORNER_PREFERENCE) failed: 0x{:08X}",
result as u32
));
}
Ok(())
}
#[cfg(not(target_os = "windows"))]
fn apply_native_window_corner(_window: &tauri::WebviewWindow) -> Result<(), String> {
Ok(())
}
#[tauri::command]
async fn open_popup_window(app: tauri::AppHandle, route: String) -> Result<(), String> {
let route = if route.starts_with('/') {
route
} else {
format!("/{route}")
};
let route_path = route.split('?').next().unwrap_or(route.as_str());
let template_label = format!(
"popup-{}",
route_path.trim_start_matches('/').replace('/', "-")
);
let window_label = format!(
"popup-{}",
route
.trim_start_matches('/')
.chars()
.map(|character| {
if character.is_ascii_alphanumeric() || matches!(character, '-' | '_') {
character
} else {
'-'
}
})
.collect::<String>()
);
if route_path == "/annotate" {
if let Some(window) = app.get_webview_window(&window_label) {
let _ = window.close();
}
} else if let Some(window) = app.get_webview_window(&window_label) {
if let Err(error) = apply_native_window_corner(&window) {
eprintln!("failed to apply native window corner: {error}");
}
window.center().map_err(|error| error.to_string())?;
window.show().map_err(|error| error.to_string())?;
window.set_focus().map_err(|error| error.to_string())?;
return Ok(());
}
let mut window_config = app
.config()
.app
.windows
.iter()
.find(|window| window.label == template_label)
.cloned()
.ok_or_else(|| format!("popup window route is not configured: {route_path}"))?;
window_config.label = window_label;
window_config.url = WebviewUrl::App(format!("index.html#{route}").into());
window_config.visible = false;
let _window = WebviewWindowBuilder::from_config(&app, &window_config)
.map_err(|error| error.to_string())?
.on_page_load(move |window, payload| {
if matches!(payload.event(), PageLoadEvent::Finished) {
if let Err(error) = apply_native_window_corner(&window) {
eprintln!("failed to apply native popup corner: {error}");
}
if let Err(error) = window.center() {
eprintln!("failed to center loaded popup window: {error}");
}
if let Err(error) = window.show() {
eprintln!("failed to show loaded popup window: {error}");
}
if let Err(error) = window.set_focus() {
eprintln!("failed to focus loaded popup window: {error}");
}
}
})
.build()
.map_err(|error| error.to_string())?;
Ok(())
}
#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
tauri::Builder::default()
.manage(AgentProcess(Mutex::new(None)))
.manage(VisionStreamProcess(Mutex::new(None)))
.manage(window_capture::OverlayState::default())
.plugin(tauri_plugin_shell::init())
.setup(|app| {
if cfg!(debug_assertions) {
app.handle().plugin(
tauri_plugin_log::Builder::default()
.level(log::LevelFilter::Info)
.build(),
)?;
}
#[cfg(target_os = "windows")]
{
window_capture::create_overlay_window(app.handle())?;
if let Some(main_window) = app.get_webview_window("main") {
if let Err(error) = apply_native_window_corner(&main_window) {
eprintln!("failed to apply native main-window corner: {error}");
}
}
}
Ok(())
})
.invoke_handler(tauri::generate_handler![
close_current_window,
exit_application,
capture_source,
capture_screen,
list_capture_sources,
load_regions,
open_popup_window,
read_screenshot_bytes,
save_regions,
start_agent,
start_vision_stream,
stop_vision_stream,
stop_agent,
window_capture::add_annotation,
window_capture::delete_annotation,
window_capture::enter_window_select_mode,
window_capture::get_overlay_frame,
window_capture::hide_overlay,
window_capture::select_target_window,
window_capture::update_annotation,
window_capture::update_annotation_geometry
])
.run(tauri::generate_context!())
.expect("error while running tauri application");
}