Genarrative/apps/ai-game-creator-shell/src-tauri/src/project/checkpoint.rs

use super::*;

use super::filesystem::validate_portable_project_path_component;

#[cfg(windows)]
use super::filesystem::PROJECT_FILE_FLAG_OPEN_REPARSE_POINT;

pub(crate) fn create_local_project_checkpoint_at(
    root: &Path,
) -> Result<LocalProjectCheckpointResult, String> {
    validate_project_root(root)?;
    let checkpoint_id = format!("checkpoint-{}", unix_millis());
    let checkpoint_root =
        resolve_local_project_path(root, &checkpoint_root_relative_path(&checkpoint_id))?;
    let files = collect_project_index_files(root)?;
    for file in &files {
        let normalized_path =
            validate_checkpoint_manifest_file_path(root, &checkpoint_id, &file.path)?;
        let source = resolve_local_project_path(root, &normalized_path)?;
        let target = resolve_local_project_path(
            root,
            &checkpoint_file_relative_path(&checkpoint_id, &normalized_path),
        )?;
        if let Some(parent) = target.parent() {
            fs::create_dir_all(parent).map_err(|error| {
                format!("创建 checkpoint 目录失败：{}: {error}", parent.display())
            })?;
        }
        fs::copy(&source, &target).map_err(|error| {
            format!(
                "写入 checkpoint 文件失败：{} -> {}: {error}",
                source.display(),
                target.display()
            )
        })?;
    }
    let total_bytes = files.iter().map(|file| file.size).sum::<u64>();
    let manifest = serde_json::json!({
        "checkpointId": checkpoint_id,
        "createdAt": unix_timestamp(),
        "files": files,
    });
    let manifest_path =
        resolve_local_project_path(root, &checkpoint_manifest_relative_path(&checkpoint_id))?;
    fs::write(
        &manifest_path,
        format!(
            "{}\n",
            serde_json::to_string_pretty(&manifest)
                .map_err(|error| format!("序列化 checkpoint 失败：{error}"))?
        ),
    )
    .map_err(|error| {
        format!(
            "写入 checkpoint manifest 失败：{}: {error}",
            manifest_path.display()
        )
    })?;
    append_agent_db_record(
        root,
        serde_json::json!({
            "recordType": "project.checkpoint",
            "checkpointId": checkpoint_id,
            "fileCount": files.len(),
            "totalBytes": total_bytes,
        }),
    )?;
    Ok(LocalProjectCheckpointResult {
        checkpoint_id,
        checkpoint_path: checkpoint_root.to_string_lossy().into_owned(),
        file_count: files.len(),
        total_bytes,
    })
}

const PROJECT_CONTENT_DIFF_CONTEXT_LINES: usize = 3;
const PROJECT_CONTENT_DIFF_MAX_FILE_BYTES: u64 = 2 * 1024 * 1024;

#[derive(Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct LocalProjectContentDiffResult {
    pub(crate) checkpoint_id: String,
    pub(crate) content: String,
    pub(crate) file_count: usize,
    pub(crate) truncated: bool,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum LocalProjectContentDiffStatus {
    Added,
    Changed,
    Deleted,
}

impl LocalProjectContentDiffStatus {
    fn as_str(self) -> &'static str {
        match self {
            Self::Added => "added",
            Self::Changed => "changed",
            Self::Deleted => "deleted",
        }
    }
}

struct LocalProjectContentDiffFile {
    path: String,
    status: LocalProjectContentDiffStatus,
}

struct LocalProjectContentDiffSource {
    bytes: Option<Vec<u8>>,
    size: u64,
    sha256: String,
}

