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, } #[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, } #[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, #[serde(default)] pub runtime_order_by_directory: BTreeMap>, #[serde(default, skip_serializing_if = "Option::is_none")] pub selected_runtime_id: Option, #[serde(default)] pub drafts: BTreeMap, #[serde(default)] pub scroll: BTreeMap, #[serde(default)] pub last_seen: BTreeMap, #[serde(default, skip_serializing_if = "Option::is_none")] pub workspace_cursor: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub interface_scale: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub default_new_session_worktree: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub dismissed_collision_warning: Option, #[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 { 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 { 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 { 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 { 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(); } }