Tool Permission System — Architecture Report

// In commands/ai.rs — extend CompletionEvent
#[derive(Clone, serde::Serialize)]
#[serde(tag = "event", content = "data")]
pub enum CompletionEvent {
    Chunk { text: String, thinking: Option<String> },
    ToolCall { tool: String, args: serde_json::Value },
    /// NEW: Ask frontend for permission before executing this tool
    ToolPermission { request_id: u64, tool: String, args: serde_json::Value },
    ToolResult { tool: String, success: bool, output: String, path: Option<String>, content: Option<String> },
    Done { done_reason: Option<String> },
    Error { message: String },
}

/// NEW: User's decision for a tool permission request
#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub enum ToolPermissionDecision {
    Accept,
    Reject,
    AlwaysAllow,
}

/// NEW: Permission mode (stored in settings)
#[derive(Debug, Clone, Copy, serde::Serialize, serde::Deserialize, PartialEq, Eq, Default)]
pub enum ToolPermissionMode {
    #[default]
    AskEveryTime,
    AutoAcceptReadOnly,   // Allow read_file silently, gate write_file/bash
    AutoAcceptAll,        // Execute everything (for testing)
}

/// NEW: Per-request pending permission (stored in AppState)
pub struct PendingToolPermission {
    pub sender: tokio::sync::oneshot::Sender<ToolPermissionDecision>,
    pub tool: String,
    pub args: serde_json::Value,
}

pub struct AppState {
    pub active_processes: Mutex<HashMap<u32, CommandChild>>,
    pub cancellation_tokens: Mutex<HashMap<u32, CancellationToken>>,
    pub http_client: reqwest::Client,
    // NEW: Pending permission requests (keyed by auto-incrementing request_id)
    pub pending_permissions: std::sync::Mutex<HashMap<u64, PendingToolPermission>>,
    pub next_permission_id: std::sync::atomic::AtomicU64,
}

/// Resolve a pending tool permission request.
/// Called by frontend when user clicks Accept/Reject/Always Allow.
#[tauri::command]
pub fn resolve_tool_permission(
    permission_id: u64,
    decision: ToolPermissionDecision,
    app: AppHandle,
) -> Result<(), AppError> {
    // Lock held briefly — O(1) operations only
    let state = app.state::<AppState>();
    let mut permissions = state.pending_permissions.lock().unwrap();
    
    let pending = permissions.remove(&permission_id)
        .ok_or_else(|| AppError::NotFound(
            format!("Permission request {permission_id} not found or already resolved")
        ))?;
    
    // Drop lock before sending — oneshot send is sync, no need to hold mutex
    drop(permissions);
    
    // Send decision to the waiting agent loop
    let _ = pending.sender.send(decision);
    
    Ok(())
}

/// Check if a tool should be gated or auto-allowed.
/// Returns: (should_gate, always_allow_this_tool)
fn check_permission_gate(
    tool: &str,
    mode: ToolPermissionMode,
    allowlist: &HashSet<String>,
) -> (bool, bool) {
    match mode {
        ToolPermissionMode::AutoAcceptAll => (false, false),
        ToolPermissionMode::AutoAcceptReadOnly => {
            match tool {
                "read_file" => (false, false),  // Never gate read_file
                _ => {
                    let always_allow = allowlist.contains(tool);
                    (!always_allow, always_allow)
                }
            }
        }
        ToolPermissionMode::AskEveryTime => {
            let always_allow = allowlist.contains(tool);
            (!always_allow, always_allow)
        }
    }
}

/// Request permission from user and block until resolved.
/// Returns the user's decision.
async fn request_permission(
    tool: &str,
    args: &serde_json::Value,
    channel: &Channel<CompletionEvent>,
    cancel_token: &CancellationToken,
    app_handle: &AppHandle,
) -> Result<ToolPermissionDecision, AppError> {
    let state = app_handle.state::<AppState>();
    let request_id = state.next_permission_id.fetch_add(1, Ordering::SeqCst);
    
    let (tx, rx) = tokio::sync::oneshot::channel::<ToolPermissionDecision>();
    
    // Register pending permission
    {
        let mut permissions = state.pending_permissions.lock().unwrap();
        permissions.insert(request_id, PendingToolPermission {
            sender: tx,
            tool: tool.to_string(),
            args: args.clone(),
        });
    }
    
    // Send permission request to frontend
    let _ = channel.send(CompletionEvent::ToolPermission {
        request_id,
        tool: tool.to_string(),
        args: args.clone(),
    });
    
    // Wait for user decision OR cancellation
    tokio::select! {
        decision = rx => {
            // Clean up registration (may already be removed by resolve command)
            let mut permissions = state.pending_permissions.lock().unwrap();
            permissions.remove(&request_id);
            decision.map_err(|_| AppError::Process("Permission channel closed".into()))
        }
        _ = cancel_token.cancelled() => {
            // Clean up on cancellation
            let mut permissions = state.pending_permissions.lock().unwrap();
            if let Some(pending) = permissions.remove(&request_id) {
                let _ = pending.sender.send(ToolPermissionDecision::Reject);
            }
            Ok(ToolPermissionDecision::Reject)
        }
    }
}

// In run_agent_loop, around lines 313-347:
let futures: Vec<_> = (0..tool_calls.len())
    .map(|idx| {
        let name = names[idx].clone();
        let arg = args[idx].clone();
        let proj = proj_dir.clone();
        let wc = Arc::clone(&write_count);
        async move {
            // Atomic write limit check
            let skip = if name == "write_file" {
                wc.load(Ordering::SeqCst) >= MAX_WRITES
            } else { false };
            
            if skip {
                return (idx, ToolExecutionResult {
                    success: false,
                    output: "write_file limit reached".into(),
                    written_path: None,
                    written_content: None,
                });
            }
            
            // --- NEW: Permission gate ---
            let (should_gate, always_allow) = check_permission_gate(
                &name, permission_mode, &allowlist
            );
            
            if should_gate {
                let decision = match request_permission(
                    &name, &arg, channel, cancel_token, app_handle
                ).await {
                    Ok(d) => d,
                    Err(e) => {
                        return (idx, ToolExecutionResult {
                            success: false,
                            output: format!("Permission error: {e}"),
                            written_path: None,
                            written_content: None,
                        });
                    }
                };
                
                match decision {
                    ToolPermissionDecision::Reject => {
                        return (idx, ToolExecutionResult {
                            success: false,
                            output: format!("User rejected {name}"),
                            written_path: None,
                            written_content: None,
                        });
                    }
                    ToolPermissionDecision::AlwaysAllow => {
                        // Persist to allowlist (via tauri-plugin-store)
                        // This is a one-time add, not blocking the execution
                        let _ = persist_allowlist_add(&name).await;
                        // Fall through to execute
                    }
                    ToolPermissionDecision::Accept => {
                        // Fall through to execute
                    }
                }
            }
            
            // Execute tool (existing logic)
            let result = execute_tool(&name, &arg, app_data_dir, output_path, &proj).await;
            if name == "write_file" && result.success {
                wc.fetch_add(1, Ordering::SeqCst);
            }
            (idx, result)
        }
    })
    .collect();

// In request_permission, add timeout branch:
tokio::select! {
    decision = rx => { ... }
    _ = cancel_token.cancelled() => { ... }
    _ = tokio::time::sleep(Duration::from_secs(300)) => {
        // Auto-reject after 5 minutes of inactivity
        permissions.remove(&request_id);
        Ok(ToolPermissionDecision::Reject)
    }
}