Genarrative/apps/ai-game-creator-shell/src-tauri/src/runner/state.rs

use super::{endpoint::*, project_owner::*, protocol::*};
use std::collections::{BTreeMap, BTreeSet};
use std::fs::{self, File};
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Condvar, Mutex, MutexGuard};

pub(super) struct ExternalAgentRunnerServerState {
    pub(super) endpoint_path: PathBuf,
    pub(super) endpoint: Mutex<ExternalAgentRunnerEndpoint>,
    pub(super) shutdown_requested: AtomicBool,
    pub(super) force_shutdown_requested: AtomicBool,
    pub(super) gui_owner_attached: AtomicBool,
    pub(super) gui_owner_platform_session_claim: Mutex<Option<(String, u64)>>,
    pub(super) draining: AtomicBool,
    pub(super) active_connections: AtomicUsize,
    pub(super) known_roots: Mutex<BTreeSet<PathBuf>>,
    pub(super) gui_owner_lock_path: PathBuf,
    pub(super) project_execution_owners:
        Mutex<BTreeMap<PathBuf, ExternalAgentRunnerProjectExecutionOwnerEntry>>,
    project_execution_owner_recovery_changed: Condvar,
    pub(super) write_request_cache: Mutex<ExternalAgentRunnerRequestCache>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub(super) struct ExternalAgentRunnerProjectExecutionOwnerClaim {
    pub(super) root: PathBuf,
    pub(super) recovered_from_boot_id: Option<String>,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus {
    Pending,
    Running,
    Completed,
}

pub(super) struct ExternalAgentRunnerProjectExecutionOwnerEntry {
    _owner: ExternalAgentRunnerProjectExecutionOwner,
    recovered_from_boot_id: Option<String>,
    recovery_status: ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus,
}

struct ExternalAgentRunnerProjectExecutionOwnerRecoveryGuard<'a> {
    state: &'a ExternalAgentRunnerServerState,
    root: PathBuf,
    running: bool,
}

impl ExternalAgentRunnerProjectExecutionOwnerRecoveryGuard<'_> {
    fn complete(mut self) {
        self.update_status(ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus::Completed);
        self.running = false;
    }

    fn update_status(&self, status: ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus) {
        let mut owners = lock_unpoisoned(&self.state.project_execution_owners);
        if let Some(entry) = owners.get_mut(&self.root) {
            if entry.recovery_status
                == ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus::Running
            {
                entry.recovery_status = status;
            }
        }
        drop(owners);
        self.state
            .project_execution_owner_recovery_changed
            .notify_all();
    }
}

impl Drop for ExternalAgentRunnerProjectExecutionOwnerRecoveryGuard<'_> {
    fn drop(&mut self) {
        if self.running {
            self.update_status(ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus::Pending);
        }
    }
}

impl ExternalAgentRunnerServerState {
    pub(super) fn new(endpoint_path: PathBuf, endpoint: ExternalAgentRunnerEndpoint) -> Self {
        let gui_owner_lock_path = endpoint_path
            .parent()
            .map(external_agent_runner_gui_owner_lock_path)
            .unwrap_or_else(|| PathBuf::from(EXTERNAL_AGENT_RUNNER_GUI_OWNER_LOCK_FILE_NAME));
        Self {
            endpoint_path,
            endpoint: Mutex::new(endpoint),
            shutdown_requested: AtomicBool::new(false),
            force_shutdown_requested: AtomicBool::new(false),
            gui_owner_attached: AtomicBool::new(false),
            gui_owner_platform_session_claim: Mutex::new(None),
            draining: AtomicBool::new(false),
            active_connections: AtomicUsize::new(0),
            known_roots: Mutex::new(BTreeSet::new()),
            gui_owner_lock_path,
            project_execution_owners: Mutex::new(BTreeMap::new()),
            project_execution_owner_recovery_changed: Condvar::new(),
            write_request_cache: Mutex::new(ExternalAgentRunnerRequestCache::default()),
        }
    }

    pub(super) fn endpoint_snapshot(&self) -> ExternalAgentRunnerEndpoint {
        lock_unpoisoned(&self.endpoint).clone()
    }

    pub(super) fn public_status(&self) -> ExternalAgentRunnerStatus {
        ExternalAgentRunnerStatus::from_endpoint(&self.endpoint_snapshot(), true)
    }

    pub(super) fn remember_root(&self, root: &Path) {
        lock_unpoisoned(&self.known_roots).insert(root.to_path_buf());
    }

    pub(super) fn known_roots_snapshot(&self) -> Vec<PathBuf> {
        lock_unpoisoned(&self.known_roots).iter().cloned().collect()
    }

    pub(super) fn claim_project_execution_owner(
        &self,
        root: &Path,
    ) -> Result<ExternalAgentRunnerProjectExecutionOwnerClaim, String> {
        self.claim_project_execution_owner_with_recovery(root, |root| {
            crate::resume_game_creator_agent_background_tasks_at(root).map(|_| ())
        })
    }

