agent-infra-sdk 0.2.0

use crate::operation::{OperationHandle, OperationObservation, OperationPoller, OperationProgress};
use crate::transport::{CallOptions, HttpTransport, InfraClientError, ServiceEndpoint};
use agent_workspace_contract::{
    AckResponse, BytesResponse, CloneWorkspaceRequest, CommandResource,
    CreateCommandRequest, CreatePreviewRequest, CreateSnapshotRequest, CreateUploadRequest,
    KeepAliveRequest, MigrateWorkspaceRequest, OperationState, OutputPage,
    PutWorkspaceFileRequest, StoredPreview, UploadChunkRequest, UploadSession,
    WORKSPACE_COMMAND_CANCEL_PATH, WORKSPACE_COMMAND_OUTPUT_PATH, WORKSPACE_COMMAND_PATH,
    WORKSPACE_COMMANDS_PATH, WORKSPACE_LEASE_PATH, WORKSPACE_MIGRATE_PATH,
    WORKSPACE_OPERATION_CANCEL_PATH, WORKSPACE_OPERATIONS_PATH, WORKSPACE_PATH,
    WORKSPACE_RECONCILE_PATH, WORKSPACE_RESOURCE_FILE_PATH, WORKSPACE_RESOURCE_SEARCH_PATH,
    WORKSPACE_RESUME_PATH, WORKSPACE_SERVICE_NAME, WORKSPACE_SNAPSHOT_RESTORE_PATH,
    WORKSPACE_SNAPSHOTS_PATH, WORKSPACE_SUSPEND_PATH, WORKSPACE_UPLOAD_PATH, WORKSPACE_UPLOADS_PATH,
    WORKSPACE_USAGE_PATH, WorkspaceActionRequest, WorkspaceLease, WorkspaceOperation,
    WorkspaceRecord, WorkspaceResourceUsage, WorkspaceSearchRequest, WorkspaceSearchResponse,
};
use reqwest::Client;
use std::sync::Arc;

mod support;
use support::*;

/// Client for leftover managed Workspace HTTP. New callers should use
/// [`crate::EnvironmentClient`].
#[derive(Clone, Debug)]
pub struct WorkspaceClient {
    transport: HttpTransport,
}

impl WorkspaceClient {
    pub(crate) fn new_with_endpoint(
        http: Client,
        endpoint: ServiceEndpoint,
        options: crate::ClientOptions,
    ) -> Self {
        let endpoint = endpoint.with_default_credential_audience("workspace");
        Self {
            transport: HttpTransport::new_with_options(
                http,
                WORKSPACE_SERVICE_NAME,
                endpoint,
                options,
            ),
        }
    }

    pub async fn workspace(&self, workspace_id: &str) -> Result<WorkspaceRecord, InfraClientError> {
        self.workspace_with_options(workspace_id, CallOptions::default())
            .await
    }

    pub async fn workspace_with_options(
        &self,
        workspace_id: &str,
        options: CallOptions,
    ) -> Result<WorkspaceRecord, InfraClientError> {
        self.transport
            .get_json_with_options(&expand(WORKSPACE_PATH, &[workspace_id]), options)
            .await
    }

    pub async fn suspend_workspace(
        &self,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        self.workspace_action(WORKSPACE_SUSPEND_PATH, workspace_id, request, key)
            .await
    }

    pub async fn resume_workspace(
        &self,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        self.workspace_action(WORKSPACE_RESUME_PATH, workspace_id, request, key)
            .await
    }

    pub async fn reconcile_workspace(
        &self,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        self.workspace_action(WORKSPACE_RECONCILE_PATH, workspace_id, request, key)
            .await
    }

