codecraft/server/src/app.rs

//! Shared application state.

use crate::config::{self, Config};
use crate::mcp::McpManager;
use crate::provider::Registry;
use crate::store::Store;
use crate::types::{Event, Question, Session, now_ms};
use anyhow::{Result, anyhow};
use serde_json::{Map, Value};
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, RwLock};
use tokio::sync::{OwnedSemaphorePermit, Semaphore, broadcast, oneshot};
use tokio_util::sync::CancellationToken;

#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub enum PermissionReply {
    /// Allow this one call.
    Once,
    /// Allow this tool (or bash command prefix) for the rest of the session.
    Always,
    Reject,
}

struct PendingPermission {
    event: Event,
    tx: oneshot::Sender<(PermissionReply, Option<String>, Option<Vec<usize>>)>,
}

struct PendingQuestion {
    event: Event,
    tx: oneshot::Sender<Option<Map<String, Value>>>,
}

/// Per-session runtime state (not persisted).
#[derive(Default)]
pub struct Runtime {
    pub cancel: Option<CancellationToken>,
    pub read_files: HashSet<PathBuf>,
    pub approved: HashSet<String>,
    /// The model request currently running (for clients that connect mid-step).
    pub step: Option<Value>,
    /// Reasoning effort chosen for the current turn (overrides the model's default).
    pub effort: Option<String>,
}

pub struct App {
    /// Serializes read-modify-write of sessions (see `update_session`).
    session_write: Mutex<()>,
    config_path: PathBuf,
    config: RwLock<Arc<Config>>,
    registry: RwLock<Arc<Registry>>,
    pub store: Store,
    events: broadcast::Sender<Event>,
    pub mcp: Arc<McpManager>,
    runtime: Mutex<HashMap<String, Runtime>>,
    permissions: Mutex<HashMap<String, PendingPermission>>,
    questions: Mutex<HashMap<String, PendingQuestion>>,
    /// Subagent concurrency limits: key → (size, semaphore). Recreated when the size changes.
    subagent_limits: Mutex<HashMap<String, (usize, Arc<Semaphore>)>>,
    /// One-time code for creating the first admin (None once users exist).
    pub setup_code: Mutex<Option<String>>,
    pub oidc_pending: Mutex<HashMap<String, crate::auth::OidcPending>>,
    pub oidc_discovery: Mutex<Option<(String, Value)>>,
    /// session id → (root session id, owner) – both never change.
    owners: Mutex<HashMap<String, (String, Option<String>)>>,
    /// root session id → users it is shared with.
    shares: Mutex<HashMap<String, Vec<String>>>,
    /// Failed sign-ins per key (ip or username): timestamps in ms.
    pub login_failures: Mutex<HashMap<String, Vec<i64>>>,
    /// State of an admin-triggered tool installation (Ghidra).
    pub tool_install: crate::reverse::ToolInstall,
    /// Background shells of all sessions.
    pub shells: crate::shells::Shells,
    /// Language servers (diagnostics after edits).
    pub lsp: crate::lsp::Lsp,
    /// Web Push (VAPID key, window visibility per user).
    pub push: crate::push::Push,
    /// Web terminals (PTYs).
    pub terminals: crate::terminal::Terminals,
    /// Newer release found by the update check (admins see it).
    pub update: crate::update::UpdateState,
    /// Provider reachability (checked every 5 minutes).
    pub health: crate::health::Health,
    /// Automatic backups (last result, run lock).
    pub backup: crate::backup::BackupState,
}

/// What a user may do with a session.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Access {
    None,
    Read,
    Full,
}