    pub(super) fn claim_project_execution_owner_with_recovery(
        &self,
        root: &Path,
        recover: impl Fn(&Path) -> Result<(), String>,
    ) -> Result<ExternalAgentRunnerProjectExecutionOwnerClaim, String> {
        let root = canonicalize_external_agent_runner_project_root(root)?;
        let config_dir = self
            .endpoint_path
            .parent()
            .ok_or_else(|| "Agent Runner endpoint 缺少 AppData 父目录".to_string())?;
        crate::validate_game_creator_runtime_config_dir_outside_project(config_dir, &root)?;
        let endpoint = self.endpoint_snapshot();

        loop {
            let mut owners = lock_unpoisoned(&self.project_execution_owners);
            if !owners.contains_key(&root) {
                let owner = acquire_external_agent_runner_project_execution_owner(
                    &root,
                    &endpoint.boot_id,
                    endpoint.protocol_version,
                )?;
                let recovered_from_boot_id = owner.recovered_from_boot_id().map(str::to_string);
                let recovery_status = if recovered_from_boot_id.is_some() {
                    ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus::Pending
                } else {
                    ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus::Completed
                };
                owners.insert(
                    root.clone(),
                    ExternalAgentRunnerProjectExecutionOwnerEntry {
                        _owner: owner,
                        recovered_from_boot_id,
                        recovery_status,
                    },
                );
            }

            let entry = owners
                .get_mut(&root)
                .expect("project execution owner entry must exist after acquisition");
            match entry.recovery_status {
                ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus::Completed => {
                    drop(owners);
                    self.remember_root(&root);
                    return Ok(ExternalAgentRunnerProjectExecutionOwnerClaim {
                        root,
                        recovered_from_boot_id: None,
                    });
                }
                ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus::Running => {
                    owners =
                        wait_unpoisoned(&self.project_execution_owner_recovery_changed, owners);
                    drop(owners);
                }
                ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus::Pending => {
                    entry.recovery_status =
                        ExternalAgentRunnerProjectExecutionOwnerRecoveryStatus::Running;
                    let recovered_from_boot_id = entry
                        .recovered_from_boot_id
                        .clone()
                        .expect("pending project owner recovery must name the previous boot");
                    drop(owners);
                    self.remember_root(&root);

                    let recovery_guard = ExternalAgentRunnerProjectExecutionOwnerRecoveryGuard {
                        state: self,
                        root: root.clone(),
                        running: true,
                    };
                    match catch_unwind(AssertUnwindSafe(|| recover(&root))) {
                        Ok(Ok(())) => {
                            recovery_guard.complete();
                            return Ok(ExternalAgentRunnerProjectExecutionOwnerClaim {
                                root,
                                recovered_from_boot_id: Some(recovered_from_boot_id),
                            });
                        }
                        Ok(Err(error)) => return Err(error),
                        Err(_) => {
                            return Err(
                                "Agent Runner 项目恢复回调异常终止，可复用同一 requestId 重试"
                                    .to_string(),
                            )
                        }
                    }
                }
            }
        }
    }
}

pub(super) struct ExternalAgentRunnerActiveConnection<'a> {
    pub(super) state: &'a ExternalAgentRunnerServerState,
}

impl Drop for ExternalAgentRunnerActiveConnection<'_> {
    fn drop(&mut self) {
        self.state.active_connections.fetch_sub(1, Ordering::AcqRel);
    }
}

pub(super) struct ExternalAgentRunnerInstanceLock {
    pub(super) _file: File,
}

#[derive(Debug)]
pub(crate) struct ExternalAgentRunnerGuiOwnerLock {
    pub(super) _file: File,
    pub(super) owner_epoch: String,
}

impl ExternalAgentRunnerGuiOwnerLock {
    pub(crate) fn owner_epoch(&self) -> &str {
        &self.owner_epoch
    }
}

pub(super) struct ExternalAgentRunnerProjectOwnerStorage {
    pub(super) lock_file: File,
    pub(super) directory_handles: Vec<File>,
    pub(super) lock_path: PathBuf,
    pub(super) diagnostic_path: PathBuf,
}

impl ExternalAgentRunnerProjectOwnerStorage {
    pub(super) fn runtime_directory(&self) -> &File {
        self.directory_handles
            .last()
            .expect("project owner storage always holds the Runtime directory")
    }
}

pub(super) struct ExternalAgentRunnerProjectExecutionOwner {
    pub(super) _file: File,
    pub(super) _directory_handles: Vec<File>,
    pub(super) _record: ExternalAgentRunnerProjectExecutionOwnerRecord,
}

impl ExternalAgentRunnerProjectExecutionOwner {
    pub(super) fn recovered_from_boot_id(&self) -> Option<&str> {
        self._record.recovered_from_boot_id.as_deref()
    }
}

pub(super) struct ExternalAgentRunnerEndpointGuard {
    pub(super) path: PathBuf,
    pub(super) boot_id: String,
}

pub(super) fn remove_external_agent_runner_endpoint_if_boot_matches(path: &Path, boot_id: &str) {
    let Ok(endpoint) = read_external_agent_runner_endpoint(path) else {
        return;
    };
    if endpoint.boot_id == boot_id {
        let _ = fs::remove_file(path);
    }
}

impl Drop for ExternalAgentRunnerEndpointGuard {
    fn drop(&mut self) {
        remove_external_agent_runner_endpoint_if_boot_matches(&self.path, &self.boot_id);
    }
}

pub(super) struct ExternalAgentRunnerTempFileGuard {
    pub(super) path: PathBuf,
    pub(super) installed: bool,
}

impl Drop for ExternalAgentRunnerTempFileGuard {
    fn drop(&mut self) {
        if !self.installed {
            let _ = fs::remove_file(&self.path);
        }
    }
}

pub(super) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
    mutex
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
}

fn wait_unpoisoned<'a, T>(condvar: &Condvar, guard: MutexGuard<'a, T>) -> MutexGuard<'a, T> {
    condvar
        .wait(guard)
        .unwrap_or_else(|poisoned| poisoned.into_inner())
}