pub(crate) fn open_project_snapshot_regular_file(
    path: &Path,
    label: &str,
) -> Result<(File, fs::Metadata), String> {
    let mut options = fs::OpenOptions::new();
    options.read(true);
    #[cfg(unix)]
    {
        use std::os::unix::fs::OpenOptionsExt;
        options.custom_flags(libc::O_NOFOLLOW);
    }
    #[cfg(windows)]
    {
        use std::os::windows::fs::OpenOptionsExt;
        options.custom_flags(PROJECT_FILE_FLAG_OPEN_REPARSE_POINT);
    }
    let file = options
        .open(path)
        .map_err(|error| format!("打开{label}失败：{}: {error}", path.display()))?;
    let metadata = file
        .metadata()
        .map_err(|error| format!("读取{label}元数据失败：{}: {error}", path.display()))?;
    if !metadata.is_file() {
        return Err(format!("{label}必须是普通文件：{}", path.display()));
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::MetadataExt;
        if metadata.nlink() != 1 {
            return Err(format!("{label}不能是硬链接文件：{}", path.display()));
        }
    }
    #[cfg(windows)]
    validate_windows_regular_file_handle(&file, label)?;
    Ok((file, metadata))
}

fn read_local_project_content_diff_source(
    path: &Path,
) -> Result<LocalProjectContentDiffSource, String> {
    let (mut file, metadata) = open_project_snapshot_regular_file(path, "内容 diff 文件")?;
    let mut hasher = Sha256::new();
    let mut bytes = (metadata.len() <= PROJECT_CONTENT_DIFF_MAX_FILE_BYTES)
        .then(|| Vec::with_capacity(metadata.len() as usize));
    let mut size = 0_u64;
    let mut buffer = [0_u8; 64 * 1024];
    loop {
        let read = file
            .read(&mut buffer)
            .map_err(|error| format!("读取内容 diff 文件失败：{}: {error}", path.display()))?;
        if read == 0 {
            break;
        }
        hasher.update(&buffer[..read]);
        size = size.saturating_add(read as u64);
        if let Some(content) = bytes.as_mut() {
            if size <= PROJECT_CONTENT_DIFF_MAX_FILE_BYTES {
                content.extend_from_slice(&buffer[..read]);
            } else {
                bytes = None;
            }
        }
    }
    let final_metadata = file
        .metadata()
        .map_err(|error| format!("复核内容 diff 文件失败：{}: {error}", path.display()))?;
    if final_metadata.len() != size {
        return Err(format!(
            "内容 diff 读取期间文件发生漂移：{}",
            path.display()
        ));
    }
    #[cfg(unix)]
    {
        use std::os::unix::fs::MetadataExt;
        if final_metadata.nlink() != 1 {
            return Err(format!(
                "内容 diff 文件不能是硬链接文件：{}",
                path.display()
            ));
        }
    }
    #[cfg(windows)]
    validate_windows_regular_file_handle(&file, "内容 diff 文件")?;

    Ok(LocalProjectContentDiffSource {
        bytes,
        size,
        sha256: format!("{:x}", hasher.finalize()),
    })
}

fn local_project_content_diff_source_at(
    root: &Path,
    checkpoint_id: &str,
    file: &LocalProjectContentDiffFile,
    checkpoint: bool,
) -> Result<Option<LocalProjectContentDiffSource>, String> {
    if (checkpoint && file.status == LocalProjectContentDiffStatus::Added)
        || (!checkpoint && file.status == LocalProjectContentDiffStatus::Deleted)
    {
        return Ok(None);
    }
    let relative_path = if checkpoint {
        checkpoint_file_relative_path(checkpoint_id, &file.path)
    } else {
        file.path.clone()
    };
    let path = resolve_local_project_path(root, &relative_path)?;
    read_local_project_content_diff_source(&path).map(Some)
}

