pi-gui/ui/src-tauri/src/lib.rs

mod bridge;

use serde_json::Value;
use std::sync::Mutex;
use tauri::{async_runtime::JoinHandle, AppHandle, Emitter, Manager, State};

struct Subscription(Mutex<Option<JoinHandle<()>>>);

#[derive(Debug, PartialEq, Eq)]
enum WindowAction {
    Hide,
    ShowAndFocus,
}

fn window_action(args: &[String], is_visible: bool) -> WindowAction {
    if args.iter().any(|arg| arg == "--hide")
        || (args.iter().any(|arg| arg == "--toggle") && is_visible)
    {
        WindowAction::Hide
    } else {
        WindowAction::ShowAndFocus
    }
}

fn socket_path() -> Result<String, String> {
    bridge::default_socket_path()
}

#[tauri::command]
async fn list_agents() -> Result<Value, String> {
    bridge::list_agents(&socket_path()?).await
}

#[tauri::command]
async fn list_directories() -> Result<Value, String> {
    bridge::list_directories(&socket_path()?).await
}

#[tauri::command]
async fn select_worktree(worktree_path: String) -> Result<Value, String> {
    bridge::select_worktree(&socket_path()?, &worktree_path).await
}

#[tauri::command]
async fn forget_directory(worktree_path: String) -> Result<Value, String> {
    bridge::forget_directory(&socket_path()?, &worktree_path).await
}

#[tauri::command]
async fn load_agent(agent_id: String) -> Result<Value, String> {
    bridge::load_agent(&socket_path()?, &agent_id).await
}

#[tauri::command]
async fn list_sessions(agent_id: String) -> Result<Value, String> {
    bridge::list_sessions(&socket_path()?, &agent_id).await
}

#[tauri::command]
async fn switch_session(agent_id: String, session_path: String) -> Result<Value, String> {
    bridge::switch_session(&socket_path()?, &agent_id, &session_path).await
}

#[tauri::command]
async fn new_session(agent_id: String) -> Result<Value, String> {
    bridge::new_session(&socket_path()?, &agent_id).await
}

#[tauri::command]
async fn submit_prompt(agent_id: String, message: String) -> Result<(), String> {
    bridge::submit_prompt(&socket_path()?, &agent_id, &message).await
}

#[tauri::command]
async fn abort(agent_id: String) -> Result<(), String> {
    bridge::abort(&socket_path()?, &agent_id).await
}

#[tauri::command]
async fn retry(agent_id: String) -> Result<(), String> {
    bridge::command(&socket_path()?, &agent_id, "retry").await
}

#[tauri::command]
async fn restart(agent_id: String) -> Result<(), String> {
    bridge::command(&socket_path()?, &agent_id, "restart").await
}

#[tauri::command]
async fn set_model(agent_id: String, provider: String, model_id: String) -> Result<(), String> {
    bridge::set_model(&socket_path()?, &agent_id, &provider, &model_id).await
}

#[tauri::command]
async fn set_thinking_level(agent_id: String, level: String) -> Result<(), String> {
    bridge::set_thinking_level(&socket_path()?, &agent_id, &level).await
}

#[tauri::command]
async fn respond_to_extension(
    agent_id: String,
    request_id: String,
    response: Value,
) -> Result<(), String> {
    bridge::respond_to_extension(&socket_path()?, &agent_id, &request_id, response).await
}

#[tauri::command]
fn subscribe_agent(
    app: AppHandle,
    subscriptions: State<'_, Subscription>,
    agent_id: String,
    cursor: u64,
) -> Result<(), String> {
    let socket = socket_path()?;
    if let Some(subscription) = subscriptions
        .0
        .lock()
        .map_err(|_| "Could not update bridge subscription".to_owned())?
        .take()
    {
        subscription.abort();
    }
    let event_app = app.clone();
    let task = tauri::async_runtime::spawn(async move {
        if let Err(error) =
            bridge::subscribe(socket, agent_id.clone(), cursor, event_app.clone()).await
        {
            let _ = event_app.emit(
                "bridge-error",
                serde_json::json!({ "agentId": agent_id, "message": error }),
            );
        }
    });
    let mut current_subscription = subscriptions
        .0
        .lock()
        .map_err(|_| "Could not retain bridge subscription".to_owned())?;
    *current_subscription = Some(task);
    Ok(())
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    let builder = tauri::Builder::default()
        .plugin(tauri_plugin_dialog::init())
        .manage(Subscription(Mutex::new(None)))
        .plugin(tauri_plugin_single_instance::init(|app, args, _cwd| {
            if let Some(window) = app.get_webview_window("main") {
                match window_action(&args, window.is_visible().unwrap_or(false)) {
                    WindowAction::Hide => {
                        let _ = window.hide();
                    }
                    WindowAction::ShowAndFocus => {
                        let _ = window.show();
                        let _ = window.set_focus();
                    }
                }
            }
        }))
        .invoke_handler(tauri::generate_handler![
            list_agents,
            list_directories,
            select_worktree,
            forget_directory,
            load_agent,
            list_sessions,
            switch_session,
            new_session,
            submit_prompt,
            abort,
            retry,
            restart,
            set_model,
            set_thinking_level,
            respond_to_extension,
            subscribe_agent
        ]);
    let context = tauri::generate_context!();
    if let Err(error) = builder.run(context) {
        eprintln!("Pi Status UI exited: {error}");
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn toggle_hides_a_visible_window_and_shows_a_hidden_window() {
        let toggle = vec!["--toggle".to_owned()];

        assert_eq!(window_action(&toggle, true), WindowAction::Hide);
        assert_eq!(window_action(&toggle, false), WindowAction::ShowAndFocus);
        assert_eq!(
            window_action(&vec!["--show".to_owned()], true),
            WindowAction::ShowAndFocus
        );
    }
}