agent-workspace-contract 0.1.0

use super::*;
use sha2::{Digest, Sha256};

pub const WORKSPACE_UPLOADS_PATH: &str = "/internal/v1/workspaces/{workspaceId}/uploads";
pub const WORKSPACE_UPLOAD_PATH: &str = "/internal/v1/workspaces/{workspaceId}/uploads/{uploadId}";
pub const WORKSPACE_COMMANDS_PATH: &str = "/internal/v1/workspaces/{workspaceId}/commands";
pub const WORKSPACE_COMMAND_PATH: &str =
    "/internal/v1/workspaces/{workspaceId}/commands/{commandId}";
pub const WORKSPACE_COMMAND_OUTPUT_PATH: &str =
    "/internal/v1/workspaces/{workspaceId}/commands/{commandId}/output";
pub const WORKSPACE_COMMAND_CANCEL_PATH: &str =
    "/internal/v1/workspaces/{workspaceId}/commands/{commandId}/cancel";
pub const WORKSPACE_SNAPSHOTS_PATH: &str = "/internal/v1/workspaces/{workspaceId}/snapshots";
pub const WORKSPACE_SNAPSHOT_RESTORE_PATH: &str =
    "/internal/v1/workspaces/{workspaceId}/snapshots/{snapshotId}/restore";
pub const WORKSPACE_LEASE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/lease";
pub const WORKSPACE_CLONE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/clone";
pub const WORKSPACE_MIGRATE_PATH: &str = "/internal/v1/workspaces/{workspaceId}/migrate";
pub const WORKSPACE_OPERATIONS_PATH: &str = "/internal/v1/workspace-operations/{operationId}";

pub const WORKSPACE_MAX_UPLOAD_BYTES: u64 = WORKSPACE_MAX_FILE_BYTES as u64;
/// Resumable uploads deliberately use small chunks so request memory is
/// bounded independently from the maximum completed file size.
pub const WORKSPACE_UPLOAD_CHUNK_BYTES: usize = 1024 * 1024;
pub const WORKSPACE_UPLOAD_TTL_MS: u64 = 60 * 60 * 1000;
pub const WORKSPACE_PREVIEW_TTL_MS: u64 = 60 * 60 * 1000;
pub const WORKSPACE_MAX_OUTPUT_CHUNKS: usize = 64;
pub const WORKSPACE_OUTPUT_CHUNK_BYTES: usize = 32 * 1024;
pub const WORKSPACE_MAX_COMMAND_TIMEOUT_MS: u64 = 60_000;
pub const WORKSPACE_MAX_COMMAND_MEMORY_BYTES: u64 = 16 * 1024 * 1024 * 1024;
pub const WORKSPACE_MAX_COMMAND_CPU_MILLIS: u64 = 60_000;
pub const WORKSPACE_MAX_COMMAND_PROCESSES: u32 = 1024;
pub const WORKSPACE_MAX_COMMAND_DISK_BYTES: u64 = 64 * 1024 * 1024 * 1024;
pub const WORKSPACE_MAX_SNAPSHOT_BYTES: usize = 64 * 1024 * 1024;
pub const WORKSPACE_MAX_RETENTION_MS: u64 = 30 * 24 * 60 * 60 * 1000;

pub fn snapshot_digest(bytes: &[u8]) -> String {
    format!("sha256:{:x}", Sha256::digest(bytes))
}