impl App {
    pub async fn new(config_path: PathBuf, cfg: Config) -> Result<Arc<Self>> {
        let data = cfg.data_dir();
        std::fs::create_dir_all(&data)?;
        let store = Store::open(&data.join("codecraft.db"))?;
        let registry = Registry::build(&cfg).await;
        let (events, _) = broadcast::channel(4096);
        let app = Arc::new(Self {
            session_write: Mutex::new(()),
            config_path,
            config: RwLock::new(Arc::new(cfg)),
            registry: RwLock::new(Arc::new(registry)),
            store,
            events,
            mcp: Arc::new(McpManager::default()),
            runtime: Mutex::default(),
            permissions: Mutex::default(),
            questions: Mutex::default(),
            subagent_limits: Mutex::default(),
            setup_code: Mutex::default(),
            oidc_pending: Mutex::default(),
            oidc_discovery: Mutex::default(),
            owners: Mutex::default(),
            shares: Mutex::default(),
            login_failures: Mutex::default(),
            tool_install: Mutex::default(),
            shells: Mutex::default(),
            lsp: Default::default(),
            push: crate::push::Push::new(&data),
            terminals: Default::default(),
            health: Default::default(),
            backup: Default::default(),
            update: Default::default(),
        });
        crate::auth::bootstrap(&app)?;
        app.clone().start_mcp();
        app.clone().start_image_migration();
        Ok(app)
    }

    /// Move inline images of older chats to the image store, drop unused images and
    /// shrink the database file (in the background; chats load meanwhile as before).
    fn start_image_migration(self: Arc<Self>) {
        tokio::task::spawn_blocking(move || {
            match self.store.migrate_inline_images() {
                Ok(0) => {}
                Ok(n) => tracing::info!("moved {n} images of older chats to the image store"),
                Err(e) => tracing::warn!("moving images to the image store: {e:#}"),
            }
            match self.store.gc_blobs() {
                Ok(0) => {}
                Ok(n) => tracing::info!("removed {n} unused images"),
                Err(e) => tracing::warn!("image cleanup: {e:#}"),
            }
            match self.store.vacuum_if_worthwhile() {
                Ok(true) => tracing::info!("database compacted"),
                Ok(false) => {}
                Err(e) => tracing::warn!("database compaction: {e:#}"),
            }
        });
    }

    fn start_mcp(self: Arc<Self>) {
        let servers = self.config().mcp.clone();
        let mcp = self.mcp.clone();
        let weak = Arc::downgrade(&self);
        tokio::spawn(async move { mcp.start_all(weak, servers).await });
    }

    pub fn config(&self) -> Arc<Config> {
        self.config.read().unwrap().clone()
    }

    pub fn registry(&self) -> Arc<Registry> {
        self.registry.read().unwrap().clone()
    }

    pub fn workspace(&self) -> PathBuf {
        let ws = self.config().workspace();
        std::fs::canonicalize(&ws).unwrap_or(ws)
    }

    /// Directories a session owner may work in: admins (and legacy sessions
    /// without owner) the whole workspace, users their own folder plus shared folders.
    pub fn user_roots(&self, owner_id: Option<&str>) -> Vec<PathBuf> {
        let ws = self.workspace();
        let Some(user) = owner_id.and_then(|id| self.store.user(id).ok().flatten()) else {
            return vec![ws];
        };
        if user.is_admin() {
            return vec![ws];
        }
        let own = ws.join(&user.username);
        let _ = std::fs::create_dir_all(&own);
        let mut roots = vec![own];
        for p in self.store.folders_for_user(&user.id).unwrap_or_default() {
            let p = PathBuf::from(p);
            if p.starts_with(&ws) && p.is_dir() {
                roots.push(p);
            }
        }
        roots
    }

    /// The root that contains `cwd` (most specific first).
    pub fn root_for(&self, owner_id: Option<&str>, cwd: &Path) -> PathBuf {
        let mut roots = self.user_roots(owner_id);
        roots.sort_by_key(|r| std::cmp::Reverse(r.as_os_str().len()));
        roots.iter().find(|r| cwd.starts_with(r)).cloned().unwrap_or_else(|| roots.pop().unwrap_or_else(|| self.workspace()))
    }

    fn owner_of(&self, session_id: &str) -> Option<(String, Option<String>)> {
        if let Some(v) = self.owners.lock().unwrap().get(session_id) {
            return Some(v.clone());
        }
        let root = self.root_session_id(session_id);
        let owner = self.store.get_session(&root).ok().flatten()?.owner_id;
        let v = (root, owner);
        self.owners.lock().unwrap().insert(session_id.to_string(), v.clone());
        Some(v)
    }

