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

use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fs::{self, OpenOptions};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};

const VERSION: u32 = 1;
const MAX_ITEMS: usize = 256;
const MAX_TOTAL_ITEMS: usize = 2048;
const MAX_DRAFT_BYTES: usize = 128 * 1024;
const MAX_TOTAL_DRAFT_BYTES: usize = 1024 * 1024;
const MAX_TOTAL_STRING_BYTES: usize = 1536 * 1024;
const MAX_FILE_BYTES: usize = 2 * 1024 * 1024;

#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct WorkspaceCursor {
    pub bridge_instance_id: String,
    pub seq: u64,
}

#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ScrollState {
    pub at_bottom: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub scroll_top: Option<f64>,
}

#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct LastSeenState {
    pub bridge_instance_id: String,
    pub seq: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub last_activity: Option<String>,
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct QuickPaneSettings {
    pub workspace: String,
    pub provider: String,
    pub model: String,
    pub follow_up_provider: String,
    pub follow_up_model: String,
    pub instructions: String,
    pub maximum_seconds: u32,
    pub maximum_tool_calls: u32,
}

impl Default for QuickPaneSettings {
    fn default() -> Self {
        Self {
            workspace: String::new(),
            provider: String::new(),
            model: String::new(),
            follow_up_provider: String::new(),
            follow_up_model: String::new(),
            instructions: "Answer quickly. Do not use subagents. Prefer read-only inspection. Escalate mutations, multi-step work, or work likely to exceed budget.".to_owned(),
            maximum_seconds: 60,
            maximum_tool_calls: 6,
        }
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct UiStateV1 {
    pub version: u32,
    #[serde(default)]
    pub directory_order: Vec<String>,
    #[serde(default)]
    pub runtime_order_by_directory: BTreeMap<String, Vec<String>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub selected_runtime_id: Option<String>,
    #[serde(default)]
    pub drafts: BTreeMap<String, String>,
    #[serde(default)]
    pub scroll: BTreeMap<String, ScrollState>,
    #[serde(default)]
    pub last_seen: BTreeMap<String, LastSeenState>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub workspace_cursor: Option<WorkspaceCursor>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub interface_scale: Option<f64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub default_new_session_worktree: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub dismissed_collision_warning: Option<bool>,
    #[serde(default)]
    pub quick_pane: QuickPaneSettings,
}

impl Default for UiStateV1 {
    fn default() -> Self {
        Self {
            version: VERSION,
            directory_order: Vec::new(),
            runtime_order_by_directory: BTreeMap::new(),
            selected_runtime_id: None,
            drafts: BTreeMap::new(),
            scroll: BTreeMap::new(),
            last_seen: BTreeMap::new(),
            workspace_cursor: None,
            interface_scale: None,
            default_new_session_worktree: None,
            dismissed_collision_warning: None,
            quick_pane: QuickPaneSettings::default(),
        }
    }
}

fn bounded(value: &str, field: &str) -> Result<(), String> {
    if value.is_empty() || value.len() > 4096 {
        Err(format!("{field} must contain 1 to 4096 bytes"))
    } else {
        Ok(())
    }
}

pub fn validate(state: &UiStateV1) -> Result<(), String> {
    if state.version != VERSION {
        return Err(format!("Unsupported UI state version: {}", state.version));
    }
    if state.directory_order.len() > MAX_ITEMS
        || state.runtime_order_by_directory.len() > MAX_ITEMS
        || state.drafts.len() > MAX_ITEMS
        || state.scroll.len() > MAX_ITEMS
        || state.last_seen.len() > MAX_ITEMS
    {
        return Err("UI state contains too many entries".to_owned());
    }
    let mut total_items = state.directory_order.len()
        + state.runtime_order_by_directory.len()
        + state.drafts.len()
        + state.scroll.len()
        + state.last_seen.len();
    let mut total_string_bytes = 0usize;
    for path in &state.directory_order {
        bounded(path, "directory path")?;
        total_string_bytes = total_string_bytes.saturating_add(path.len());
    }
    for (path, runtimes) in &state.runtime_order_by_directory {
        bounded(path, "directory path")?;
        total_string_bytes = total_string_bytes.saturating_add(path.len());
        if runtimes.len() > MAX_ITEMS {
            return Err("UI state contains too many runtime entries".to_owned());
        }
        total_items = total_items.saturating_add(runtimes.len());
        for runtime_id in runtimes {
            bounded(runtime_id, "runtime id")?;
            total_string_bytes = total_string_bytes.saturating_add(runtime_id.len());
        }
    }
    if total_items > MAX_TOTAL_ITEMS {
        return Err("UI state contains too many aggregate entries".to_owned());
    }
    if let Some(runtime_id) = &state.selected_runtime_id {
        bounded(runtime_id, "selected runtime id")?;
        total_string_bytes = total_string_bytes.saturating_add(runtime_id.len());
    }
    if let Some(worktree) = &state.default_new_session_worktree {
        bounded(worktree, "default new-session worktree")?;
        if !Path::new(worktree).is_absolute() {
            return Err("Default new-session worktree must be absolute".to_owned());
        }
        total_string_bytes = total_string_bytes.saturating_add(worktree.len());
    }
    let mut total_draft_bytes = 0usize;
    for (runtime_id, draft) in &state.drafts {
        bounded(runtime_id, "draft runtime id")?;
        if draft.len() > MAX_DRAFT_BYTES {
            return Err("A UI draft exceeds the size limit".to_owned());
        }
        total_draft_bytes = total_draft_bytes.saturating_add(draft.len());
        total_string_bytes = total_string_bytes
            .saturating_add(runtime_id.len())
            .saturating_add(draft.len());
    }
    if total_draft_bytes > MAX_TOTAL_DRAFT_BYTES {
        return Err("UI drafts exceed the total size limit".to_owned());
    }
    for (runtime_id, scroll) in &state.scroll {
        bounded(runtime_id, "runtime id")?;
        total_string_bytes = total_string_bytes.saturating_add(runtime_id.len());
        if scroll
            .scroll_top
            .is_some_and(|value| !value.is_finite() || value < 0.0)
        {
            return Err("Scroll position is invalid".to_owned());
        }
    }
    for (runtime_id, last_seen) in &state.last_seen {
        bounded(runtime_id, "runtime id")?;
        bounded(
            &last_seen.bridge_instance_id,
            "last-seen bridge instance id",
        )?;
        total_string_bytes = total_string_bytes
            .saturating_add(runtime_id.len())
            .saturating_add(last_seen.bridge_instance_id.len());
        if let Some(activity) = &last_seen.last_activity {
            bounded(activity, "last activity")?;
            total_string_bytes = total_string_bytes.saturating_add(activity.len());
        }
    }
    if let Some(cursor) = &state.workspace_cursor {
        bounded(&cursor.bridge_instance_id, "bridge instance id")?;
        total_string_bytes = total_string_bytes.saturating_add(cursor.bridge_instance_id.len());
    }
    if total_string_bytes > MAX_TOTAL_STRING_BYTES {
        return Err("UI state strings exceed the total size limit".to_owned());
    }
    let quick = &state.quick_pane;
    let quick_is_default = quick.workspace.is_empty()
        && quick.provider.is_empty()
        && quick.model.is_empty()
        && quick.follow_up_provider.is_empty()
        && quick.follow_up_model.is_empty();
    if !quick_is_default {
        bounded(&quick.workspace, "quick-pane workspace")?;
        if !Path::new(&quick.workspace).is_absolute() || !Path::new(&quick.workspace).is_dir() {
            return Err("Quick-pane workspace must be an existing absolute directory".to_owned());
        }
        bounded(&quick.provider, "quick-pane provider")?;
        bounded(&quick.model, "quick-pane model")?;
        bounded(&quick.follow_up_provider, "quick-pane follow-up provider")?;
        bounded(&quick.follow_up_model, "quick-pane follow-up model")?;
    }
    if quick.instructions.is_empty() || quick.instructions.len() > 4096 {
        return Err("Quick-pane instructions must contain 1 to 4096 bytes".to_owned());
    }
    if !(1..=600).contains(&quick.maximum_seconds) || !(1..=64).contains(&quick.maximum_tool_calls)
    {
        return Err("Quick-pane limits are outside the supported range".to_owned());
    }
    total_string_bytes = total_string_bytes
        .saturating_add(quick.workspace.len())
        .saturating_add(quick.provider.len())
        .saturating_add(quick.model.len())
        .saturating_add(quick.follow_up_provider.len())
        .saturating_add(quick.follow_up_model.len())
        .saturating_add(quick.instructions.len());
    if total_string_bytes > MAX_TOTAL_STRING_BYTES {
        return Err("UI state strings exceed the total size limit".to_owned());
    }
    if let Some(scale) = state.interface_scale {
        if !scale.is_finite() || !(0.5..=2.0).contains(&scale) {
            return Err("Interface scale is outside the supported range".to_owned());
        }
    }
    Ok(())
}

pub fn state_path(app_data_dir: &Path) -> PathBuf {
    app_data_dir.join("ui-workspace-v1.json")
}

pub fn load_from(app_data_dir: &Path) -> Result<UiStateV1, String> {
    let path = state_path(app_data_dir);
    let mut file = match fs::File::open(&path) {
        Ok(file) => file,
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
            return Ok(UiStateV1::default())
        }
        Err(error) => return Err(format!("Could not read UI state: {error}")),
    };
    let metadata = file
        .metadata()
        .map_err(|error| format!("Could not inspect UI state: {error}"))?;
    if metadata.len() > MAX_FILE_BYTES as u64 {
        return Err("UI state file exceeds the size limit".to_owned());
    }
    let mut content = Vec::with_capacity(metadata.len() as usize);
    Read::by_ref(&mut file)
        .take((MAX_FILE_BYTES + 1) as u64)
        .read_to_end(&mut content)
        .map_err(|error| format!("Could not read UI state: {error}"))?;
    if content.len() > MAX_FILE_BYTES {
        return Err("UI state file exceeds the size limit".to_owned());
    }
    let state: UiStateV1 = serde_json::from_slice(&content)
        .map_err(|error| format!("UI state is corrupt or unsupported: {error}"))?;
    validate(&state)?;
    Ok(state)
}

pub fn save_to(app_data_dir: &Path, state: &UiStateV1) -> Result<(), String> {
    validate(state)?;
    fs::create_dir_all(app_data_dir)
        .map_err(|error| format!("Could not create app data: {error}"))?;
    set_mode(app_data_dir, 0o700)?;
    let target = state_path(app_data_dir);
    let nonce = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_nanos();
    let temporary = app_data_dir.join(format!(".ui-workspace-{}-{nonce}.tmp", std::process::id()));
    let encoded = serde_json::to_vec_pretty(state).map_err(|error| error.to_string())?;
    if encoded.len().saturating_add(1) > MAX_FILE_BYTES {
        return Err("UI state file exceeds the size limit".to_owned());
    }
    let result = (|| {
        let mut file = create_private_file(&temporary)?;
        file.write_all(&encoded)
            .map_err(|error| format!("Could not write UI state: {error}"))?;
        file.write_all(b"\n")
            .map_err(|error| format!("Could not write UI state: {error}"))?;
        file.sync_all()
            .map_err(|error| format!("Could not sync UI state: {error}"))?;
        fs::rename(&temporary, &target)
            .map_err(|error| format!("Could not replace UI state: {error}"))?;
        set_mode(&target, 0o600)?;
        sync_directory(app_data_dir)?;
        Ok(())
    })();
    if result.is_err() {
        let _ = fs::remove_file(&temporary);
    }
    result
}

#[cfg(unix)]
fn sync_directory(path: &Path) -> Result<(), String> {
    let directory = OpenOptions::new()
        .read(true)
        .open(path)
        .map_err(|error| format!("Could not open app data directory for sync: {error}"))?;
    directory
        .sync_all()
        .map_err(|error| format!("Could not sync app data directory: {error}"))
}

#[cfg(not(unix))]
fn sync_directory(_path: &Path) -> Result<(), String> {
    Ok(())
}

#[cfg(unix)]
fn create_private_file(path: &Path) -> Result<std::fs::File, String> {
    use std::os::unix::fs::OpenOptionsExt;
    OpenOptions::new()
        .create_new(true)
        .write(true)
        .mode(0o600)
        .open(path)
        .map_err(|error| format!("Could not create temporary UI state: {error}"))
}

#[cfg(not(unix))]
fn create_private_file(path: &Path) -> Result<std::fs::File, String> {
    OpenOptions::new()
        .create_new(true)
        .write(true)
        .open(path)
        .map_err(|error| format!("Could not create temporary UI state: {error}"))
}

#[cfg(unix)]
fn set_mode(path: &Path, mode: u32) -> Result<(), String> {
    use std::os::unix::fs::PermissionsExt;
    fs::set_permissions(path, fs::Permissions::from_mode(mode))
        .map_err(|error| format!("Could not secure UI state path: {error}"))
}

#[cfg(not(unix))]
fn set_mode(_path: &Path, _mode: u32) -> Result<(), String> {
    Ok(())
}

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