pub const WORKSPACE_SNAPSHOT_ARCHIVE_MAGIC: &[u8; 7] = b"AWSNP1\0";

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UploadState {
    Open,
    Completed,
    Expired,
    Aborted,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UploadSession {
    pub id: String,
    pub tenant_id: String,
    pub workspace_id: String,
    pub path: String,
    pub offset: u64,
    pub max_bytes: u64,
    pub expires_at_ms: u64,
    pub state: UploadState,
    pub version: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub if_match: Option<String>,
    #[serde(default)]
    pub request_digest: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CreateUploadRequest {
    pub path: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub if_match: Option<String>,
    #[serde(default)]
    pub ttl_ms: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields, rename_all = "camelCase")]
pub struct CreatePreviewRequest {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    pub port: u16,
    #[serde(default)]
    pub ttl_ms: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UploadChunkRequest {
    pub offset: u64,
    pub content_base64: String,
    #[serde(default)]
    pub final_chunk: bool,
    pub version: u64,
}

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

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CommandResource {
    pub id: String,
    pub tenant_id: String,
    pub workspace_id: String,
    pub spec: CommandSpec,
    pub state: CommandState,
    pub operation_id: String,
    pub output_end_cursor: u64,
    pub output_truncated: bool,
    pub version: u64,
    pub created_at_ms: u64,
    pub updated_at_ms: u64,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub exit_code: Option<i32>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error_code: Option<String>,
}

#[derive(Debug, Clone)]
pub struct CommandTransition {
    pub tenant_id: String,
    pub id: String,
    pub expected_version: u64,
    pub state: CommandState,
    pub exit_code: Option<i32>,
    pub error_code: Option<String>,
    pub now_ms: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CreateCommandRequest {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    pub command: CommandSpec,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum OutputStream {
    Stdout,
    Stderr,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OutputChunk {
    pub command_id: String,
    pub cursor: u64,
    pub stream: OutputStream,
    pub content: String,
    pub bytes: usize,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct OutputPage {
    pub chunks: Vec<OutputChunk>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub next_cursor: Option<u64>,
    pub truncated: bool,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Snapshot {
    pub id: String,
    pub tenant_id: String,
    pub workspace_id: String,
    pub digest: String,
    pub source_version: u64,
    pub size_bytes: u64,
    pub file_count: usize,
    pub created_at_ms: u64,
    pub expires_at_ms: u64,
    #[serde(default)]
    pub request_digest: String,
    /// Durable object-store reference. Legacy rows can omit this while they
    /// are migrated from the former inline SQL blob representation.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub object_ref: Option<SnapshotObjectRef>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SnapshotObjectRef {
    pub key: String,
    pub digest: String,
    pub size_bytes: u64,
}

#[derive(Clone)]
pub struct StoredSnapshot {
    pub snapshot: Snapshot,
    /// Read-only compatibility payload for rows created before object storage
    /// was introduced. New writes must always leave this empty.
    pub legacy_blob: Option<Vec<u8>>,
}

impl Debug for StoredSnapshot {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("StoredSnapshot")
            .field("snapshot", &self.snapshot)
            .field(
                "legacy_blob",
                &self.legacy_blob.as_ref().map(|bytes| bytes.len()),
            )
            .finish()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CreateSnapshotRequest {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,
    #[serde(default)]
    pub retention_ms: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CloneWorkspaceRequest {
    pub destination_workspace_id: String,
    pub destination_project_id: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub snapshot_id: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MigrateWorkspaceRequest {
    pub destination_backend: String,
    #[serde(default)]
    pub config: WorkspaceConfig,
}

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

#[derive(Debug, Clone)]
pub struct WorkspaceLeaseKeepAlive {
    pub tenant_id: String,
    pub id: String,
    pub owner: String,
    pub fencing_token: u64,
    pub now_ms: u64,
    pub lease_ms: u64,
}

#[derive(Debug, Clone)]
pub struct CommandPermitRequest {
    pub command_id: String,
    pub tenant_id: String,
    pub owner: String,
    pub now_ms: u64,
    pub lease_ms: u64,
    pub global_limit: usize,
    pub tenant_limit: usize,
}

/// Bounded, random-access reader used across application and object-store
/// ports. Implementations must never return more than `max_bytes`.
#[async_trait]
pub trait SnapshotObjectReader: Send + Sync + Debug {
    fn len(&self) -> u64;
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
    fn digest(&self) -> &str;
    async fn read_chunk(&self, offset: u64, max_bytes: usize) -> Result<Vec<u8>>;
}

#[async_trait]
pub trait SnapshotObjectUpload: Send + Debug {
    async fn write_chunk(&mut self, bytes: &[u8]) -> Result<()>;
    async fn commit(self: Box<Self>, digest: &str, size_bytes: u64) -> Result<SnapshotObjectRef>;
    async fn abort(self: Box<Self>) -> Result<()>;
}

#[async_trait]
pub trait SnapshotObjectStore: Send + Sync + Debug {
    async fn begin_upload(&self, key: &str) -> Result<Box<dyn SnapshotObjectUpload>>;
    async fn open(&self, object: &SnapshotObjectRef) -> Result<Arc<dyn SnapshotObjectReader>>;
    async fn delete(&self, object: &SnapshotObjectRef) -> Result<()>;
    async fn check_readiness(&self) -> Result<()> {
        Ok(())
    }
}

#[async_trait]
pub trait SnapshotProvider: Send + Sync + Debug {
    async fn export_snapshot(
        &self,
        durable_backend_id: &str,
    ) -> Result<Arc<dyn SnapshotObjectReader>>;
    async fn import_snapshot(
        &self,
        workspace_id: &str,
        config: &WorkspaceConfig,
        snapshot: Arc<dyn SnapshotObjectReader>,
        idempotency_key: &str,
    ) -> Result<ProvisionedWorkspace>;
}

#[async_trait]
pub trait PreviewLifecycleProvider: Send + Sync + Debug {
    async fn create_preview(
        &self,
        durable_backend_id: &str,
        request_id: &str,
        port: u16,
        ttl_ms: u64,
    ) -> Result<StoredPreview>;
    async fn stop_preview(&self, provider_ref: &str) -> Result<()>;
}

#[derive(Clone)]
pub struct WorkspaceSession {
    pub workspace: Arc<dyn Workspace>,
    pub upload_writer: Arc<dyn WorkspaceUploadWriter>,
    pub executor: Arc<dyn WorkspaceCommandExecutor>,
    pub snapshot_provider: Option<Arc<dyn SnapshotProvider>>,
    pub snapshot_objects: Option<Arc<dyn SnapshotObjectStore>>,
    pub preview_provider: Option<Arc<dyn PreviewLifecycleProvider>>,
}

impl Debug for WorkspaceSession {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("WorkspaceSession")
            .finish_non_exhaustive()
    }
}

#[async_trait]
pub trait WorkspaceSessionResolver: Send + Sync + Debug {
    async fn resolve(&self, workspace: &WorkspaceRecord) -> Result<WorkspaceSession>;
    /// Resolve the importer owned by the requested destination backend. A
    /// migration must never reuse the source session's importer implicitly.
    async fn resolve_snapshot_destination(
        &self,
        backend: &str,
    ) -> Result<Arc<dyn SnapshotProvider>>;
    async fn check_readiness(&self) -> Result<()> {
        Ok(())
    }
}

#[async_trait]
pub trait WorkspaceAdvancedRepository: Send + Sync + Debug {
    async fn put_upload(&self, value: &UploadSession) -> Result<()>;
    async fn get_upload(&self, tenant_id: &str, id: &str) -> Result<Option<UploadSession>>;
    async fn advance_upload(
        &self,
        tenant_id: &str,
        id: &str,
        expected_offset: u64,
        expected_version: u64,
        new_offset: u64,
        state: UploadState,
    ) -> Result<UploadSession>;
    async fn put_command(&self, value: &CommandResource) -> Result<()>;
    async fn put_command_operation(
        &self,
        command: &CommandResource,
        operation: &WorkspaceOperation,
        payload_json: &str,
    ) -> Result<()>;
    async fn get_command(&self, tenant_id: &str, id: &str) -> Result<Option<CommandResource>>;
    async fn transition_command(&self, request: CommandTransition) -> Result<CommandResource>;
    async fn append_output(
        &self,
        tenant_id: &str,
        command_id: &str,
        chunks: &[OutputChunk],
        truncated: bool,
    ) -> Result<()>;
    async fn output_page(
        &self,
        tenant_id: &str,
        command_id: &str,
        after: u64,
        limit: usize,
    ) -> Result<OutputPage>;
    async fn put_snapshot(&self, value: &Snapshot, legacy_blob: Option<&[u8]>) -> Result<()>;
    async fn get_snapshot(&self, tenant_id: &str, id: &str) -> Result<Option<StoredSnapshot>>;
    async fn list_expired_snapshots(
        &self,
        tenant_id: &str,
        now_ms: u64,
        limit: usize,
    ) -> Result<Vec<Snapshot>>;
    async fn delete_expired_snapshot_manifest(
        &self,
        tenant_id: &str,
        id: &str,
        expected_expires_at_ms: u64,
    ) -> Result<bool>;
    /// Persist a failed object deletion so poison objects back off and cannot
    /// occupy every slot in subsequent bounded retention batches.
    async fn record_snapshot_cleanup_failure(
        &self,
        tenant_id: &str,
        id: &str,
        now_ms: u64,
        error: &str,
    ) -> Result<()> {
        let _ = (tenant_id, id, now_ms, error);
        Ok(())
    }
    async fn list_expired_previews(
        &self,
        tenant_id: &str,
        workspace_id: &str,
        now_ms: u64,
        limit: usize,
    ) -> Result<Vec<StoredPreview>>;
    async fn reconcile_expired(&self, tenant_id: &str, now_ms: u64) -> Result<usize>;
    async fn keep_alive_workspace(
        &self,
        request: WorkspaceLeaseKeepAlive,
    ) -> Result<WorkspaceLease>;
}

/// Aggregate metadata port used by the application layer. Concrete SQL/storage
/// technology belongs in adapter crates.
pub trait WorkspaceRepository: WorkspaceControlRepository + WorkspaceAdvancedRepository {}
impl<T> WorkspaceRepository for T where T: WorkspaceControlRepository + WorkspaceAdvancedRepository {}

#[async_trait]
pub trait WorkspaceUploadWriter: Send + Sync + Debug {
    async fn write_upload_chunk(&self, path: &str, offset: u64, bytes: &[u8]) -> Result<u64>;
}

#[async_trait]
pub trait CommandPermitRepository: Send + Sync + Debug {
    async fn acquire_command_permit(&self, request: CommandPermitRequest) -> Result<bool>;
    async fn release_command_permit(&self, command_id: &str, owner: &str) -> Result<()>;
    async fn cancel_command_permit(&self, command_id: &str) -> Result<bool>;
    async fn command_canceled(&self, command_id: &str) -> Result<bool>;
}