    fn shared_with(&self, root: &str) -> Vec<String> {
        if let Some(v) = self.shares.lock().unwrap().get(root) {
            return v.clone();
        }
        let v = self.store.shares(root).unwrap_or_default();
        self.shares.lock().unwrap().insert(root.to_string(), v.clone());
        v
    }

    pub fn set_shares(&self, root: &str, users: &[String]) -> Result<()> {
        self.store.set_shares(root, users)?;
        self.shares.lock().unwrap().insert(root.to_string(), users.to_vec());
        Ok(())
    }

    pub fn access(&self, user: &crate::store::User, session_id: &str) -> Access {
        let Some((root, owner)) = self.owner_of(session_id) else { return Access::None };
        if user.is_admin() || owner.as_deref() == Some(user.id.as_str()) {
            Access::Full
        } else if self.shared_with(&root).contains(&user.id) {
            Access::Read
        } else {
            Access::None
        }
    }

    /// Whether an event may be sent to this user's connections.
    pub fn event_visible(&self, user: &crate::store::User, ev: &Event) -> bool {
        match ev.session_id() {
            None => user.is_admin(),
            Some(_) if user.is_admin() => true,
            Some(id) => self.access(user, id) != Access::None,
        }
    }

    /// Record an action in the audit log.
    pub fn audit(&self, user: Option<&crate::store::User>, action: &str, detail: impl AsRef<str>, ip: &str) {
        self.store.audit(user.map(|u| u.id.as_str()), user.map(|u| u.username.as_str()).unwrap_or(""), action, detail.as_ref(), ip);
    }

    pub fn data_dir(&self) -> PathBuf {
        self.config().data_dir()
    }

    pub fn config_text(&self) -> Result<String> {
        Ok(std::fs::read_to_string(&self.config_path)?)
    }

    /// Validate, write and apply a new config text. Providers and MCP servers reload.
    pub async fn save_config(self: &Arc<Self>, text: &str) -> Result<()> {
        let cfg = Config::parse(text)?;
        let tmp = self.config_path.with_extension("toml.tmp");
        std::fs::write(&tmp, text)?;
        std::fs::rename(&tmp, &self.config_path)?;
        let registry = Registry::build(&cfg).await;
        let mcp_changed = serde_json::to_value(&cfg.mcp)? != serde_json::to_value(&self.config().mcp)?;
        let providers_changed = serde_json::to_value(&cfg.providers)? != serde_json::to_value(&self.config().providers)?;
        *self.config.write().unwrap() = Arc::new(cfg);
        *self.registry.write().unwrap() = Arc::new(registry);
        if mcp_changed {
            self.clone().start_mcp();
        }
        // A new key or URL should show its status right away, not at the next 5-minute check.
        if providers_changed {
            crate::health::check_all(self).await;
        }
        Ok(())
    }

    pub async fn patch_config(self: &Arc<Self>, patches: &[config::ConfigPatch]) -> Result<()> {
        let text = config::apply_patches(&self.config_text()?, patches)?;
        self.save_config(&text).await
    }

    pub async fn reload_providers(&self) {
        let registry = Registry::build(&self.config()).await;
        *self.registry.write().unwrap() = Arc::new(registry);
        crate::health::check_all(self).await;
    }

    // --- events ---------------------------------------------------------

    pub fn emit(&self, ev: Event) {
        match &ev {
            Event::StepStarted { session_id, .. } => {
                let v = serde_json::to_value(&ev).ok();
                self.with_runtime(session_id, |r| r.step = v);
            }
            Event::SessionStatus { session_id, busy: false } => self.with_runtime(session_id, |r| r.step = None),
            _ => {}
        }
        let _ = self.events.send(ev);
    }

    pub fn subscribe(&self) -> broadcast::Receiver<Event> {
        self.events.subscribe()
    }

