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