fn render_local_project_content_diff_section(
    file: &LocalProjectContentDiffFile,
    checkpoint: Option<&LocalProjectContentDiffSource>,
    current: Option<&LocalProjectContentDiffSource>,
) -> (String, bool) {
    let checkpoint_sha256 = checkpoint
        .map(|source| source.sha256.as_str())
        .unwrap_or("-");
    let current_sha256 = current.map(|source| source.sha256.as_str()).unwrap_or("-");
    let mut output = format!(
        "diff --git a/{0} b/{0}\nstatus: {1}\ncheckpoint-sha256: {2}\ncurrent-sha256: {3}\n",
        file.path,
        file.status.as_str(),
        checkpoint_sha256,
        current_sha256,
    );

    if checkpoint.is_some_and(|source| source.bytes.is_none())
        || current.is_some_and(|source| source.bytes.is_none())
    {
        let checkpoint_size = checkpoint
            .map(|source| source.size.to_string())
            .unwrap_or_else(|| "-".to_string());
        let current_size = current
            .map(|source| source.size.to_string())
            .unwrap_or_else(|| "-".to_string());
        output.push_str(&format!(
            "[content omitted: per-file limit exceeded; limit={PROJECT_CONTENT_DIFF_MAX_FILE_BYTES} bytes; checkpoint={checkpoint_size} bytes; current={current_size} bytes; truncated]\n\n"
        ));
        return (output, true);
    }

    let checkpoint_bytes = checkpoint
        .and_then(|source| source.bytes.as_deref())
        .unwrap_or_default();
    let current_bytes = current
        .and_then(|source| source.bytes.as_deref())
        .unwrap_or_default();
    if checkpoint_bytes.contains(&0) || current_bytes.contains(&0) {
        output.push_str("[content omitted: binary data detected; truncated]\n\n");
        return (output, true);
    }
    let (Ok(checkpoint_text), Ok(current_text)) = (
        std::str::from_utf8(checkpoint_bytes),
        std::str::from_utf8(current_bytes),
    ) else {
        output.push_str("[content omitted: non-UTF-8 data detected; truncated]\n\n");
        return (output, true);
    };

    let diff = TextDiff::from_lines(checkpoint_text, current_text);
    let checkpoint_header = if file.status == LocalProjectContentDiffStatus::Added {
        "/dev/null".to_string()
    } else {
        format!("a/{}", file.path)
    };
    let current_header = if file.status == LocalProjectContentDiffStatus::Deleted {
        "/dev/null".to_string()
    } else {
        format!("b/{}", file.path)
    };
    let mut unified = diff.unified_diff();
    unified.context_radius(PROJECT_CONTENT_DIFF_CONTEXT_LINES);
    unified.header(&checkpoint_header, &current_header);
    output.push_str(&unified.to_string());
    if !output.ends_with('\n') {
        output.push('\n');
    }
    output.push('\n');
    (output, false)
}

fn append_local_project_content_diff_truncation(
    output: &mut String,
    output_chars: &mut usize,
    max_chars: usize,
    detail: &str,
) {
    let marker = format!("[content diff truncated: {detail}]\n");
    let marker_chars = marker.chars().count();
    if output_chars.saturating_add(marker_chars) <= max_chars {
        output.push_str(&marker);
        *output_chars += marker_chars;
    } else if *output_chars < max_chars {
        output.push('…');
        *output_chars += 1;
    }
}