    fn temp_dir(name: &str) -> PathBuf {
        std::env::temp_dir().join(format!(
            "pi-ui-state-{name}-{}-{}",
            std::process::id(),
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ))
    }

    #[test]
    fn owner_only_round_trip_and_missing_default() {
        let root = temp_dir("roundtrip");
        assert_eq!(load_from(&root).unwrap(), UiStateV1::default());
        let mut state = UiStateV1 {
            selected_runtime_id: Some("runtime-1".into()),
            interface_scale: Some(1.1),
            ..UiStateV1::default()
        };
        state.drafts.insert("runtime-1".into(), "draft".into());
        save_to(&root, &state).unwrap();
        assert_eq!(load_from(&root).unwrap(), state);
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            assert_eq!(
                fs::metadata(&root).unwrap().permissions().mode() & 0o777,
                0o700
            );
            assert_eq!(
                fs::metadata(state_path(&root))
                    .unwrap()
                    .permissions()
                    .mode()
                    & 0o777,
                0o600
            );
        }
        fs::remove_dir_all(root).unwrap();
    }

    fn temporary_files(root: &Path) -> Vec<PathBuf> {
        fs::read_dir(root)
            .unwrap()
            .filter_map(Result::ok)
            .map(|entry| entry.path())
            .filter(|path| {
                path.file_name()
                    .and_then(|name| name.to_str())
                    .is_some_and(|name| {
                        name.starts_with(".ui-workspace-") && name.ends_with(".tmp")
                    })
            })
            .collect()
    }

    #[test]
    fn quick_pane_settings_validate_complete_existing_workspace() {
        let root = temp_dir("quick-pane");
        fs::create_dir_all(&root).unwrap();
        let mut state = UiStateV1::default();
        state.quick_pane.workspace = root.to_string_lossy().into_owned();
        state.quick_pane.provider = "provider".into();
        state.quick_pane.model = "quick-model".into();
        state.quick_pane.follow_up_provider = "provider".into();
        state.quick_pane.follow_up_model = "follow-up-model".into();
        assert!(validate(&state).is_ok());
        state.quick_pane.maximum_seconds = 0;
        assert!(validate(&state).is_err());
        fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn second_save_replaces_valid_state_without_leaving_temporary_files() {
        let root = temp_dir("replace");
        let first = UiStateV1 {
            selected_runtime_id: Some("runtime-1".into()),
            ..UiStateV1::default()
        };
        let second = UiStateV1 {
            selected_runtime_id: Some("runtime-2".into()),
            interface_scale: Some(1.2),
            ..UiStateV1::default()
        };
        save_to(&root, &first).unwrap();
        save_to(&root, &second).unwrap();
        assert_eq!(load_from(&root).unwrap(), second);
        assert!(temporary_files(&root).is_empty());
        fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn rejects_corruption_versions_unknown_fields_and_limits_without_overwrite() {
        let root = temp_dir("invalid");
        fs::create_dir_all(&root).unwrap();
        let path = state_path(&root);
        for invalid in [
            "{broken",
            r#"{"version":2}"#,
            r#"{"version":1,"transcript":[]}"#,
        ] {
            fs::write(&path, invalid).unwrap();
            assert!(load_from(&root).is_err());
            assert_eq!(fs::read_to_string(&path).unwrap(), invalid);
        }

        let baseline = UiStateV1 {
            selected_runtime_id: Some("baseline".into()),
            ..UiStateV1::default()
        };
        save_to(&root, &baseline).unwrap();
        let original = fs::read(&path).unwrap();
        let mut oversized = UiStateV1::default();
        oversized
            .drafts
            .insert("runtime".into(), "x".repeat(MAX_DRAFT_BYTES + 1));
        assert!(save_to(&root, &oversized).is_err());
        assert_eq!(fs::read(&path).unwrap(), original);
        assert!(temporary_files(&root).is_empty());
        fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn rejects_aggregate_and_raw_file_size_limits() {
        let root = temp_dir("aggregate");
        fs::create_dir_all(&root).unwrap();
        let mut aggregate = UiStateV1::default();
        for directory in 0..9 {
            aggregate.runtime_order_by_directory.insert(
                format!("/directory-{directory}"),
                (0..MAX_ITEMS)
                    .map(|runtime| format!("runtime-{directory}-{runtime}"))
                    .collect(),
            );
        }
        assert!(validate(&aggregate)
            .unwrap_err()
            .contains("aggregate entries"));

        let path = state_path(&root);
        fs::write(&path, vec![b' '; MAX_FILE_BYTES + 1]).unwrap();
        assert!(load_from(&root)
            .unwrap_err()
            .contains("file exceeds the size limit"));
        assert_eq!(
            fs::metadata(&path).unwrap().len(),
            (MAX_FILE_BYTES + 1) as u64
        );
        assert!(temporary_files(&root).is_empty());
        fs::remove_dir_all(root).unwrap();
    }
}