    /// Pending permission and question requests, re-sent to newly connected clients.
    pub fn pending_events(&self) -> Vec<Event> {
        let mut out: Vec<Event> = self.permissions.lock().unwrap().values().map(|p| p.event.clone()).collect();
        out.extend(self.questions.lock().unwrap().values().map(|q| q.event.clone()));
        out
    }

    // --- runtime --------------------------------------------------------

    pub fn with_runtime<R>(&self, session_id: &str, f: impl FnOnce(&mut Runtime) -> R) -> R {
        let mut rt = self.runtime.lock().unwrap();
        f(rt.entry(session_id.to_string()).or_default())
    }

    pub fn is_busy(&self, session_id: &str) -> bool {
        self.with_runtime(session_id, |r| r.cancel.is_some())
    }

    pub fn mark_read(&self, session_id: &str, path: &Path) {
        self.with_runtime(session_id, |r| r.read_files.insert(path.to_path_buf()));
    }

    pub fn was_read(&self, session_id: &str, path: &Path) -> bool {
        self.with_runtime(session_id, |r| r.read_files.contains(path))
    }

    pub fn session(&self, id: &str) -> Result<Session> {
        self.store.get_session(id)?.ok_or_else(|| anyhow!("session not found"))
    }

    pub fn save_session(&self, mut s: Session) -> Result<Session> {
        s.updated_at = now_ms();
        self.store.save_session(&s)?;
        self.emit(Event::SessionUpdated { session: s.clone() });
        Ok(s)
    }

    /// Load, change and save a session atomically: background tasks (title, usage, phase,
    /// todos) change different fields at the same time and must not undo each other.
    pub fn update_session(&self, id: &str, f: impl FnOnce(&mut Session)) -> Result<Session> {
        self.modify_session(id, true, f)
    }

    /// Like `update_session`; `touch = false` keeps `updated_at` (organising the list).
    pub fn modify_session(&self, id: &str, touch: bool, f: impl FnOnce(&mut Session)) -> Result<Session> {
        let _guard = self.session_write.lock().unwrap();
        let mut s = self.session(id)?;
        f(&mut s);
        if touch {
            return self.save_session(s);
        }
        self.store.save_session(&s)?;
        self.emit(Event::SessionUpdated { session: s.clone() });
        Ok(s)
    }

    pub fn set_todos(&self, session_id: &str, todos: Value) -> Result<()> {
        self.update_session(session_id, |s| s.todos = todos)?;
        Ok(())
    }

    /// Walk up the parent chain to the top-level session.
    pub fn root_session_id(&self, session_id: &str) -> String {
        let mut id = session_id.to_string();
        for _ in 0..16 {
            match self.store.get_session(&id).ok().flatten().and_then(|s| s.parent_id) {
                Some(p) => id = p,
                None => break,
            }
        }
        id
    }

    fn limit(&self, key: &str, size: usize) -> Arc<Semaphore> {
        let mut m = self.subagent_limits.lock().unwrap();
        match m.get(key) {
            Some((n, sem)) if *n == size => sem.clone(),
            _ => {
                let sem = Arc::new(Semaphore::new(size));
                m.insert(key.to_string(), (size, sem.clone()));
                sem
            }
        }
    }

    /// Wait until a subagent of this type may run (per-agent limit, then the global
    /// limit – always in this order). Returns None if cancelled while waiting, and
    /// whether the call had to wait.
    pub async fn subagent_slot(
        &self,
        agent: &str,
        max_parallel: Option<usize>,
        session_id: &str,
        cancel: &CancellationToken,
    ) -> (Option<Vec<OwnedSemaphorePermit>>, bool) {
        let mut sems = Vec::new();
        if let Some(n) = max_parallel.filter(|n| *n > 0) {
            sems.push((format!("subagent:{agent}"), self.limit(&format!("agent:{agent}"), n)));
        }
        let global = self.config().server.max_parallel_subagents;
        if global > 0 {
            sems.push(("subagent:*".to_string(), self.limit("*", global)));
        }
        let mut waited = false;
        let mut permits = Vec::new();
        for (reason, sem) in sems {
            if sem.available_permits() == 0 {
                waited = true;
                self.emit(Event::Waiting { session_id: session_id.to_string(), reason: Some(reason) });
            }
            tokio::select! {
                p = sem.acquire_owned() => match p {
                    Ok(p) => permits.push(p),
                    Err(_) => return (None, waited),
                },
                _ = cancel.cancelled() => return (None, waited),
            }
        }
        (Some(permits), waited)
    }