pub(crate) fn diff_local_project_checkpoint_content_at(
    root: &Path,
    checkpoint_id: &str,
    max_files: usize,
    max_chars: usize,
) -> Result<LocalProjectContentDiffResult, String> {
    let initial_diff = diff_local_project_checkpoint_at(root, checkpoint_id)?;
    let checkpoint_id = initial_diff.checkpoint_id.clone();
    let mut files = Vec::with_capacity(
        initial_diff.added.len() + initial_diff.changed.len() + initial_diff.deleted.len(),
    );
    files.extend(
        initial_diff
            .added
            .iter()
            .map(|entry| LocalProjectContentDiffFile {
                path: entry.path.clone(),
                status: LocalProjectContentDiffStatus::Added,
            }),
    );
    files.extend(
        initial_diff
            .changed
            .iter()
            .map(|entry| LocalProjectContentDiffFile {
                path: entry.path.clone(),
                status: LocalProjectContentDiffStatus::Changed,
            }),
    );
    files.extend(
        initial_diff
            .deleted
            .iter()
            .map(|entry| LocalProjectContentDiffFile {
                path: entry.path.clone(),
                status: LocalProjectContentDiffStatus::Deleted,
            }),
    );
    files.sort_by(|left, right| left.path.cmp(&right.path));

    let mut content = String::new();
    let mut content_chars = 0_usize;
    let mut file_count = 0_usize;
    let mut truncated = false;
    let mut character_budget_exhausted = false;
    let file_limit = max_files.min(files.len());
    for (index, file) in files.iter().take(file_limit).enumerate() {
        let checkpoint = local_project_content_diff_source_at(root, &checkpoint_id, file, true)?;
        let current = local_project_content_diff_source_at(root, &checkpoint_id, file, false)?;
        let (section, section_truncated) =
            render_local_project_content_diff_section(file, checkpoint.as_ref(), current.as_ref());
        let section_chars = section.chars().count();
        let has_unprocessed_files = index + 1 < files.len();
        let truncation_reserve = usize::from(has_unprocessed_files && max_chars > 0);
        if content_chars
            .saturating_add(section_chars)
            .saturating_add(truncation_reserve)
            > max_chars
        {
            truncated = true;
            character_budget_exhausted = true;
            append_local_project_content_diff_truncation(
                &mut content,
                &mut content_chars,
                max_chars,
                &format!(
                    "character budget reached; {} file diff(s) omitted",
                    files.len() - file_count
                ),
            );
            break;
        }
        content.push_str(&section);
        content_chars += section_chars;
        file_count += 1;
        truncated |= section_truncated;
    }
    if !character_budget_exhausted && file_limit < files.len() {
        truncated = true;
        append_local_project_content_diff_truncation(
            &mut content,
            &mut content_chars,
            max_chars,
            &format!(
                "file budget reached; {} file diff(s) omitted",
                files.len() - file_count
            ),
        );
    }
    let final_diff = diff_local_project_checkpoint_at(root, &checkpoint_id)?;
    if final_diff != initial_diff {
        return Err("内容 diff 读取期间项目文件发生变化，请重新执行 project.diff".to_string());
    }

    Ok(LocalProjectContentDiffResult {
        checkpoint_id,
        content,
        file_count,
        truncated,
    })
}

pub(crate) fn diff_local_project_checkpoint_at(
    root: &Path,
    checkpoint_id: &str,
) -> Result<LocalProjectDiffResult, String> {
    validate_project_root(root)?;
    let checkpoint_id = normalize_checkpoint_id(checkpoint_id)?;
    let checkpoint_files = read_checkpoint_files(root, &checkpoint_id)?;
    let current_files = collect_project_index_files(root)?;
    let mut added = Vec::new();
    let mut changed = Vec::new();
    let mut deleted = Vec::new();
    for file in &current_files {
        match checkpoint_files
            .iter()
            .find(|candidate| candidate.path == file.path)
        {
            None => added.push(LocalProjectDiffEntry {
                path: file.path.clone(),
                status: "added".to_string(),
            }),
            Some(previous) if previous.checksum != file.checksum || previous.size != file.size => {
                changed.push(LocalProjectDiffEntry {
                    path: file.path.clone(),
                    status: "changed".to_string(),
                });
            }
            Some(_) => {}
        }
    }
    for file in &checkpoint_files {
        if !current_files
            .iter()
            .any(|candidate| candidate.path == file.path)
        {
            deleted.push(LocalProjectDiffEntry {
                path: file.path.clone(),
                status: "deleted".to_string(),
            });
        }
    }
    Ok(LocalProjectDiffResult {
        checkpoint_id,
        added,
        changed,
        deleted,
    })
}

