1988 lines · 79.1 KB · Rust
RawDownload
use std::borrow::Cow;use std::collections::{HashMap, HashSet};use std::sync::Arc;use std::time::Duration; use chrono::{DateTime, Utc};// Re-exported from cloud_objects.pub use cloud_objects::cloud_object::SerializedModel;use derivative::Derivative;use http::StatusCode;use lazy_static::lazy_static;use uuid::Uuid;use warp_errors::report_error;use warp_graphql::scalars::time::ServerTimestamp;use warpui::r#async::FutureId;use warpui::{Entity, ModelContext, RequestState, RetryOption, SingletonEntity}; use super::graphql::GraphQLError;use super::ids::{ClientId, HashableId, ObjectUid, ServerId, SyncId, ToServerId};use super::server_api::auth::UserAuthenticationError;use super::server_api::object::ObjectClient;use crate::ai::ambient_agents::scheduled::CloudScheduledAmbientAgentModel;use crate::ai::cloud_agent_config::CloudAgentConfigModel;use crate::ai::cloud_environments::CloudAmbientAgentEnvironmentModel;use crate::ai::execution_profiles::CloudAIExecutionProfileModel;use crate::ai::facts::CloudAIFactModel;use crate::ai::mcp::CloudMCPServerModel;use crate::ai::mcp::templatable::CloudTemplatableMCPServerModel;use crate::cloud_object::model::actions::{    ObjectAction, ObjectActionHistory, ObjectActionSubtype, ObjectActionType,};use crate::cloud_object::model::generic_string_model::GenericStringObjectId;use crate::cloud_object::{    BulkCreateCloudObjectResult, BulkCreateGenericStringObjectsRequest, CloudModelType,    CloudObject, CloudObjectEventEntrypoint, CreateCloudObjectResult, CreateObjectRequest,    GenericCloudObject, GenericStringObjectFormat, GenericStringObjectUniqueKey, JsonObjectType,    ObjectType, Owner, Revision, RevisionAndLastEditor, ServerCloudObject, ServerCreationInfo,    UpdateCloudObjectResult,};use crate::drive::CloudObjectTypeAndId;use crate::drive::folders::CloudFolderModel;use crate::env_vars::CloudEnvVarCollectionModel;use crate::notebooks::CloudNotebookModel;use crate::server::cloud_objects::update_manager::InitiatedBy;use crate::settings::cloud_preferences::CloudPreferenceModel;use crate::workflows::CloudWorkflowModel;use crate::workflows::workflow_enum::CloudWorkflowEnumModel; lazy_static! {    static ref DEFAULT_RETRY_OPTION: RetryOption =        RetryOption::exponential(Duration::from_secs(1), 2., 3);} /// Type of successful update response from server.#[allow(clippy::large_enum_variant)]enum ResponseType {    Creation {        creation_result: CreationResponseType,    },    Update {        update_result: UpdateResponseType,    },    ObjectAction {        action_timestamp: DateTime<Utc>,        action_history: ObjectActionHistory,    },} /// Allow large enum variant because success is the most common by far#[allow(clippy::large_enum_variant)]enum CreationResponseType {    Success {        client_id: ClientId,        revision_and_editor: RevisionAndLastEditor,        metadata_ts: ServerTimestamp,        server_creation_info: ServerCreationInfo,    },    UserFacingError {        message: String,        client_id: ClientId,    },} enum UpdateResponseType {    Success {        revision_and_editor: RevisionAndLastEditor,    },    /// The update was rejected because the update was not sent from the current revision in    /// storage. The object and revision in storage are returned.    Rejected { object: Box<ServerCloudObject> },} #[derive(Debug, PartialEq, Clone)]pub struct GenericStringObjectToCreate {    pub id: ClientId,    pub format: GenericStringObjectFormat,    pub serialized_model: Arc<SerializedModel>,    pub initial_folder_id: Option<SyncId>,    pub entrypoint: CloudObjectEventEntrypoint,    pub uniqueness_key: Option<GenericStringObjectUniqueKey>,    pub initiated_by: InitiatedBy,} /// An ID for a `QueueItem` in the sync queue.#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]pub struct QueueItemId(Uuid); impl QueueItemId {    pub fn new() -> Self {        Self(Uuid::new_v4())    }} #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]enum QueueDependency {    QueueItem(QueueItemId),    BulkCreateGenericStringObject {        queue_item_id: QueueItemId,        client_id: ClientId,    },} /// Item that sync queue supports.#[derive(Derivative, Debug)]#[derivative(PartialEq, Eq, Clone)]pub enum QueueItem {    CreateObject {        object_type: ObjectType,        owner: Owner,        id: ClientId,        title: Option<Arc<String>>,        serialized_model: Option<Arc<SerializedModel>>,        initial_folder_id: Option<SyncId>,        entrypoint: CloudObjectEventEntrypoint,        initiated_by: InitiatedBy,    },    // Separate CreateWorkflow type that should be removed when we do the SyncId refactor.    // Stores a reference to a CloudWorkflowModel, rather than a SerializedModel, which is needed    // for updating an enqueued workflow with the server IDs of enums it references as they are created.    CreateWorkflow {        object_type: ObjectType,        owner: Owner,        id: ClientId,        #[derivative(PartialEq = "ignore")]        model: Arc<CloudWorkflowModel>,        initial_folder_id: Option<SyncId>,        entrypoint: CloudObjectEventEntrypoint,        initiated_by: InitiatedBy,    },    BulkCreateGenericStringObjects {        owner: Owner,        objects: Vec<GenericStringObjectToCreate>,    },    // Note, we continue to do UpdateXXX items per object type here    // because it's the most type safe way to handle the different id types    // and the different update payloads.  Most of the logic is still shared    // via a single update_object method.    UpdateNotebook {        model: Arc<CloudNotebookModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateWorkflow {        model: Arc<CloudWorkflowModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateFolder {        id: SyncId,        model: Arc<CloudFolderModel>,    },    UpdateCloudPreferences {        model: Arc<CloudPreferenceModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateEnvVarCollection {        model: Arc<CloudEnvVarCollectionModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateWorkflowEnum {        model: Arc<CloudWorkflowEnumModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateAIFact {        model: Arc<CloudAIFactModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateMCPServer {        model: Arc<CloudMCPServerModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateAIExecutionProfile {        model: Arc<CloudAIExecutionProfileModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateTemplatableMCPServer {        model: Arc<CloudTemplatableMCPServerModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateCloudEnvironment {        model: Arc<CloudAmbientAgentEnvironmentModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateScheduledAmbientAgent {        model: Arc<CloudScheduledAmbientAgentModel>,        id: SyncId,        revision: Option<Revision>,    },    UpdateCloudAgentConfig {        model: Arc<CloudAgentConfigModel>,        id: SyncId,        revision: Option<Revision>,    },    RecordObjectAction {        id_and_type: CloudObjectTypeAndId,        action_type: ObjectActionType,        action_timestamp: DateTime<Utc>,        data: Option<String>,    },} impl QueueItem {    pub fn from_cached_objects(        objects: impl Iterator<Item = Box<dyn CloudObject>>,    ) -> Vec<QueueItem> {        objects            .map(|object| {                match object.create_object_queue_item(                    CloudObjectEventEntrypoint::default(),                    // InitiatedBy::User was added as a default value since we do not save the initiated_by values in the Sqlite cache.                    // InitiatedBy::User is a safer default option because it will show toasts.                    // In the future, if System events are common, we may want to save the initiated_by field in Sqlite.                    InitiatedBy::User,                ) {                    Some(create_object_queue_item) => create_object_queue_item,                    _ => object.update_object_queue_item(None),                }            })            .collect::<Vec<_>>()    }     // Converts a list of pending actions into SyncQueue items that can be sent to the server.    pub fn from_unsynced_actions(        actions: impl Iterator<Item = (CloudObjectTypeAndId, ObjectAction)>,    ) -> Vec<QueueItem> {        actions            .filter_map(|(id_and_type, action)| match action.action_subtype {                ObjectActionSubtype::SingleAction {                    timestamp,                    data,                    pending: true,                    ..                } => Some(QueueItem::RecordObjectAction {                    id_and_type,                    action_type: action.action_type,                    action_timestamp: timestamp,                    data,                }),                _ => None,            })            .collect::<Vec<_>>()    }} #[derive(Derivative, Clone, Debug)]#[derivative(PartialEq, Eq)]#[allow(clippy::enum_variant_names)]pub enum CreationFailureReason {    UniqueKeyConflict {        id: String,        initiated_by: InitiatedBy,    },    Denied {        message: String,        client_id: ClientId,        initiated_by: InitiatedBy,    },    Other {        id: String,        initiated_by: InitiatedBy,    },} #[derive(Derivative, Clone, Debug)]#[derivative(PartialEq, Eq)]#[allow(clippy::enum_variant_names)]#[allow(clippy::large_enum_variant)]pub enum SyncQueueEvent {    /// Emitted when we receive a successful response on creating a shareable object.    ObjectCreationSuccessful {        server_creation_info: ServerCreationInfo,        /// The client id of the object that was successfully created.        client_id: ClientId,        /// The revision of the object after creation.        revision_and_editor: RevisionAndLastEditor,        /// The timestamp of the object metadata after creation.        metadata_ts: ServerTimestamp,        /// Whether the creation was initiated by the user or the system.        initiated_by: InitiatedBy,    },    /// Emitted when we receive a successful response on updating a shareable object.    ObjectUpdateSuccessful {        /// The server id of the object        server_id: ServerId,        /// The revision of the object after updating.        revision_and_editor: RevisionAndLastEditor,    },    ObjectUpdateRejected {        id: String,        #[derivative(PartialEq = "ignore")]        object: Arc<ServerCloudObject>,    },    #[allow(dead_code)]    ObjectUpdateFeatureNotAvailable { id: String },    /// Request to server for creating a queue item has failed.    ObjectCreationFailure { reason: CreationFailureReason },    /// Request to server for updating a queue item has failed.    ObjectUpdateFailure { id: SyncId },    ReportObjectActionFailed {        uid: ObjectUid,        action_timestamp: DateTime<Utc>,    },    ReportObjectActionSucceeded {        uid: ObjectUid,        action_timestamp: DateTime<Utc>,        action_history: ObjectActionHistory,    },} /// Central mechanism for syncing shareable objects to the server. All CRUD actions/// on team objects should go through here to make sure we update the server in correct/// serial order and handle errors with exponential backoff.pub struct SyncQueue {    queue: Vec<(QueueItemId, QueueItem)>,    waiting_response: HashMap<String, HashSet<QueueItemId>>,    in_flight_bulk_create_objects: HashMap<QueueItemId, HashSet<ClientId>>,    object_client: Arc<dyn ObjectClient>,    client_id_to_server: HashMap<ClientId, String>,    server_id_to_client_hash: HashMap<String, String>,    should_dequeue: bool,    // Vec of futures that have been spawned after dequeueing items. Used within tests to ensure    // calls to `ctx#spawn` have finished before asserting.    spawned_futures: Vec<FutureId>,     // For each QueueItem, store the set of queue dependencies it depends upon finishing first.    queue_dependencies: HashMap<QueueItemId, HashSet<QueueDependency>>,} impl SyncQueue {    #[cfg(any(test, all(feature = "tui", feature = "test-util")))]    pub fn mock(ctx: &mut ModelContext<Self>) -> Self {        use super::server_api::ServerApiProvider;         Self::new(            Default::default(),            ServerApiProvider::new_for_test().get(),            ctx,        )    }     pub fn new(        queue_items: Vec<QueueItem>,        object_client: Arc<dyn ObjectClient>,        ctx: &mut ModelContext<Self>,    ) -> Self {        let mut sync_queue = Self {            queue: queue_items                .into_iter()                .map(|queue_item| (QueueItemId::new(), queue_item))                .collect(),            waiting_response: Default::default(),            in_flight_bulk_create_objects: Default::default(),            object_client,            client_id_to_server: Default::default(),            server_id_to_client_hash: Default::default(),            should_dequeue: false,            spawned_futures: vec![],            queue_dependencies: Default::default(),        };         sync_queue.initialize_queue_dependencies(ctx);        sync_queue    }     pub fn is_dequeueing(&self) -> bool {        self.should_dequeue    }     pub fn stop_dequeueing(&mut self) {        self.should_dequeue = false    }     pub fn start_dequeueing(&mut self, ctx: &mut ModelContext<Self>) {        self.should_dequeue = true;        self.dequeue(ctx)    }     // Clear the SyncQueue. This is used during Logout.    pub fn clear(&mut self) {        self.queue.clear();        self.waiting_response.clear();        self.in_flight_bulk_create_objects.clear();        self.queue_dependencies.clear();        self.client_id_to_server.clear();        self.server_id_to_client_hash.clear();    }     // Remove the (id, item) pair that matches `remove_id` from the SyncQueue    fn remove_id_from_queue(&mut self, remove_id: &QueueItemId) {        self.queue.retain(|(id, _)| id != remove_id);    }     /// Enqueue a new request.    pub fn enqueue(&mut self, item: QueueItem, ctx: &mut ModelContext<Self>) -> QueueItemId {        let queue_id = QueueItemId::new();        let mut queue_item = item;         self.add_inferred_dependencies(&mut queue_item, &queue_id, ctx);        self.queue.push((queue_id, queue_item));        self.dequeue(ctx);        queue_id    }     /// Given a queue item, infer its dependencies based on the current state of the queue and add them to queue_dependencies.    pub fn add_inferred_dependencies(        &mut self,        item: &mut QueueItem,        item_id: &QueueItemId,        ctx: &mut ModelContext<Self>,    ) {        let mut dependencies = match item {            // Update requests will depend on any existing create/updates to the same object            QueueItem::UpdateNotebook { id, .. }            | QueueItem::UpdateFolder { id, .. }            | QueueItem::UpdateCloudPreferences { id, .. }            | QueueItem::UpdateEnvVarCollection { id, .. }            | QueueItem::UpdateWorkflowEnum { id, .. }            | QueueItem::UpdateAIFact { id, .. }            | QueueItem::UpdateMCPServer { id, .. }            | QueueItem::UpdateAIExecutionProfile { id, .. }            | QueueItem::UpdateTemplatableMCPServer { id, .. }            | QueueItem::UpdateCloudEnvironment { id, .. }            | QueueItem::UpdateScheduledAmbientAgent { id, .. }            | QueueItem::UpdateCloudAgentConfig { id, .. } => self.get_update_dependencies(id),             // Update workflow requests should depend on existing requests to that object, as well as            // any enums or env vars they reference.            QueueItem::UpdateWorkflow {                id: workflow_id,                model,                ..            } => {                let mut dependencies = self.get_update_dependencies(workflow_id);                let enum_dependencies = self.get_workflow_object_dependencies(                    Arc::make_mut(model),                    *workflow_id,                    item_id,                    ctx,                );                match enum_dependencies {                    Ok(deps) => dependencies.extend(deps),                    Err(_) => self.handle_update_failure_response(*workflow_id, *item_id, ctx),                }                 dependencies            }             // Workflow creation requests should depend on every enum they reference.            // We should never dequeue a Workflow that references ClientIds rather than ServerIds.            QueueItem::CreateWorkflow {                model,                id: workflow_id,                initiated_by,                ..            } => {                let mut dependencies = HashSet::new();                let enum_dependencies = self.get_workflow_object_dependencies(                    Arc::make_mut(model),                    SyncId::ClientId(*workflow_id),                    item_id,                    ctx,                );                match enum_dependencies {                    Ok(deps) => dependencies.extend(deps),                    Err(_) => self.handle_creation_failure_response(                        workflow_id.to_string(),                        *item_id,                        *initiated_by,                        ctx,                    ),                } 

Lines 1–500 of 1988.