    // --- user interaction -------------------------------------------------

    /// Forward questions to the UI and wait. None = dismissed.
    pub async fn ask_user(
        &self,
        session_id: &str,
        source: &str,
        questions: Vec<Question>,
        cancel: CancellationToken,
    ) -> Result<Option<Map<String, Value>>> {
        let id = crate::types::new_id();
        let event = Event::QuestionRequest {
            id: id.clone(),
            session_id: session_id.to_string(),
            root_session_id: self.root_session_id(session_id),
            source: source.to_string(),
            questions,
        };
        let (tx, rx) = oneshot::channel();
        self.questions.lock().unwrap().insert(id.clone(), PendingQuestion { event: event.clone(), tx });
        self.emit(event);
        let res = tokio::select! {
            r = rx => r.unwrap_or(None),
            _ = cancel.cancelled() => None,
        };
        self.questions.lock().unwrap().remove(&id);
        self.emit(Event::QuestionResolved { id, session_id: session_id.to_string() });
        Ok(res)
    }

    /// "tool: title" of a pending permission request (for the audit log).
    pub fn pending_title(&self, id: &str) -> Option<String> {
        self.permissions.lock().unwrap().get(id).and_then(|p| match &p.event {
            Event::PermissionRequest { tool, title, .. } => Some(format!("{tool}: {title}")),
            _ => None,
        })
    }

    /// Session of a pending permission request or question.
    pub fn pending_session(&self, id: &str) -> Option<String> {
        let from_perm = self.permissions.lock().unwrap().get(id).and_then(|p| p.event.session_id().map(String::from));
        from_perm.or_else(|| self.questions.lock().unwrap().get(id).and_then(|q| q.event.session_id().map(String::from)))
    }

    pub fn answer_question(&self, id: &str, answers: Option<Map<String, Value>>) -> Result<()> {
        let p = self.questions.lock().unwrap().remove(id).ok_or_else(|| anyhow!("question not pending"))?;
        let _ = p.tx.send(answers);
        Ok(())
    }

    /// Ask the user to approve a tool call. Returns the reply and an optional
    /// message (feedback on rejection).
    pub async fn request_permission(
        &self,
        session_id: &str,
        tool: &str,
        title: String,
        input: Value,
        diff: Option<String>,
        cancel: CancellationToken,
    ) -> (PermissionReply, Option<String>, Option<Vec<usize>>) {
        let id = crate::types::new_id();
        let event = Event::PermissionRequest {
            id: id.clone(),
            session_id: session_id.to_string(),
            root_session_id: self.root_session_id(session_id),
            tool: tool.to_string(),
            title,
            input,
            diff,
        };
        let (tx, rx) = oneshot::channel();
        self.permissions.lock().unwrap().insert(id.clone(), PendingPermission { event: event.clone(), tx });
        self.emit(event);
        let res = tokio::select! {
            r = rx => r.unwrap_or((PermissionReply::Reject, None, None)),
            _ = cancel.cancelled() => (PermissionReply::Reject, Some("aborted".into()), None),
        };
        self.permissions.lock().unwrap().remove(&id);
        self.emit(Event::PermissionResolved { id, session_id: session_id.to_string() });
        res
    }

    /// `hunks`: for file edits, the indices of the diff hunks the user accepted (None = all).
    pub fn reply_permission(&self, id: &str, reply: PermissionReply, message: Option<String>, hunks: Option<Vec<usize>>) -> Result<()> {
        let p = self.permissions.lock().unwrap().remove(id).ok_or_else(|| anyhow!("permission not pending"))?;
        let _ = p.tx.send((reply, message, hunks));
        Ok(())
    }
}