pub(crate) fn restore_local_project_checkpoint_at(
    root: &Path,
    checkpoint_id: &str,
) -> Result<LocalProjectRestoreResult, String> {
    validate_project_root(root)?;
    let checkpoint_id = normalize_checkpoint_id(checkpoint_id)?;
    let files = read_checkpoint_files(root, &checkpoint_id)?;
    let current_files = collect_project_index_files(root)?;
    let mut restore_plan = Vec::new();
    for file in &files {
        if should_skip_project_restore_path(&file.path) {
            continue;
        }
        let source = resolve_local_project_path(
            root,
            &checkpoint_file_relative_path(&checkpoint_id, &file.path),
        )?;
        let source_metadata = fs::symlink_metadata(&source).map_err(|error| {
            format!(
                "读取 checkpoint 文件类型失败：{}: {error}",
                source.display()
            )
        })?;
        if !source_metadata.is_file() {
            return Err(format!(
                "checkpoint 内容必须是普通文件：{}",
                source.display()
            ));
        }
        let target = resolve_local_project_path(root, &file.path)?;
        if target.exists()
            && !fs::symlink_metadata(&target)
                .map_err(|error| format!("读取恢复目标类型失败：{}: {error}", target.display()))?
                .is_file()
        {
            return Err(format!("checkpoint 只能恢复普通文件：{}", file.path));
        }
        restore_plan.push((source, target));
    }
    let mut delete_plan = Vec::new();
    for file in &current_files {
        if should_skip_project_restore_path(&file.path)
            || files.iter().any(|candidate| candidate.path == file.path)
        {
            continue;
        }
        let target = resolve_local_project_path(root, &file.path)?;
        if target.is_file() {
            delete_plan.push(target);
        }
    }

    let restored_count = restore_plan.len();
    for (source, target) in restore_plan {
        if let Some(parent) = target.parent() {
            fs::create_dir_all(parent)
                .map_err(|error| format!("创建恢复目录失败：{}: {error}", parent.display()))?;
        }
        fs::copy(&source, &target).map_err(|error| {
            format!(
                "恢复 checkpoint 文件失败：{} -> {}: {error}",
                source.display(),
                target.display()
            )
        })?;
    }
    let deleted_count = delete_plan.len();
    for target in delete_plan {
        fs::remove_file(&target).map_err(|error| {
            format!(
                "删除 checkpoint 外新增文件失败：{}: {error}",
                target.display()
            )
        })?;
    }
    let project_index_path = root.join(PROJECT_INDEX_PATH);
    if project_index_path.is_file() {
        fs::remove_file(&project_index_path).map_err(|error| {
            format!(
                "删除恢复后的旧项目索引失败：{}: {error}",
                project_index_path.display()
            )
        })?;
    }
    append_agent_db_record(
        root,
        serde_json::json!({
            "recordType": "project.restore",
            "checkpointId": checkpoint_id,
            "restoredCount": restored_count,
            "deletedCount": deleted_count,
        }),
    )?;
    Ok(LocalProjectRestoreResult {
        checkpoint_id,
        restored_count,
        deleted_count,
    })
}

pub(crate) fn should_skip_project_restore_path(relative_path: &str) -> bool {
    should_skip_project_snapshot_path(relative_path)
        || relative_path == ".agent/agent.db"
        || relative_path == PROJECT_PERMISSION_POLICY_PATH
        || relative_path == PROJECT_WRITE_LOCK_PATH
        || relative_path == PROJECT_INDEX_PATH
        || relative_path.starts_with(".agent/logs/")
        || relative_path.starts_with(".agent/conversations/")
        || relative_path.starts_with(".agent/runtime/")
}

pub(crate) fn normalize_checkpoint_id(checkpoint_id: &str) -> Result<String, String> {
    let normalized = checkpoint_id.trim();
    if normalized != checkpoint_id
        || normalized.is_empty()
        || normalized.chars().any(char::is_control)
        || Path::new(normalized).is_absolute()
        || normalized == "."
        || normalized.contains(':')
        || normalized.contains('/')
        || normalized.contains('\\')
        || normalized.contains("..")
    {
        return Err("checkpoint id 非法".to_string());
    }
    validate_portable_project_path_component(normalized)
        .map_err(|_| "checkpoint id 非法".to_string())?;
    Ok(normalized.to_string())
}