    pub async fn delete_workspace(
        &self,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation = self
            .transport
            .delete_json_with_options(
                &expand(WORKSPACE_PATH, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn workspace_usage(
        &self,
        workspace_id: &str,
    ) -> Result<WorkspaceResourceUsage, InfraClientError> {
        self.transport
            .get_json(&expand(WORKSPACE_USAGE_PATH, &[workspace_id]))
            .await
    }

    async fn workspace_action(
        &self,
        path: &str,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation = self
            .transport
            .post_json_with_options(
                &expand(path, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    /// Read a file from an explicit managed workspace. New consumers should
    /// prefer Environment Space file APIs.
    pub async fn read_workspace_file(
        &self,
        workspace_id: &str,
        path: &str,
    ) -> Result<BytesResponse, InfraClientError> {
        self.read_workspace_file_with_options(workspace_id, path, CallOptions::default())
            .await
    }

    pub async fn read_workspace_file_with_options(
        &self,
        workspace_id: &str,
        path: &str,
        options: CallOptions,
    ) -> Result<BytesResponse, InfraClientError> {
        self.transport
            .get_json_with_options(&resource_file_path(workspace_id, path), options)
            .await
    }

    /// Read a managed-workspace file without turning an ordinary miss into an
    /// error. This avoids the exists-then-read race and the second network
    /// round trip for callers that need create/update preconditions.
    pub async fn read_workspace_file_if_exists_with_options(
        &self,
        workspace_id: &str,
        path: &str,
        options: CallOptions,
    ) -> Result<Option<BytesResponse>, InfraClientError> {
        match self
            .read_workspace_file_with_options(workspace_id, path, options)
            .await
        {
            Ok(response) => Ok(Some(response)),
            Err(InfraClientError::HttpStatus { status: 404, .. }) => Ok(None),
            Err(error) => Err(error),
        }
    }

    pub async fn put_workspace_file(
        &self,
        workspace_id: &str,
        path: &str,
        request: &PutWorkspaceFileRequest,
        options: CallOptions,
    ) -> Result<Option<String>, InfraClientError> {
        let response: AckResponse = self
            .transport
            .put_json_with_options(&resource_file_path(workspace_id, path), request, options)
            .await?;
        ensure_ack(response)
    }

    /// Remove one file from an explicit managed workspace.
    pub async fn remove_workspace_file_with_options(
        &self,
        workspace_id: &str,
        path: &str,
        options: CallOptions,
    ) -> Result<(), InfraClientError> {
        let response: AckResponse = self
            .transport
            .delete_json_with_options(
                &resource_file_path(workspace_id, path),
                &serde_json::json!({}),
                options,
            )
            .await?;
        ensure_ack(response).map(|_| ())
    }

    pub async fn search_workspace_files(
        &self,
        workspace_id: &str,
        request: &WorkspaceSearchRequest,
    ) -> Result<WorkspaceSearchResponse, InfraClientError> {
        self.search_workspace_files_with_options(workspace_id, request, CallOptions::default())
            .await
    }

    pub async fn search_workspace_files_with_options(
        &self,
        workspace_id: &str,
        request: &WorkspaceSearchRequest,
        options: CallOptions,
    ) -> Result<WorkspaceSearchResponse, InfraClientError> {
        self.transport
            .post_json_with_options(
                &expand(WORKSPACE_RESOURCE_SEARCH_PATH, &[workspace_id]),
                request,
                options.idempotent(true),
            )
            .await
    }

    pub async fn next_workspace_search(
        &self,
        workspace_id: &str,
        request: &WorkspaceSearchRequest,
        current: &WorkspaceSearchResponse,
    ) -> Result<Option<WorkspaceSearchResponse>, InfraClientError> {
        let Some(cursor) = current.next_cursor.as_deref() else {
            return Ok(None);
        };
        let mut next = request.clone();
        next.cursor = Some(cursor.to_string());
        self.search_workspace_files(workspace_id, &next)
            .await
            .map(Some)
    }

    pub async fn create_upload(
        &self,
        workspace_id: &str,
        request: &CreateUploadRequest,
        idempotency_key: &str,
    ) -> Result<UploadSession, InfraClientError> {
        self.transport
            .post_json_with_options(
                &expand(WORKSPACE_UPLOADS_PATH, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await
    }

    pub async fn create_preview(
        &self,
        workspace_id: &str,
        request: &CreatePreviewRequest,
        idempotency_key: &str,
    ) -> Result<StoredPreview, InfraClientError> {
        self.transport
            .post_json_with_options(
                &format!(
                    "/internal/v1/workspaces/{}/previews",
                    encode_path_segment(workspace_id)
                ),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await
    }

    pub async fn upload(
        &self,
        workspace_id: &str,
        upload_id: &str,
    ) -> Result<UploadSession, InfraClientError> {
        self.transport
            .get_json(&expand(WORKSPACE_UPLOAD_PATH, &[workspace_id, upload_id]))
            .await
    }

    pub async fn upload_chunk(
        &self,
        workspace_id: &str,
        upload_id: &str,
        request: &UploadChunkRequest,
    ) -> Result<UploadSession, InfraClientError> {
        self.transport
            .put_json_with_options(
                &expand(WORKSPACE_UPLOAD_PATH, &[workspace_id, upload_id]),
                request,
                CallOptions::default(),
            )
            .await
    }

    pub async fn create_command(
        &self,
        workspace_id: &str,
        request: &CreateCommandRequest,
        idempotency_key: &str,
    ) -> Result<CommandResource, InfraClientError> {
        self.create_command_with_options(
            workspace_id,
            request,
            CallOptions::default().idempotency_key(idempotency_key),
        )
        .await
    }

    pub async fn create_command_with_options(
        &self,
        workspace_id: &str,
        request: &CreateCommandRequest,
        options: CallOptions,
    ) -> Result<CommandResource, InfraClientError> {
        self.transport
            .post_json_with_options(
                &expand(WORKSPACE_COMMANDS_PATH, &[workspace_id]),
                request,
                options,
            )
            .await
    }

    pub async fn command(
        &self,
        workspace_id: &str,
        command_id: &str,
    ) -> Result<CommandResource, InfraClientError> {
        self.command_with_options(workspace_id, command_id, CallOptions::default())
            .await
    }

    pub async fn command_with_options(
        &self,
        workspace_id: &str,
        command_id: &str,
        options: CallOptions,
    ) -> Result<CommandResource, InfraClientError> {
        self.transport
            .get_json_with_options(
                &expand(WORKSPACE_COMMAND_PATH, &[workspace_id, command_id]),
                options,
            )
            .await
    }

    pub async fn command_output(
        &self,
        workspace_id: &str,
        command_id: &str,
        after: u64,
        limit: Option<usize>,
    ) -> Result<OutputPage, InfraClientError> {
        self.command_output_with_options(
            workspace_id,
            command_id,
            after,
            limit,
            CallOptions::default(),
        )
        .await
    }

    pub async fn command_output_with_options(
        &self,
        workspace_id: &str,
        command_id: &str,
        after: u64,
        limit: Option<usize>,
        options: CallOptions,
    ) -> Result<OutputPage, InfraClientError> {
        let mut path = expand(WORKSPACE_COMMAND_OUTPUT_PATH, &[workspace_id, command_id]);
        path.push_str(&format!(
            "?after={after}&limit={}",
            limit.unwrap_or(64).min(64)
        ));
        self.transport.get_json_with_options(&path, options).await
    }

    pub async fn next_command_output(
        &self,
        workspace_id: &str,
        command_id: &str,
        current: &OutputPage,
        limit: Option<usize>,
    ) -> Result<Option<OutputPage>, InfraClientError> {
        let Some(cursor) = current.next_cursor else {
            return Ok(None);
        };
        self.command_output(workspace_id, command_id, cursor, limit)
            .await
            .map(Some)
    }

    pub async fn cancel_command(
        &self,
        workspace_id: &str,
        command_id: &str,
    ) -> Result<CommandResource, InfraClientError> {
        self.cancel_command_with_options(
            workspace_id,
            command_id,
            CallOptions::default()
                .idempotency_key(format!("cancel-command:{workspace_id}:{command_id}")),
        )
        .await
    }

    pub async fn cancel_command_with_options(
        &self,
        workspace_id: &str,
        command_id: &str,
        options: CallOptions,
    ) -> Result<CommandResource, InfraClientError> {
        self.transport
            .post_json_with_options(
                &expand(WORKSPACE_COMMAND_CANCEL_PATH, &[workspace_id, command_id]),
                &serde_json::json!({}),
                options,
            )
            .await
    }

    pub async fn create_snapshot(
        &self,
        workspace_id: &str,
        request: &CreateSnapshotRequest,
        idempotency_key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation: WorkspaceOperation = self
            .transport
            .post_json_with_options(
                &expand(WORKSPACE_SNAPSHOTS_PATH, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn restore_snapshot(
        &self,
        workspace_id: &str,
        snapshot_id: &str,
        idempotency_key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation: WorkspaceOperation = self
            .transport
            .post_json_with_options(
                &expand(
                    WORKSPACE_SNAPSHOT_RESTORE_PATH,
                    &[workspace_id, snapshot_id],
                ),
                &serde_json::json!({}),
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn clone_workspace(
        &self,
        workspace_id: &str,
        request: &CloneWorkspaceRequest,
        idempotency_key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation: WorkspaceOperation = self
            .transport
            .post_json_with_options(
                &expand(
                    agent_workspace_contract::WORKSPACE_CLONE_PATH,
                    &[workspace_id],
                ),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn migrate_workspace(
        &self,
        workspace_id: &str,
        request: &MigrateWorkspaceRequest,
        idempotency_key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation: WorkspaceOperation = self
            .transport
            .post_json_with_options(
                &expand(WORKSPACE_MIGRATE_PATH, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn keep_alive(
        &self,
        workspace_id: &str,
        request: &KeepAliveRequest,
    ) -> Result<WorkspaceLease, InfraClientError> {
        self.transport
            .put_json_with_options(
                &expand(WORKSPACE_LEASE_PATH, &[workspace_id]),
                request,
                CallOptions::default(),
            )
            .await
    }

    pub async fn operation(
        &self,
        operation_id: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation = WorkspaceOperationPoller {
            transport: self.transport.clone(),
        }
        .poll(operation_id)
        .await?;
        Ok(self.operation_handle(operation))
    }

    fn operation_handle(
        &self,
        operation: WorkspaceOperation,
    ) -> OperationHandle<WorkspaceOperation> {
        OperationHandle::new(
            operation.id.clone(),
            Some(operation),
            Arc::new(WorkspaceOperationPoller {
                transport: self.transport.clone(),
            }),
        )
    }
}