agent_workspace_adapter_local/
snapshot_provider.rs

1use crate::{FilesystemSnapshotObjectStore, LocalWorkspace};
2use agent_workspace_contract::*;
3use anyhow::{Context, Result};
4use async_trait::async_trait;
5use sha2::{Digest, Sha256};
6use std::path::PathBuf;
7use std::sync::Arc;
8
9/// Explicit local/test snapshot importer. The configured root is the only
10/// directory under which destination workspace IDs may be created.
11#[derive(Debug, Clone)]
12pub struct LocalSnapshotProvider {
13    destination_root: PathBuf,
14    objects: FilesystemSnapshotObjectStore,
15}
16
17impl LocalSnapshotProvider {
18    pub fn for_local_development(destination_root: PathBuf) -> Result<Self> {
19        std::fs::create_dir_all(&destination_root)?;
20        let objects = FilesystemSnapshotObjectStore::for_local_development(
21            destination_root.join(".snapshot-objects"),
22        )?;
23        Ok(Self {
24            destination_root,
25            objects,
26        })
27    }
28
29    fn destination(&self, workspace_id: &str) -> Result<PathBuf> {
30        anyhow::ensure!(
31            !workspace_id.is_empty()
32                && workspace_id.len() <= 128
33                && workspace_id
34                    .bytes()
35                    .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')),
36            "invalid destination workspace id"
37        );
38        Ok(self.destination_root.join(workspace_id))
39    }
40}
41
42async fn exact(reader: &dyn SnapshotObjectReader, offset: &mut u64, len: usize) -> Result<Vec<u8>> {
43    let mut output = Vec::with_capacity(len);
44    while output.len() < len {
45        let bytes = reader
46            .read_chunk(
47                *offset,
48                (len - output.len()).min(WORKSPACE_UPLOAD_CHUNK_BYTES),
49            )
50            .await?;
51        anyhow::ensure!(!bytes.is_empty(), "truncated snapshot object");
52        *offset = offset.saturating_add(bytes.len() as u64);
53        output.extend_from_slice(&bytes);
54    }
55    Ok(output)
56}
57
58async fn integer<const N: usize>(
59    reader: &dyn SnapshotObjectReader,
60    offset: &mut u64,
61) -> Result<[u8; N]> {
62    exact(reader, offset, N)
63        .await?
64        .try_into()
65        .map_err(|_| anyhow::anyhow!("invalid snapshot integer"))
66}
67
68#[async_trait]
69impl SnapshotProvider for LocalSnapshotProvider {
70    async fn export_snapshot(
71        &self,
72        durable_backend_id: &str,
73    ) -> Result<Arc<dyn SnapshotObjectReader>> {
74        let workspace = LocalWorkspace::new(PathBuf::from(durable_backend_id));
75        let key = format!("exports/{}.awsnp", uuid::Uuid::new_v4().simple());
76        let mut upload = self.objects.begin_upload(&key).await?;
77        let mut digest = Sha256::new();
78        let mut total = WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC.len();
79        upload.write_chunk(WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC).await?;
80        digest.update(WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC);
81        for path in workspace.walk_tree(".", 64).await? {
82            if workspace.is_dir(&path).await.unwrap_or(false) {
83                continue;
84            }
85            let bytes = workspace.read_file_bytes(&path).await?;
86            total = total
87                .saturating_add(12)
88                .saturating_add(path.len())
89                .saturating_add(bytes.len());
90            anyhow::ensure!(
91                total <= WORKSPACE_MAX_SNAPSHOT_BYTES,
92                "snapshot byte quota exceeded"
93            );
94            let path_len = u32::try_from(path.len())?.to_be_bytes();
95            let content_len = u64::try_from(bytes.len())?.to_be_bytes();
96            for chunk in [path_len.as_slice(), content_len.as_slice(), path.as_bytes()] {
97                upload.write_chunk(chunk).await?;
98                digest.update(chunk);
99            }
100            for chunk in bytes.chunks(WORKSPACE_UPLOAD_CHUNK_BYTES) {
101                upload.write_chunk(chunk).await?;
102                digest.update(chunk);
103            }
104        }
105        let digest = format!("sha256:{:x}", digest.finalize());
106        let object = upload.commit(&digest, total as u64).await?;
107        self.objects.open(&object).await
108    }
109
110    async fn import_snapshot(
111        &self,
112        workspace_id: &str,
113        config: &WorkspaceConfig,
114        snapshot: Arc<dyn SnapshotObjectReader>,
115        _idempotency_key: &str,
116    ) -> Result<ProvisionedWorkspace> {
117        anyhow::ensure!(
118            matches!(config.backend.as_str(), "local" | "local-test"),
119            "unsupported local snapshot destination backend"
120        );
121        anyhow::ensure!(
122            snapshot.len() <= WORKSPACE_MAX_SNAPSHOT_BYTES as u64,
123            "snapshot byte quota exceeded"
124        );
125        let mut hasher = Sha256::new();
126        let mut verify_offset = 0;
127        while verify_offset < snapshot.len() {
128            let bytes = snapshot
129                .read_chunk(verify_offset, WORKSPACE_UPLOAD_CHUNK_BYTES)
130                .await?;
131            anyhow::ensure!(!bytes.is_empty(), "truncated snapshot object");
132            hasher.update(&bytes);
133            verify_offset = verify_offset.saturating_add(bytes.len() as u64);
134        }
135        anyhow::ensure!(
136            format!("sha256:{:x}", hasher.finalize()) == snapshot.digest(),
137            "snapshot digest mismatch"
138        );
139        let destination = self.destination(workspace_id)?;
140        let workspace = LocalWorkspace::new(destination.clone());
141        let mut offset = 0;
142        anyhow::ensure!(
143            exact(
144                snapshot.as_ref(),
145                &mut offset,
146                WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC.len()
147            )
148            .await?
149                == WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC,
150            "unsupported snapshot format"
151        );
152        while offset < snapshot.len() {
153            let path_len =
154                u32::from_be_bytes(integer(snapshot.as_ref(), &mut offset).await?) as usize;
155            let content_len = usize::try_from(u64::from_be_bytes(
156                integer(snapshot.as_ref(), &mut offset).await?,
157            ))?;
158            anyhow::ensure!(
159                path_len <= WORKSPACE_MAX_PATH_BYTES && content_len <= WORKSPACE_MAX_FILE_BYTES,
160                "snapshot entry exceeds policy"
161            );
162            let path = String::from_utf8(exact(snapshot.as_ref(), &mut offset, path_len).await?)?;
163            let content = exact(snapshot.as_ref(), &mut offset, content_len).await?;
164            workspace.write_file_bytes(&path, &content).await?;
165        }
166        let root = destination
167            .to_str()
168            .context("snapshot destination path is not UTF-8")?
169            .to_owned();
170        Ok(ProvisionedWorkspace {
171            durable_backend_id: root.clone(),
172            root,
173        })
174    }
175}
agent_workspace_adapter_local/snapshot_provider.rs

agent_workspace_adapter_local/
snapshot_provider.rs