fn checkpoint_root_relative_path(checkpoint_id: &str) -> String {
    format!(".agent/checkpoints/{checkpoint_id}")
}

fn checkpoint_manifest_relative_path(checkpoint_id: &str) -> String {
    format!(
        "{}/manifest.json",
        checkpoint_root_relative_path(checkpoint_id)
    )
}

fn checkpoint_file_relative_path(checkpoint_id: &str, relative_path: &str) -> String {
    format!(
        "{}/files/{relative_path}",
        checkpoint_root_relative_path(checkpoint_id)
    )
}

fn normalize_checkpoint_manifest_file_path(relative_path: &str) -> Result<String, String> {
    if relative_path.chars().any(char::is_control) {
        return Err("checkpoint 文件路径不能包含控制字符".to_string());
    }
    let normalized = normalize_relative_path(relative_path)
        .map_err(|error| format!("checkpoint 文件路径非法：{error}"))?;
    if normalized != relative_path {
        return Err("checkpoint 文件路径必须是规范相对路径".to_string());
    }
    Ok(normalized)
}

fn validate_checkpoint_manifest_file_path(
    root: &Path,
    checkpoint_id: &str,
    relative_path: &str,
) -> Result<String, String> {
    let normalized = normalize_checkpoint_manifest_file_path(relative_path)?;
    resolve_local_project_path(root, &normalized)
        .map_err(|error| format!("checkpoint 文件路径不安全：{normalized}: {error}"))?;
    resolve_local_project_path(
        root,
        &checkpoint_file_relative_path(checkpoint_id, &normalized),
    )
    .map_err(|error| format!("checkpoint 内容路径不安全：{normalized}: {error}"))?;
    Ok(normalized)
}

pub(crate) fn read_checkpoint_files(
    root: &Path,
    checkpoint_id: &str,
) -> Result<Vec<LocalProjectIndexedFile>, String> {
    validate_project_root(root)?;
    let checkpoint_id = normalize_checkpoint_id(checkpoint_id)?;
    let manifest_path =
        resolve_local_project_path(root, &checkpoint_manifest_relative_path(&checkpoint_id))?;
    let manifest_metadata = fs::symlink_metadata(&manifest_path)
        .map_err(|error| format!("读取 checkpoint 失败：{}: {error}", manifest_path.display()))?;
    if !manifest_metadata.is_file() {
        return Err(format!(
            "checkpoint manifest 必须是普通文件：{}",
            manifest_path.display()
        ));
    }
    let content = fs::read_to_string(&manifest_path)
        .map_err(|error| format!("读取 checkpoint 失败：{}: {error}", manifest_path.display()))?;
    let payload = serde_json::from_str::<serde_json::Value>(&content)
        .map_err(|error| format!("解析 checkpoint 失败：{}: {error}", manifest_path.display()))?;
    let mut files = serde_json::from_value::<Vec<LocalProjectIndexedFile>>(
        payload
            .get("files")
            .cloned()
            .ok_or_else(|| "checkpoint 缺少 files".to_string())?,
    )
    .map_err(|error| format!("解析 checkpoint 文件清单失败：{error}"))?;
    let mut validated_paths = BTreeMap::new();
    for file in &mut files {
        let normalized = validate_checkpoint_manifest_file_path(root, &checkpoint_id, &file.path)?;
        let windows_alias_key = normalized.to_ascii_lowercase();
        if validated_paths
            .insert(windows_alias_key, normalized.clone())
            .is_some()
        {
            return Err(format!("checkpoint 文件路径重复：{normalized}"));
        }
        file.path = normalized;
    }
    files.retain(|file| !should_skip_project_snapshot_path(&file.path));
    Ok(files)
}

#[cfg(test)]
mod security_tests;