Skip to main content

agent_workspace_contract/
advanced.rs

1use super::*;
2use sha2::{Digest, Sha256};
3
4pub const WORKSPACE_UPLOADS_PATH: &str = "/internal/v1/workspaces/{workspaceId}/uploads";
5pub const WORKSPACE_UPLOAD_PATH: &str = "/internal/v1/workspaces/{workspaceId}/uploads/{uploadId}";
6pub const WORKSPACE_COMMANDS_PATH: &str = "/internal/v1/workspaces/{workspaceId}/commands";
7pub const WORKSPACE_COMMAND_PATH: &str =
8    "/internal/v1/workspaces/{workspaceId}/commands/{commandId}";
9pub const WORKSPACE_COMMAND_OUTPUT_PATH: &str =
10    "/internal/v1/workspaces/{workspaceId}/commands/{commandId}/output";
11pub const WORKSPACE_COMMAND_CANCEL_PATH: &str =
12    "/internal/v1/workspaces/{workspaceId}/commands/{commandId}/cancel";
13pub const WORKSPACE_SNAPSHOTS_PATH: &str = "/internal/v1/workspaces/{workspaceId}/snapshots";
14pub const WORKSPACE_SNAPSHOT_RESTORE_PATH: &str =
15    "/internal/v1/workspaces/{workspaceId}/snapshots/{snapshotId}/restore";
16pub const WORKSPACE_LEASE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/lease";
17pub const WORKSPACE_CLONE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/clone";
18pub const WORKSPACE_MIGRATE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/migrate";
19pub const WORKSPACE_OPERATIONS_PATH: &str = "/internal/v1/workspace-operations/{operationId}";
20
21pub const WORKSPACE_MAX_UPLOAD_BYTES: u64 = WORKSPACE_MAX_FILE_BYTES as u64;
22/// Resumable uploads deliberately use small chunks so request memory is
23/// bounded independently from the maximum completed file size.
24pub const WORKSPACE_UPLOAD_CHUNK_BYTES: usize = 1024 * 1024;
25pub const WORKSPACE_UPLOAD_TTL_MS: u64 = 60 * 60 * 1000;
26pub const WORKSPACE_PREVIEW_TTL_MS: u64 = 60 * 60 * 1000;
27pub const WORKSPACE_MAX_OUTPUT_CHUNKS: usize = 64;
28pub const WORKSPACE_OUTPUT_CHUNK_BYTES: usize = 32 * 1024;
29pub const WORKSPACE_MAX_COMMAND_TIMEOUT_MS: u64 = 60_000;
30pub const WORKSPACE_MAX_COMMAND_MEMORY_BYTES: u64 = 16 * 1024 * 1024 * 1024;
31pub const WORKSPACE_MAX_COMMAND_CPU_MILLIS: u64 = 60_000;
32pub const WORKSPACE_MAX_COMMAND_PROCESSES: u32 = 1024;
33pub const WORKSPACE_MAX_COMMAND_DISK_BYTES: u64 = 64 * 1024 * 1024 * 1024;
34pub const WORKSPACE_MAX_SNAPSHOT_BYTES: usize = 64 * 1024 * 1024;
35pub const WORKSPACE_MAX_RETENTION_MS: u64 = 30 * 24 * 60 * 60 * 1000;
36
37pub fn snapshot_digest(bytes: &[u8]) -> String {
38    format!("sha256:{:x}", Sha256::digest(bytes))
39}
40
41pub const WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC: &[u8; 7] = b"AWSNP1\0";
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
44#[serde(rename_all = "snake_case")]
45pub enum UploadState {
46    Open,
47    Completed,
48    Expired,
49    Aborted,
50}
51
52#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
53#[serde(rename_all = "camelCase")]
54pub struct UploadSession {
55    pub id: String,
56    pub tenant_id: String,
57    pub workspace_id: String,
58    pub path: String,
59    pub offset: u64,
60    pub max_bytes: u64,
61    pub expires_at_ms: u64,
62    pub state: UploadState,
63    pub version: u64,
64    #[serde(default, skip_serializing_if = "Option::is_none")]
65    pub if_match: Option<String>,
66    #[serde(default)]
67    pub request_digest: String,
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize)]
71#[serde(rename_all = "camelCase")]
72pub struct CreateUploadRequest {
73    pub path: String,
74    #[serde(default, skip_serializing_if = "Option::is_none")]
75    pub id: Option<String>,
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub if_match: Option<String>,
78    #[serde(default)]
79    pub ttl_ms: u64,
80}
81
82#[derive(Debug, Clone, Serialize, Deserialize)]
83#[serde(deny_unknown_fields, rename_all = "camelCase")]
84pub struct CreatePreviewRequest {
85    #[serde(default, skip_serializing_if = "Option::is_none")]
86    pub id: Option<String>,
87    pub port: u16,
88    #[serde(default)]
89    pub ttl_ms: u64,
90}
91
92#[derive(Debug, Clone, Serialize, Deserialize)]
93#[serde(rename_all = "camelCase")]
94pub struct UploadChunkRequest {
95    pub offset: u64,
96    pub content_base64: String,
97    #[serde(default)]
98    pub final_chunk: bool,
99    pub version: u64,
100}
101
102#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
103#[serde(rename_all = "snake_case")]
104pub enum CommandState {
105    Queued,
106    Running,
107    Succeeded,
108    Failed,
109    Canceling,
110    Canceled,
111}
112
113#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
114#[serde(rename_all = "camelCase")]
115pub struct CommandResource {
116    pub id: String,
117    pub tenant_id: String,
118    pub workspace_id: String,
119    pub spec: CommandSpec,
120    pub state: CommandState,
121    pub operation_id: String,
122    pub output_end_cursor: u64,
123    pub output_truncated: bool,
124    pub version: u64,
125    pub created_at_ms: u64,
126    pub updated_at_ms: u64,
127    #[serde(default, skip_serializing_if = "Option::is_none")]
128    pub exit_code: Option<i32>,
129    #[serde(default, skip_serializing_if = "Option::is_none")]
130    pub error_code: Option<String>,
131}
132
133#[derive(Debug, Clone)]
134pub struct CommandTransition {
135    pub tenant_id: String,
136    pub id: String,
137    pub expected_version: u64,
138    pub state: CommandState,
139    pub exit_code: Option<i32>,
140    pub error_code: Option<String>,
141    pub now_ms: u64,
142}
143
144#[derive(Debug, Clone, Serialize, Deserialize)]
145#[serde(rename_all = "camelCase")]
146pub struct CreateCommandRequest {
147    #[serde(default, skip_serializing_if = "Option::is_none")]
148    pub id: Option<String>,
149    pub command: CommandSpec,
150}
151
152#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
153#[serde(rename_all = "snake_case")]
154pub enum OutputStream {
155    Stdout,
156    Stderr,
157}
158
159#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
160#[serde(rename_all = "camelCase")]
161pub struct OutputChunk {
162    pub command_id: String,
163    pub cursor: u64,
164    pub stream: OutputStream,
165    pub content: String,
166    pub bytes: usize,
167}
168
169#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
170#[serde(rename_all = "camelCase")]
171pub struct OutputPage {
172    pub chunks: Vec<OutputChunk>,
173    #[serde(default, skip_serializing_if = "Option::is_none")]
174    pub next_cursor: Option<u64>,
175    pub truncated: bool,
176}
177
178#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
179#[serde(rename_all = "camelCase")]
180pub struct Snapshot {
181    pub id: String,
182    pub tenant_id: String,
183    pub workspace_id: String,
184    pub digest: String,
185    pub source_version: u64,
186    pub size_bytes: u64,
187    pub file_count: usize,
188    pub created_at_ms: u64,
189    pub expires_at_ms: u64,
190    #[serde(default)]
191    pub request_digest: String,
192    /// Durable object-store reference. Legacy rows can omit this while they
193    /// are migrated from the former inline SQL blob representation.
194    #[serde(default, skip_serializing_if = "Option::is_none")]
195    pub object_ref: Option<SnapshotObjectRef>,
196}
197
198#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
199#[serde(rename_all = "camelCase")]
200pub struct SnapshotObjectRef {
201    pub key: String,
202    pub digest: String,
203    pub size_bytes: u64,
204}
205
206#[derive(Clone)]
207pub struct StoredSnapshot {
208    pub snapshot: Snapshot,
209    /// Read-only compatibility payload for rows created before object storage
210    /// was introduced. New writes must always leave this empty.
211    pub legacy_blob: Option<Vec<u8>>,
212}
213
214impl Debug for StoredSnapshot {
215    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
216        formatter
217            .debug_struct("StoredSnapshot")
218            .field("snapshot", &self.snapshot)
219            .field(
220                "legacy_blob",
221                &self.legacy_blob.as_ref().map(|bytes| bytes.len()),
222            )
223            .finish()
224    }
225}
226
227#[derive(Debug, Clone, Serialize, Deserialize)]
228#[serde(rename_all = "camelCase")]
229pub struct CreateSnapshotRequest {
230    #[serde(default, skip_serializing_if = "Option::is_none")]
231    pub id: Option<String>,
232    #[serde(default)]
233    pub retention_ms: u64,
234}
235
236#[derive(Debug, Clone, Serialize, Deserialize)]
237#[serde(rename_all = "camelCase")]
238pub struct CloneWorkspaceRequest {
239    pub destination_workspace_id: String,
240    pub destination_project_id: String,
241    #[serde(default, skip_serializing_if = "Option::is_none")]
242    pub snapshot_id: Option<String>,
243}
244
245#[derive(Debug, Clone, Serialize, Deserialize)]
246#[serde(rename_all = "camelCase")]
247pub struct MigrateWorkspaceRequest {
248    pub destination_backend: String,
249    #[serde(default)]
250    pub config: WorkspaceConfig,
251}
252
253#[derive(Debug, Clone, Serialize, Deserialize)]
254#[serde(rename_all = "camelCase")]
255pub struct KeepAliveRequest {
256    pub owner: String,
257    pub fencing_token: u64,
258    pub lease_ms: u64,
259}
260
261#[derive(Debug, Clone)]
262pub struct WorkspaceLeaseKeepAlive {
263    pub tenant_id: String,
264    pub id: String,
265    pub owner: String,
266    pub fencing_token: u64,
267    pub now_ms: u64,
268    pub lease_ms: u64,
269}
270
271#[derive(Debug, Clone)]
272pub struct CommandPermitRequest {
273    pub command_id: String,
274    pub tenant_id: String,
275    pub owner: String,
276    pub now_ms: u64,
277    pub lease_ms: u64,
278    pub global_limit: usize,
279    pub tenant_limit: usize,
280}
281
282/// Bounded, random-access reader used across application and object-store
283/// ports. Implementations must never return more than `max_bytes`.
284#[async_trait]
285pub trait SnapshotObjectReader: Send + Sync + Debug {
286    fn len(&self) -> u64;
287    fn is_empty(&self) -> bool {
288        self.len() == 0
289    }
290    fn digest(&self) -> &str;
291    async fn read_chunk(&self, offset: u64, max_bytes: usize) -> Result<Vec<u8>>;
292}
293
294#[async_trait]
295pub trait SnapshotObjectUpload: Send + Debug {
296    async fn write_chunk(&mut self, bytes: &[u8]) -> Result<()>;
297    async fn commit(self: Box<Self>, digest: &str, size_bytes: u64) -> Result<SnapshotObjectRef>;
298    async fn abort(self: Box<Self>) -> Result<()>;
299}
300
301#[async_trait]
302pub trait SnapshotObjectStore: Send + Sync + Debug {
303    async fn begin_upload(&self, key: &str) -> Result<Box<dyn SnapshotObjectUpload>>;
304    async fn open(&self, object: &SnapshotObjectRef) -> Result<Arc<dyn SnapshotObjectReader>>;
305    async fn delete(&self, object: &SnapshotObjectRef) -> Result<()>;
306    async fn check_readiness(&self) -> Result<()> {
307        Ok(())
308    }
309}
310
311#[async_trait]
312pub trait SnapshotProvider: Send + Sync + Debug {
313    async fn export_snapshot(
314        &self,
315        durable_backend_id: &str,
316    ) -> Result<Arc<dyn SnapshotObjectReader>>;
317    async fn import_snapshot(
318        &self,
319        workspace_id: &str,
320        config: &WorkspaceConfig,
321        snapshot: Arc<dyn SnapshotObjectReader>,
322        idempotency_key: &str,
323    ) -> Result<ProvisionedWorkspace>;
324}
325
326#[async_trait]
327pub trait PreviewLifecycleProvider: Send + Sync + Debug {
328    async fn create_preview(
329        &self,
330        durable_backend_id: &str,
331        request_id: &str,
332        port: u16,
333        ttl_ms: u64,
334    ) -> Result<StoredPreview>;
335    async fn stop_preview(&self, provider_ref: &str) -> Result<()>;
336}
337
338#[derive(Clone)]
339pub struct WorkspaceSession {
340    pub workspace: Arc<dyn Workspace>,
341    pub upload_writer: Arc<dyn WorkspaceUploadWriter>,
342    pub executor: Arc<dyn WorkspaceCommandExecutor>,
343    pub snapshot_provider: Option<Arc<dyn SnapshotProvider>>,
344    pub snapshot_objects: Option<Arc<dyn SnapshotObjectStore>>,
345    pub preview_provider: Option<Arc<dyn PreviewLifecycleProvider>>,
346}
347
348impl Debug for WorkspaceSession {
349    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
350        formatter
351            .debug_struct("WorkspaceSession")
352            .finish_non_exhaustive()
353    }
354}
355
356#[async_trait]
357pub trait WorkspaceSessionResolver: Send + Sync + Debug {
358    async fn resolve(&self, workspace: &WorkspaceRecord) -> Result<WorkspaceSession>;
359    /// Resolve the importer owned by the requested destination backend. A
360    /// migration must never reuse the source session's importer implicitly.
361    async fn resolve_snapshot_destination(
362        &self,
363        backend: &str,
364    ) -> Result<Arc<dyn SnapshotProvider>>;
365    async fn check_readiness(&self) -> Result<()> {
366        Ok(())
367    }
368}
369
370#[async_trait]
371pub trait WorkspaceAdvancedRepository: Send + Sync + Debug {
372    async fn put_upload(&self, value: &UploadSession) -> Result<()>;
373    async fn get_upload(&self, tenant_id: &str, id: &str) -> Result<Option<UploadSession>>;
374    async fn advance_upload(
375        &self,
376        tenant_id: &str,
377        id: &str,
378        expected_offset: u64,
379        expected_version: u64,
380        new_offset: u64,
381        state: UploadState,
382    ) -> Result<UploadSession>;
383    async fn put_command(&self, value: &CommandResource) -> Result<()>;
384    async fn put_command_operation(
385        &self,
386        command: &CommandResource,
387        operation: &WorkspaceOperation,
388        payload_json: &str,
389    ) -> Result<()>;
390    async fn get_command(&self, tenant_id: &str, id: &str) -> Result<Option<CommandResource>>;
391    async fn transition_command(&self, request: CommandTransition) -> Result<CommandResource>;
392    async fn append_output(
393        &self,
394        tenant_id: &str,
395        command_id: &str,
396        chunks: &[OutputChunk],
397        truncated: bool,
398    ) -> Result<()>;
399    async fn output_page(
400        &self,
401        tenant_id: &str,
402        command_id: &str,
403        after: u64,
404        limit: usize,
405    ) -> Result<OutputPage>;
406    async fn put_snapshot(&self, value: &Snapshot, legacy_blob: Option<&[u8]>) -> Result<()>;
407    async fn get_snapshot(&self, tenant_id: &str, id: &str) -> Result<Option<StoredSnapshot>>;
408    async fn list_expired_snapshots(
409        &self,
410        tenant_id: &str,
411        now_ms: u64,
412        limit: usize,
413    ) -> Result<Vec<Snapshot>>;
414    async fn delete_expired_snapshot_manifest(
415        &self,
416        tenant_id: &str,
417        id: &str,
418        expected_expires_at_ms: u64,
419    ) -> Result<bool>;
420    /// Persist a failed object deletion so poison objects back off and cannot
421    /// occupy every slot in subsequent bounded retention batches.
422    async fn record_snapshot_cleanup_failure(
423        &self,
424        tenant_id: &str,
425        id: &str,
426        now_ms: u64,
427        error: &str,
428    ) -> Result<()> {
429        let _ = (tenant_id, id, now_ms, error);
430        Ok(())
431    }
432    async fn list_expired_previews(
433        &self,
434        tenant_id: &str,
435        workspace_id: &str,
436        now_ms: u64,
437        limit: usize,
438    ) -> Result<Vec<StoredPreview>>;
439    async fn reconcile_expired(&self, tenant_id: &str, now_ms: u64) -> Result<usize>;
440    async fn keep_alive_workspace(
441        &self,
442        request: WorkspaceLeaseKeepAlive,
443    ) -> Result<WorkspaceLease>;
444}
445
446/// Aggregate metadata port used by the application layer. Concrete SQL/storage
447/// technology belongs in adapter crates.
448pub trait WorkspaceRepository: WorkspaceControlRepository + WorkspaceAdvancedRepository {}
449impl<T> WorkspaceRepository for T where T: WorkspaceControlRepository + WorkspaceAdvancedRepository {}
450
451#[async_trait]
452pub trait WorkspaceUploadWriter: Send + Sync + Debug {
453    async fn write_upload_chunk(&self, path: &str, offset: u64, bytes: &[u8]) -> Result<u64>;
454}
455
456#[async_trait]
457pub trait CommandPermitRepository: Send + Sync + Debug {
458    async fn acquire_command_permit(&self, request: CommandPermitRequest) -> Result<bool>;
459    async fn release_command_permit(&self, command_id: &str, owner: &str) -> Result<()>;
460    async fn cancel_command_permit(&self, command_id: &str) -> Result<bool>;
461    async fn command_canceled(&self, command_id: &str) -> Result<bool>;
462}