agent-workspace-contract 0.1.0

use super::*;
use std::collections::BTreeMap;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkspaceLifecycle {
    Provisioning,
    Ready,
    Suspended,
    Deleting,
    Deleted,
    Error,
}

pub const WORKSPACES_PATH: &str = "/internal/v1/workspaces";
pub const WORKSPACE_PATH: &str = "/internal/v1/workspaces/{workspaceId}";
pub const WORKSPACE_SUSPEND_PATH: &str = "/internal/v1/workspaces/{workspaceId}/suspend";
pub const WORKSPACE_RESUME_PATH: &str = "/internal/v1/workspaces/{workspaceId}/resume";
pub const WORKSPACE_RECONCILE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/reconcile";
pub const WORKSPACE_USAGE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/usage";
pub const WORKSPACE_OPERATION_CANCEL_PATH: &str =
    "/internal/v1/workspace-operations/{operationId}/cancel";
pub const WORKSPACE_MAX_PAGE_SIZE: usize = 200;

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct CreateWorkspaceRequest {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    pub project_id: String,
    pub config: WorkspaceConfig,
    #[serde(default)]
    pub labels: BTreeMap<String, String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct WorkspaceActionRequest {
    pub expected_version: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reason: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WorkspacePage {
    pub items: Vec<WorkspaceRecord>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub next_cursor: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ProvisionedWorkspace {
    pub durable_backend_id: String,
    pub root: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceLease {
    pub owner: String,
    pub deadline_ms: u64,
    pub fencing_token: u64,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceRecord {
    pub id: String,
    pub tenant_id: String,
    pub project_id: String,
    pub backend: String,
    pub durable_backend_id: String,
    pub root: String,
    #[serde(default)]
    pub config: WorkspaceConfig,
    pub lifecycle: WorkspaceLifecycle,
    pub version: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub lease: Option<WorkspaceLease>,
    pub created_at_ms: u64,
    pub updated_at_ms: u64,
    #[serde(default)]
    pub labels: BTreeMap<String, String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum OperationState {
    Queued,
    Running,
    Succeeded,
    Failed,
    Canceling,
    Canceled,
}

impl OperationState {
    pub fn terminal(self) -> bool {
        matches!(self, Self::Succeeded | Self::Failed | Self::Canceled)
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WorkspaceOperation {
    pub id: String,
    pub tenant_id: String,
    pub workspace_id: String,
    pub kind: String,
    pub state: OperationState,
    pub phase: String,
    pub attempt: u32,
    pub version: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub lease: Option<WorkspaceLease>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error_code: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub result_resource_id: Option<String>,
    pub created_at_ms: u64,
    pub updated_at_ms: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub next_poll_after_ms: Option<u64>,
}

impl WorkspaceOperation {
    pub fn refresh_poll_hint(&mut self) {
        self.next_poll_after_ms = match self.state {
            OperationState::Queued => Some(250),
            OperationState::Running | OperationState::Canceling => Some(500),
            OperationState::Succeeded | OperationState::Failed | OperationState::Canceled => None,
        };
    }
}

#[derive(Debug, Clone)]
pub struct OperationTransition {
    pub operation_id: String,
    pub owner: String,
    pub fencing_token: u64,
    pub expected_version: u64,
    pub state: OperationState,
    pub phase: String,
    pub error_code: Option<String>,
    pub result_resource_id: Option<String>,
    pub now_ms: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct StoredChangeSet {
    pub tenant_id: String,
    pub workspace_id: String,
    pub change_set: WorkspaceChangeSet,
    pub base_version: u64,
    pub workspace_version: u64,
    pub expires_at_ms: u64,
    pub request_digest: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct StoredPreview {
    pub id: String,
    pub tenant_id: String,
    pub workspace_id: String,
    pub provider_ref: String,
    pub url: String,
    pub state: String,
    pub expires_at_ms: u64,
    pub version: u64,
    #[serde(default)]
    pub request_digest: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CommandSpec {
    pub argv: Vec<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cwd: Option<String>,
    #[serde(default)]
    pub env: BTreeMap<String, SecretRef>,
    #[serde(default)]
    pub shell: bool,
    pub timeout_ms: u64,
    pub memory_bytes: u64,
    pub cpu_millis: u64,
    pub max_processes: u32,
    /// Whether provider-side egress is denied or permitted for this command.
    /// Providers must reject unsupported policy rather than silently ignore it.
    #[serde(default)]
    pub network: CommandNetworkPolicy,
    /// Maximum additional writable bytes for the command sandbox.
    #[serde(default = "default_command_disk_bytes")]
    pub disk_bytes: u64,
    pub max_output_bytes: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum CommandNetworkPolicy {
    #[default]
    Deny,
    Allow,
}

fn default_command_disk_bytes() -> u64 {
    1024 * 1024 * 1024
}

#[async_trait]
pub trait WorkspaceControlRepository: Send + Sync + Debug {
    async fn put_workspace(&self, workspace: &WorkspaceRecord) -> Result<()>;
    async fn get_workspace(&self, tenant_id: &str, id: &str) -> Result<Option<WorkspaceRecord>>;
    async fn list_workspaces(
        &self,
        tenant_id: &str,
        project_id: &str,
        after: Option<&str>,
        limit: usize,
    ) -> Result<WorkspacePage>;
    async fn list_workspaces_global(
        &self,
        after: Option<(&str, &str)>,
        limit: usize,
    ) -> Result<WorkspacePage>;
    async fn put_workspace_operation(
        &self,
        workspace: &WorkspaceRecord,
        operation: &WorkspaceOperation,
        payload_json: &str,
        idempotency_key: &str,
        request_digest: &str,
    ) -> Result<WorkspaceOperation>;
    async fn compare_and_swap_workspace(
        &self,
        tenant_id: &str,
        id: &str,
        expected_version: u64,
        lifecycle: WorkspaceLifecycle,
        now_ms: u64,
    ) -> Result<WorkspaceRecord>;
    async fn claim_workspace_lease(
        &self,
        tenant_id: &str,
        id: &str,
        owner: &str,
        now_ms: u64,
        lease_ms: u64,
    ) -> Result<WorkspaceLease>;
    async fn release_workspace_lease(
        &self,
        tenant_id: &str,
        id: &str,
        owner: &str,
        fencing_token: u64,
        now_ms: u64,
    ) -> Result<()>;
    async fn fenced_workspace_update(
        &self,
        tenant_id: &str,
        id: &str,
        owner: &str,
        fencing_token: u64,
        lifecycle: WorkspaceLifecycle,
        now_ms: u64,
    ) -> Result<WorkspaceRecord>;
    async fn fenced_replace_workspace(
        &self,
        workspace: &WorkspaceRecord,
        owner: &str,
        fencing_token: u64,
        now_ms: u64,
    ) -> Result<WorkspaceRecord>;
    async fn put_operation(&self, operation: &WorkspaceOperation) -> Result<()>;
    async fn put_operation_with_input(
        &self,
        operation: &WorkspaceOperation,
        payload_json: &str,
    ) -> Result<()>;
    async fn put_idempotent_operation(
        &self,
        operation: &WorkspaceOperation,
        payload_json: &str,
        idempotency_key: &str,
        request_digest: &str,
    ) -> Result<WorkspaceOperation>;
    async fn get_operation(&self, operation_id: &str) -> Result<Option<WorkspaceOperation>>;
    async fn claim_operation(
        &self,
        operation_id: &str,
        owner: &str,
        now_ms: u64,
        lease_ms: u64,
    ) -> Result<WorkspaceOperation>;
    async fn renew_operation_lease(
        &self,
        operation_id: &str,
        owner: &str,
        fencing_token: u64,
        now_ms: u64,
        lease_ms: u64,
    ) -> Result<WorkspaceOperation>;
    async fn transition_operation(
        &self,
        transition: &OperationTransition,
    ) -> Result<WorkspaceOperation>;
    async fn request_operation_cancel(
        &self,
        tenant_id: &str,
        operation_id: &str,
        expected_version: u64,
        now_ms: u64,
    ) -> Result<WorkspaceOperation>;
    async fn put_operation_input(&self, operation_id: &str, payload_json: &str) -> Result<()>;
    async fn get_operation_input(&self, operation_id: &str) -> Result<Option<String>>;
    async fn list_claimable_operations(
        &self,
        tenant_id: &str,
        now_ms: u64,
        limit: usize,
    ) -> Result<Vec<WorkspaceOperation>>;
    async fn list_claimable_operations_global(
        &self,
        now_ms: u64,
        limit: usize,
    ) -> Result<Vec<WorkspaceOperation>>;
    async fn put_change_set(&self, value: &StoredChangeSet) -> Result<()>;
    async fn get_change_set(&self, tenant_id: &str, id: &str) -> Result<Option<StoredChangeSet>>;
    async fn put_preview(&self, value: &StoredPreview) -> Result<()>;
    async fn get_preview(&self, tenant_id: &str, id: &str) -> Result<Option<StoredPreview>>;
    async fn stop_preview(
        &self,
        tenant_id: &str,
        id: &str,
        expected_version: u64,
    ) -> Result<StoredPreview>;
    async fn retain_tenant(&self, tenant_id: &str, now_ms: u64, max_rows: usize) -> Result<usize>;
    async fn record_worker_heartbeat(&self, worker_id: &str, now_ms: u64) -> Result<()> {
        let _ = (worker_id, now_ms);
        Ok(())
    }
    async fn check_readiness(
        &self,
        worker_id: &str,
        now_ms: u64,
        max_worker_staleness_ms: u64,
        max_operation_backlog: usize,
    ) -> Result<()> {
        let _ = (
            worker_id,
            now_ms,
            max_worker_staleness_ms,
            max_operation_backlog,
        );
        Ok(())
    }
}

#[async_trait]
pub trait SecretResolver: Send + Sync + Debug {
    /// Secret bytes must remain inside the trusted adapter/executor boundary.
    async fn resolve(&self, tenant_id: &str, secret: &SecretRef) -> Result<Vec<u8>>;
}

#[async_trait]
pub trait WorkspaceProviderAdapter: Send + Sync + Debug {
    async fn provision(
        &self,
        tenant_id: &str,
        workspace_id: &str,
        config: &WorkspaceConfig,
        idempotency_key: &str,
    ) -> Result<ProvisionedWorkspace>;
    async fn reconnect(&self, durable_backend_id: &str) -> Result<()>;
    async fn suspend(&self, durable_backend_id: &str, idempotency_key: &str) -> Result<()>;
    async fn resume(&self, durable_backend_id: &str, idempotency_key: &str) -> Result<()>;
    async fn inspect(&self, workspace: &WorkspaceRecord) -> Result<WorkspaceResourceUsage>;
    async fn delete(&self, durable_backend_id: &str, idempotency_key: &str) -> Result<()>;
    async fn cancel_command(&self, durable_backend_id: &str, command_id: &str) -> Result<()>;
}

#[async_trait]
pub trait WorkspaceProviderRegistry: Send + Sync + Debug {
    async fn provider(&self, backend: &str) -> Result<Arc<dyn WorkspaceProviderAdapter>>;
}

#[async_trait]
pub trait WorkspaceCommandExecutor: Send + Sync + Debug {
    async fn execute(
        &self,
        tenant_id: &str,
        workspace_id: &str,
        command_id: &str,
        command: &CommandSpec,
    ) -> Result<CmdOutput>;
    async fn cancel(&self, tenant_id: &str, command_id: &str) -> Result<()>;
}