Skip to main content

agent_workspace_contract/
control.rs

1use super::*;
2use std::collections::BTreeMap;
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
5#[serde(rename_all = "snake_case")]
6pub enum WorkspaceLifecycle {
7    Provisioning,
8    Ready,
9    Suspended,
10    Deleting,
11    Deleted,
12    Error,
13}
14
15pub const WORKSPACES_PATH: &str = "/internal/v1/workspaces";
16pub const WORKSPACE_PATH: &str = "/internal/v1/workspaces/{workspaceId}";
17pub const WORKSPACE_SUSPEND_PATH: &str = "/internal/v1/workspaces/{workspaceId}/suspend";
18pub const WORKSPACE_RESUME_PATH: &str = "/internal/v1/workspaces/{workspaceId}/resume";
19pub const WORKSPACE_RECONCILE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/reconcile";
20pub const WORKSPACE_USAGE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/usage";
21pub const WORKSPACE_OPERATION_CANCEL_PATH: &str =
22    "/internal/v1/workspace-operations/{operationId}/cancel";
23pub const WORKSPACE_MAX_PAGE_SIZE: usize = 200;
24
25#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
26#[serde(deny_unknown_fields, rename_all = "camelCase")]
27pub struct CreateWorkspaceRequest {
28    #[serde(default, skip_serializing_if = "Option::is_none")]
29    pub id: Option<String>,
30    pub project_id: String,
31    pub config: WorkspaceConfig,
32    #[serde(default)]
33    pub labels: BTreeMap<String, String>,
34}
35
36#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
37#[serde(deny_unknown_fields, rename_all = "camelCase")]
38pub struct WorkspaceActionRequest {
39    pub expected_version: u64,
40    #[serde(default, skip_serializing_if = "Option::is_none")]
41    pub reason: Option<String>,
42}
43
44#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
45#[serde(rename_all = "camelCase")]
46pub struct WorkspacePage {
47    pub items: Vec<WorkspaceRecord>,
48    #[serde(default, skip_serializing_if = "Option::is_none")]
49    pub next_cursor: Option<String>,
50}
51
52#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
53#[serde(rename_all = "camelCase")]
54pub struct ProvisionedWorkspace {
55    pub durable_backend_id: String,
56    pub root: String,
57}
58
59#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
60#[serde(rename_all = "camelCase")]
61pub struct WorkspaceLease {
62    pub owner: String,
63    pub deadline_ms: u64,
64    pub fencing_token: u64,
65}
66
67#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
68#[serde(rename_all = "camelCase")]
69pub struct WorkspaceRecord {
70    pub id: String,
71    pub tenant_id: String,
72    pub project_id: String,
73    pub backend: String,
74    pub durable_backend_id: String,
75    pub root: String,
76    #[serde(default)]
77    pub config: WorkspaceConfig,
78    pub lifecycle: WorkspaceLifecycle,
79    pub version: u64,
80    #[serde(default, skip_serializing_if = "Option::is_none")]
81    pub lease: Option<WorkspaceLease>,
82    pub created_at_ms: u64,
83    pub updated_at_ms: u64,
84    #[serde(default)]
85    pub labels: BTreeMap<String, String>,
86}
87
88#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
89#[serde(rename_all = "snake_case")]
90pub enum OperationState {
91    Queued,
92    Running,
93    Succeeded,
94    Failed,
95    Canceling,
96    Canceled,
97}
98
99impl OperationState {
100    pub fn terminal(self) -> bool {
101        matches!(self, Self::Succeeded | Self::Failed | Self::Canceled)
102    }
103}
104
105#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
106#[serde(rename_all = "camelCase")]
107pub struct WorkspaceOperation {
108    pub id: String,
109    pub tenant_id: String,
110    pub workspace_id: String,
111    pub kind: String,
112    pub state: OperationState,
113    pub phase: String,
114    pub attempt: u32,
115    pub version: u64,
116    #[serde(default, skip_serializing_if = "Option::is_none")]
117    pub lease: Option<WorkspaceLease>,
118    #[serde(default, skip_serializing_if = "Option::is_none")]
119    pub error_code: Option<String>,
120    #[serde(default, skip_serializing_if = "Option::is_none")]
121    pub result_resource_id: Option<String>,
122    pub created_at_ms: u64,
123    pub updated_at_ms: u64,
124    #[serde(default, skip_serializing_if = "Option::is_none")]
125    pub next_poll_after_ms: Option<u64>,
126}
127
128impl WorkspaceOperation {
129    pub fn refresh_poll_hint(&mut self) {
130        self.next_poll_after_ms = match self.state {
131            OperationState::Queued => Some(250),
132            OperationState::Running | OperationState::Canceling => Some(500),
133            OperationState::Succeeded | OperationState::Failed | OperationState::Canceled => None,
134        };
135    }
136}
137
138#[derive(Debug, Clone)]
139pub struct OperationTransition {
140    pub operation_id: String,
141    pub owner: String,
142    pub fencing_token: u64,
143    pub expected_version: u64,
144    pub state: OperationState,
145    pub phase: String,
146    pub error_code: Option<String>,
147    pub result_resource_id: Option<String>,
148    pub now_ms: u64,
149}
150
151#[derive(Debug, Clone, Serialize, Deserialize)]
152#[serde(rename_all = "camelCase")]
153pub struct StoredChangeSet {
154    pub tenant_id: String,
155    pub workspace_id: String,
156    pub change_set: WorkspaceChangeSet,
157    pub base_version: u64,
158    pub workspace_version: u64,
159    pub expires_at_ms: u64,
160    pub request_digest: String,
161}
162
163#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
164#[serde(rename_all = "camelCase")]
165pub struct StoredPreview {
166    pub id: String,
167    pub tenant_id: String,
168    pub workspace_id: String,
169    pub provider_ref: String,
170    pub url: String,
171    pub state: String,
172    pub expires_at_ms: u64,
173    pub version: u64,
174    #[serde(default)]
175    pub request_digest: String,
176}
177
178#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
179#[serde(rename_all = "camelCase")]
180pub struct CommandSpec {
181    pub argv: Vec<String>,
182    #[serde(default, skip_serializing_if = "Option::is_none")]
183    pub cwd: Option<String>,
184    #[serde(default)]
185    pub env: BTreeMap<String, SecretRef>,
186    #[serde(default)]
187    pub shell: bool,
188    pub timeout_ms: u64,
189    pub memory_bytes: u64,
190    pub cpu_millis: u64,
191    pub max_processes: u32,
192    /// Whether provider-side egress is denied or permitted for this command.
193    /// Providers must reject unsupported policy rather than silently ignore it.
194    #[serde(default)]
195    pub network: CommandNetworkPolicy,
196    /// Maximum additional writable bytes for the command sandbox.
197    #[serde(default = "default_command_disk_bytes")]
198    pub disk_bytes: u64,
199    pub max_output_bytes: usize,
200}
201
202#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
203#[serde(rename_all = "snake_case")]
204pub enum CommandNetworkPolicy {
205    #[default]
206    Deny,
207    Allow,
208}
209
210fn default_command_disk_bytes() -> u64 {
211    1024 * 1024 * 1024
212}
213
214#[async_trait]
215pub trait WorkspaceControlRepository: Send + Sync + Debug {
216    async fn put_workspace(&self, workspace: &WorkspaceRecord) -> Result<()>;
217    async fn get_workspace(&self, tenant_id: &str, id: &str) -> Result<Option<WorkspaceRecord>>;
218    async fn list_workspaces(
219        &self,
220        tenant_id: &str,
221        project_id: &str,
222        after: Option<&str>,
223        limit: usize,
224    ) -> Result<WorkspacePage>;
225    async fn list_workspaces_global(
226        &self,
227        after: Option<(&str, &str)>,
228        limit: usize,
229    ) -> Result<WorkspacePage>;
230    async fn put_workspace_operation(
231        &self,
232        workspace: &WorkspaceRecord,
233        operation: &WorkspaceOperation,
234        payload_json: &str,
235        idempotency_key: &str,
236        request_digest: &str,
237    ) -> Result<WorkspaceOperation>;
238    async fn compare_and_swap_workspace(
239        &self,
240        tenant_id: &str,
241        id: &str,
242        expected_version: u64,
243        lifecycle: WorkspaceLifecycle,
244        now_ms: u64,
245    ) -> Result<WorkspaceRecord>;
246    async fn claim_workspace_lease(
247        &self,
248        tenant_id: &str,
249        id: &str,
250        owner: &str,
251        now_ms: u64,
252        lease_ms: u64,
253    ) -> Result<WorkspaceLease>;
254    async fn release_workspace_lease(
255        &self,
256        tenant_id: &str,
257        id: &str,
258        owner: &str,
259        fencing_token: u64,
260        now_ms: u64,
261    ) -> Result<()>;
262    async fn fenced_workspace_update(
263        &self,
264        tenant_id: &str,
265        id: &str,
266        owner: &str,
267        fencing_token: u64,
268        lifecycle: WorkspaceLifecycle,
269        now_ms: u64,
270    ) -> Result<WorkspaceRecord>;
271    async fn fenced_replace_workspace(
272        &self,
273        workspace: &WorkspaceRecord,
274        owner: &str,
275        fencing_token: u64,
276        now_ms: u64,
277    ) -> Result<WorkspaceRecord>;
278    async fn put_operation(&self, operation: &WorkspaceOperation) -> Result<()>;
279    async fn put_operation_with_input(
280        &self,
281        operation: &WorkspaceOperation,
282        payload_json: &str,
283    ) -> Result<()>;
284    async fn put_idempotent_operation(
285        &self,
286        operation: &WorkspaceOperation,
287        payload_json: &str,
288        idempotency_key: &str,
289        request_digest: &str,
290    ) -> Result<WorkspaceOperation>;
291    async fn get_operation(&self, operation_id: &str) -> Result<Option<WorkspaceOperation>>;
292    async fn claim_operation(
293        &self,
294        operation_id: &str,
295        owner: &str,
296        now_ms: u64,
297        lease_ms: u64,
298    ) -> Result<WorkspaceOperation>;
299    async fn renew_operation_lease(
300        &self,
301        operation_id: &str,
302        owner: &str,
303        fencing_token: u64,
304        now_ms: u64,
305        lease_ms: u64,
306    ) -> Result<WorkspaceOperation>;
307    async fn transition_operation(
308        &self,
309        transition: &OperationTransition,
310    ) -> Result<WorkspaceOperation>;
311    async fn request_operation_cancel(
312        &self,
313        tenant_id: &str,
314        operation_id: &str,
315        expected_version: u64,
316        now_ms: u64,
317    ) -> Result<WorkspaceOperation>;
318    async fn put_operation_input(&self, operation_id: &str, payload_json: &str) -> Result<()>;
319    async fn get_operation_input(&self, operation_id: &str) -> Result<Option<String>>;
320    async fn list_claimable_operations(
321        &self,
322        tenant_id: &str,
323        now_ms: u64,
324        limit: usize,
325    ) -> Result<Vec<WorkspaceOperation>>;
326    async fn list_claimable_operations_global(
327        &self,
328        now_ms: u64,
329        limit: usize,
330    ) -> Result<Vec<WorkspaceOperation>>;
331    async fn put_change_set(&self, value: &StoredChangeSet) -> Result<()>;
332    async fn get_change_set(&self, tenant_id: &str, id: &str) -> Result<Option<StoredChangeSet>>;
333    async fn put_preview(&self, value: &StoredPreview) -> Result<()>;
334    async fn get_preview(&self, tenant_id: &str, id: &str) -> Result<Option<StoredPreview>>;
335    async fn stop_preview(
336        &self,
337        tenant_id: &str,
338        id: &str,
339        expected_version: u64,
340    ) -> Result<StoredPreview>;
341    async fn retain_tenant(&self, tenant_id: &str, now_ms: u64, max_rows: usize) -> Result<usize>;
342    async fn record_worker_heartbeat(&self, worker_id: &str, now_ms: u64) -> Result<()> {
343        let _ = (worker_id, now_ms);
344        Ok(())
345    }
346    async fn check_readiness(
347        &self,
348        worker_id: &str,
349        now_ms: u64,
350        max_worker_staleness_ms: u64,
351        max_operation_backlog: usize,
352    ) -> Result<()> {
353        let _ = (
354            worker_id,
355            now_ms,
356            max_worker_staleness_ms,
357            max_operation_backlog,
358        );
359        Ok(())
360    }
361}
362
363#[async_trait]
364pub trait SecretResolver: Send + Sync + Debug {
365    /// Secret bytes must remain inside the trusted adapter/executor boundary.
366    async fn resolve(&self, tenant_id: &str, secret: &SecretRef) -> Result<Vec<u8>>;
367}
368
369#[async_trait]
370pub trait WorkspaceProviderAdapter: Send + Sync + Debug {
371    async fn provision(
372        &self,
373        tenant_id: &str,
374        workspace_id: &str,
375        config: &WorkspaceConfig,
376        idempotency_key: &str,
377    ) -> Result<ProvisionedWorkspace>;
378    async fn reconnect(&self, durable_backend_id: &str) -> Result<()>;
379    async fn suspend(&self, durable_backend_id: &str, idempotency_key: &str) -> Result<()>;
380    async fn resume(&self, durable_backend_id: &str, idempotency_key: &str) -> Result<()>;
381    async fn inspect(&self, workspace: &WorkspaceRecord) -> Result<WorkspaceResourceUsage>;
382    async fn delete(&self, durable_backend_id: &str, idempotency_key: &str) -> Result<()>;
383    async fn cancel_command(&self, durable_backend_id: &str, command_id: &str) -> Result<()>;
384}
385
386#[async_trait]
387pub trait WorkspaceProviderRegistry: Send + Sync + Debug {
388    async fn provider(&self, backend: &str) -> Result<Arc<dyn WorkspaceProviderAdapter>>;
389}
390
391#[async_trait]
392pub trait WorkspaceCommandExecutor: Send + Sync + Debug {
393    async fn execute(
394        &self,
395        tenant_id: &str,
396        workspace_id: &str,
397        command_id: &str,
398        command: &CommandSpec,
399    ) -> Result<CmdOutput>;
400    async fn cancel(&self, tenant_id: &str, command_id: &str) -> Result<()>;
401}