agent-workspace-adapter-local 0.3.0

use agent_workspace_contract::{
    SnapshotObjectReader, SnapshotObjectRef, SnapshotObjectStore, SnapshotObjectUpload,
    WORKSPACE_UPLOAD_CHUNK_BYTES,
};
use anyhow::{Context, Result};
use async_trait::async_trait;
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};

fn validated_key(root: &Path, key: &str) -> Result<PathBuf> {
    anyhow::ensure!(
        !key.is_empty()
            && key.len() <= 512
            && !key.starts_with('/')
            && key.split('/').all(|part| {
                !part.is_empty()
                    && part != "."
                    && part != ".."
                    && part.bytes().all(|byte| {
                        byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.')
                    })
            }),
        "invalid snapshot object key"
    );
    Ok(root.join(key))
}

#[derive(Debug, Clone)]
pub struct FilesystemSnapshotObjectStore {
    root: Arc<PathBuf>,
}

impl FilesystemSnapshotObjectStore {
    pub fn for_local_development(root: PathBuf) -> Result<Self> {
        std::fs::create_dir_all(&root)?;
        Ok(Self {
            root: Arc::new(std::fs::canonicalize(root)?),
        })
    }
}

#[derive(Debug)]
struct FilesystemUpload {
    key: String,
    destination: PathBuf,
    temporary: PathBuf,
    file: tokio::fs::File,
    digest: Sha256,
    written: u64,
}

#[async_trait]
impl SnapshotObjectUpload for FilesystemUpload {
    async fn write_chunk(&mut self, bytes: &[u8]) -> Result<()> {
        anyhow::ensure!(
            bytes.len() <= WORKSPACE_UPLOAD_CHUNK_BYTES,
            "snapshot object chunk exceeds policy"
        );
        self.file.write_all(bytes).await?;
        self.digest.update(bytes);
        self.written = self.written.saturating_add(bytes.len() as u64);
        Ok(())
    }

    async fn commit(
        mut self: Box<Self>,
        digest: &str,
        size_bytes: u64,
    ) -> Result<SnapshotObjectRef> {
        self.file.flush().await?;
        self.file.sync_all().await?;
        let actual = format!("sha256:{:x}", self.digest.finalize_reset());
        anyhow::ensure!(actual == digest, "snapshot object digest mismatch");
        anyhow::ensure!(self.written == size_bytes, "snapshot object size mismatch");
        if let Some(parent) = self.destination.parent() {
            tokio::fs::create_dir_all(parent).await?;
        }
        match tokio::fs::rename(&self.temporary, &self.destination).await {
            Ok(()) => {}
            Err(_error) if self.destination.exists() => {
                tokio::fs::remove_file(&self.temporary).await.ok();
                let existing = FilesystemReader {
                    path: self.destination.clone(),
                    digest: digest.to_owned(),
                    len: size_bytes,
                };
                verify_reader(&existing).await?;
            }
            Err(error) => return Err(error.into()),
        }
        Ok(SnapshotObjectRef {
            key: self.key.clone(),
            digest: digest.to_owned(),
            size_bytes,
        })
    }

    async fn abort(self: Box<Self>) -> Result<()> {
        match tokio::fs::remove_file(&self.temporary).await {
            Ok(()) => Ok(()),
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
            Err(error) => Err(error.into()),
        }
    }
}

#[derive(Debug)]
struct FilesystemReader {
    path: PathBuf,
    digest: String,
    len: u64,
}

#[async_trait]
impl SnapshotObjectReader for FilesystemReader {
    fn len(&self) -> u64 {
        self.len
    }

    fn digest(&self) -> &str {
        &self.digest
    }

    async fn read_chunk(&self, offset: u64, max_bytes: usize) -> Result<Vec<u8>> {
        anyhow::ensure!(
            max_bytes <= WORKSPACE_UPLOAD_CHUNK_BYTES,
            "snapshot read chunk exceeds policy"
        );
        if offset >= self.len {
            return Ok(Vec::new());
        }
        let remaining = usize::try_from((self.len - offset).min(max_bytes as u64))?;
        let mut file = tokio::fs::File::open(&self.path).await?;
        file.seek(std::io::SeekFrom::Start(offset)).await?;
        let mut bytes = vec![0; remaining];
        file.read_exact(&mut bytes).await?;
        Ok(bytes)
    }
}

async fn verify_reader(reader: &FilesystemReader) -> Result<()> {
    let mut hasher = Sha256::new();
    let mut offset = 0;
    while offset < reader.len {
        let bytes = reader
            .read_chunk(offset, WORKSPACE_UPLOAD_CHUNK_BYTES)
            .await?;
        anyhow::ensure!(!bytes.is_empty(), "truncated snapshot object");
        hasher.update(&bytes);
        offset = offset.saturating_add(bytes.len() as u64);
    }
    anyhow::ensure!(
        format!("sha256:{:x}", hasher.finalize()) == reader.digest,
        "existing snapshot object digest mismatch"
    );
    Ok(())
}

#[async_trait]
impl SnapshotObjectStore for FilesystemSnapshotObjectStore {
    async fn begin_upload(&self, key: &str) -> Result<Box<dyn SnapshotObjectUpload>> {
        let destination = validated_key(&self.root, key)?;
        let parent = destination
            .parent()
            .context("snapshot object parent missing")?;
        tokio::fs::create_dir_all(parent).await?;
        let temporary = parent.join(format!(".upload-{}", uuid::Uuid::new_v4().simple()));
        let file = tokio::fs::OpenOptions::new()
            .create_new(true)
            .write(true)
            .open(&temporary)
            .await?;
        Ok(Box::new(FilesystemUpload {
            key: key.to_owned(),
            destination,
            temporary,
            file,
            digest: Sha256::new(),
            written: 0,
        }))
    }

    async fn open(&self, object: &SnapshotObjectRef) -> Result<Arc<dyn SnapshotObjectReader>> {
        let path = validated_key(&self.root, &object.key)?;
        let len = tokio::fs::metadata(&path).await?.len();
        anyhow::ensure!(len == object.size_bytes, "snapshot object size mismatch");
        Ok(Arc::new(FilesystemReader {
            path,
            digest: object.digest.clone(),
            len,
        }))
    }

    async fn delete(&self, object: &SnapshotObjectRef) -> Result<()> {
        let path = validated_key(&self.root, &object.key)?;
        match tokio::fs::remove_file(path).await {
            Ok(()) => Ok(()),
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
            Err(error) => Err(error.into()),
        }
    }
}