4833 lines · 198.5 KB · Rust
RawDownload
use std::collections::HashSet;use std::future::Future;use std::sync::Arc;use std::sync::mpsc::SyncSender;use std::time::Duration; use chrono::{DateTime, Utc};#[cfg(test)]pub use cloud_object_client::GetCloudObjectResponse;pub use cloud_object_client::InitialLoadResponse;use futures::channel::oneshot::{self, Receiver};use futures::stream::AbortHandle;use itertools::Itertools;use lazy_static::lazy_static;use regex::Regex;use warp_core::features::FeatureFlag;use warp_errors::report_error;use warp_graphql::mcp_gallery_template::MCPGalleryTemplate;use warp_graphql::object_permissions::AccessLevel;use warp_graphql::scalars::time::ServerTimestamp;use warp_util::sync::Condition;use warpui::r#async::{FutureId, Timer};use warpui::{    AppContext, Entity, ModelContext, ModelHandle, RequestState, RetryOption, SingletonEntity,    duration_with_jitter,}; use super::listener::ObjectUpdateMessage;use crate::ai::agent::conversation::AIConversationId;use crate::ai::ambient_agents::AmbientAgentTaskId;use crate::ai::ambient_agents::scheduled::{    CloudScheduledAmbientAgentModel, ScheduledAmbientAgent,};use crate::ai::blocklist::BlocklistAIHistoryModel;use crate::ai::cloud_environments::{AmbientAgentEnvironment, CloudAmbientAgentEnvironmentModel};use crate::ai::execution_profiles::profiles::AIExecutionProfilesModel;use crate::ai::execution_profiles::{AIExecutionProfile, CloudAIExecutionProfileModel};use crate::ai::facts::{AIFact, CloudAIFactModel};#[cfg(not(target_family = "wasm"))]use crate::ai::mcp::templatable::{CloudTemplatableMCPServerModel, TemplatableMCPServer};use crate::auth::AuthStateProvider;use crate::auth::auth_manager::AuthManager;use crate::cloud_object::model::actions::{    ObjectAction, ObjectActionHistory, ObjectActionType, ObjectActions,};use crate::cloud_object::model::generic_string_model::{    GenericStringModel, GenericStringObjectId, Serializer, StringModel,};use crate::cloud_object::model::persistence::{CloudModel, CloudModelEvent, UpdateSource};use crate::cloud_object::model::view::{CloudViewModel, Editor, EditorState};use crate::cloud_object::{    CloudLinkSharing, CloudModelType, CloudObject, CloudObjectEventEntrypoint, CloudObjectLocation,    CloudObjectSyncStatus, CreateCloudObjectResult, CreateObjectRequest, GenericCloudObject,    GenericServerObject, GenericStringObjectFormat, JsonObjectType, NumInFlightRequests,    ObjectDeleteResult, ObjectIdType, ObjectMetadataUpdateResult, ObjectPermissionsUpdateData,    ObjectType, Owner, Revision, RevisionAndLastEditor, ServerAIExecutionProfile, ServerAIFact,    ServerAmbientAgentEnvironment, ServerCloudAgentConfig, ServerCloudObject,    ServerEnvVarCollection, ServerMCPServer, ServerMetadata, ServerPermissions, ServerPreference,    ServerScheduledAmbientAgent, ServerTemplatableMCPServer, ServerWorkflowEnum, Space,    UpdateCloudObjectResult,};use crate::drive::CloudObjectTypeAndId;use crate::drive::drive_helpers::{    is_feature_gated_anonymous_user_past_env_var_limit,    is_feature_gated_anonymous_user_past_notebook_limit,    is_feature_gated_anonymous_user_past_workflow_limit,};use crate::drive::folders::{CloudFolderModel, FolderId};use crate::drive::sharing::SharingAccessLevel;use crate::env_vars::{CloudEnvVarCollectionModel, EnvVarCollection};use crate::network::{NetworkStatus, NetworkStatusEvent, NetworkStatusKind};use crate::notebooks::{CloudNotebookModel, NotebookId};use crate::persistence::ModelEvent;use crate::server::ids::{    ClientId, HashableId, HashedSqliteId, ObjectUid, ServerId, SyncId, ToServerId,    parse_sqlite_id_to_uid,};use crate::server::retry_strategies::{    OUT_OF_BAND_REQUEST_RETRY_STRATEGY, PERIODIC_POLL, PERIODIC_POLL_RETRY_STRATEGY,};use crate::server::server_api::object::{GuestIdentifier, ObjectClient};use crate::server::sync_queue::{    CreationFailureReason, GenericStringObjectToCreate, QueueItem, SyncQueue, SyncQueueEvent,};use crate::settings::cloud_preferences::Preference;use crate::workflows::workflow::Workflow;use crate::workflows::workflow_enum::{CloudWorkflowEnum, CloudWorkflowEnumModel, WorkflowEnum};use crate::workflows::{CloudWorkflowModel, WorkflowId};use crate::workspaces::team_tester::{TeamTesterStatus, TeamTesterStatusEvent};use crate::workspaces::update_manager::TeamUpdateManager;use crate::workspaces::user_profiles::{UserProfileWithUID, UserProfiles};use crate::workspaces::user_workspaces::UserWorkspaces; lazy_static! {    /// For online-only operations, we want to quickly determine if the operation can succeed,    /// so that if it can't, we can put the user back into the known good state.    /// So we try 3 times to prevent any transient failures.    static ref ONLINE_ONLY_OPERATION_RETRY_STRATEGY: RetryOption =        RetryOption::exponential(Duration::from_millis(500) /* interval */, 2. /* exponential factor */, 3 /* max retry count */);     static ref DUPLICATE_OBJECT_NAME_REGEX: Regex = Regex::new(r" \((\d+)\)$").expect("regex should not fail to compile"); } #[derive(Debug, PartialEq)]pub enum OperationSuccessType {    Success,    Failure,    Rejection,    Denied(String),    FeatureNotAvailable,} #[derive(Debug, PartialEq)]pub enum ObjectOperation {    Create { initiated_by: InitiatedBy },    Update,    MoveToFolder,    MoveToDrive,    Trash,    TakeEditAccess,    Untrash,    Delete { initiated_by: InitiatedBy },    EmptyTrash,    UpdatePermissions,    Leave,} #[derive(Debug)]pub struct ObjectOperationResult {    pub success_type: OperationSuccessType,    pub operation: ObjectOperation,    pub client_id: Option<ClientId>,    pub server_id: Option<ServerId>,    pub num_objects: Option<i32>, // counts number of objects (including descendants) deleted for permadeletion} #[derive(Debug)]pub enum UpdateManagerEvent {    ObjectOperationComplete {        result: ObjectOperationResult,    },    CloudPreferencesUpdated {        updated: Vec<Preference>,    },    MCPGalleryUpdated {        templates: Vec<MCPGalleryTemplate>,    },    AmbientTaskUpdated {        task_id: AmbientAgentTaskId,        timestamp: DateTime<Utc>,    },} /// An enum for choosing the behavior of the fetch_single_cloud_object function.pub enum FetchSingleObjectOption {    /// Perform the normal upsert behavior.    None,    /// Perform the normal upsert behavior, but additionally force overwrite the    /// in-memory object to whatever the server object is.    ForceOverwrite,    /// Only perform the normal upsert behavior if the object doesn't already    /// exist in-memory.    IgnoreIfExists,} /// An enum that defines whether the action was initiated by the user or the system.#[derive(Debug, Clone, Copy, PartialEq, Eq)]pub enum InitiatedBy {    User,    System,} #[derive(Debug)]pub struct GenericStringObjectInput<T, S>where    T: StringModel<            CloudObjectType = GenericCloudObject<GenericStringObjectId, GenericStringModel<T, S>>,        > + 'static,    S: Serializer<T> + 'static,{    pub id: ClientId,    pub model: GenericStringModel<T, S>,    pub initial_folder_id: Option<SyncId>,    pub entrypoint: CloudObjectEventEntrypoint,} /// The UpdateManager is responsible for delegating work/// when there is an update to an object (e.g. via a user interaction or/// a message from the server). Specifically, it will/// - write to SQLite/// - interact with the CloudModel to update the in-memory state used by the object views/// - interact with the SyncQueue by enqueueing an eventpub struct UpdateManager {    model_event_sender: Option<SyncSender<ModelEvent>>,    object_client: Arc<dyn ObjectClient>,    next_poll_abort_handle: Option<AbortHandle>,    in_flight_request_abort_handle: Option<AbortHandle>,    should_poll_for_updated_objects: bool,    spawned_futures: Vec<FutureId>,    has_initial_load: Condition,} impl UpdateManager {    pub fn new(        model_event_sender: Option<SyncSender<ModelEvent>>,        object_client: Arc<dyn ObjectClient>,        ctx: &mut ModelContext<Self>,    ) -> Self {        let network_status = NetworkStatus::handle(ctx);        ctx.subscribe_to_model(&network_status, |me, _, event, ctx| {            me.handle_network_status_changed(event, ctx);        });         let team_tester_status = TeamTesterStatus::handle(ctx);        ctx.subscribe_to_model(&team_tester_status, Self::handle_team_tester_status_changed);         let sync_queue = SyncQueue::handle(ctx);        ctx.subscribe_to_model(&sync_queue, |me, _, event, ctx| {            me.handle_model_event(event, ctx);        });         Self {            model_event_sender,            object_client,            next_poll_abort_handle: None,            in_flight_request_abort_handle: None,            should_poll_for_updated_objects: false,            spawned_futures: Default::default(),            has_initial_load: Condition::new(),        }    }     #[cfg(test)]    pub fn mock(ctx: &mut ModelContext<Self>) -> Self {        use crate::server::server_api::ServerApiProvider;         Self::new(            None,            ServerApiProvider::new_for_test().get_cloud_objects_client(),            ctx,        )    }     /// Simulate the initial load of changed objects completing.    #[cfg(test)]    pub fn mock_initial_load(        &mut self,        response: InitialLoadResponse,        ctx: &mut ModelContext<Self>,    ) {        self.on_changed_objects_fetched(response, false /* force_refresh */, ctx);    }     #[cfg(any(test, feature = "integration_tests"))]    pub fn spawned_futures(&self) -> &[FutureId] {        &self.spawned_futures    }     fn save_to_db(&self, events: impl IntoIterator<Item = ModelEvent>) {        let model_event_sender = self.model_event_sender.clone();        if let Some(model_event_sender) = &model_event_sender {            for event in events {                if let Err(e) = model_event_sender.send(event) {                    report_error!(anyhow::Error::new(e).context("Error saving to database"));                }            }        }    }     /// Remove team-owned objects in response to leaving a team.    pub fn remove_team_objects(&mut self, left_team_uid: ServerId, ctx: &mut ModelContext<Self>) {        let cloud_model = CloudModel::handle(ctx);        let objects_to_remove = cloud_model            .as_ref(ctx)            .all_cloud_objects_in_space(                Space::Team {                    team_uid: left_team_uid,                },                ctx,            )            .map(|object| object.cloud_object_type_and_id())            .collect_vec();         // First, delete in-memory from CloudModel and object actions.        cloud_model.update(ctx, |cloud_model, ctx| {            for object in objects_to_remove.iter() {                cloud_model.delete_object(object.sync_id(), ctx);            }        });        ObjectActions::handle(ctx).update(ctx, |object_actions, ctx| {            for object in objects_to_remove.iter() {                object_actions.delete_actions_for_object(&object.uid(), ctx);            }        });         // Then, delete from SQLite.        let object_ids_and_types = objects_to_remove            .into_iter()            .map(|object| (object.sync_id(), object.object_id_type()))            .collect();        self.save_to_db([ModelEvent::DeleteObjects {            ids: object_ids_and_types,        }]);    }     fn handle_team_tester_status_changed(        &mut self,        _: ModelHandle<TeamTesterStatus>,        event: &TeamTesterStatusEvent,        ctx: &mut ModelContext<Self>,    ) {        let TeamTesterStatusEvent::InitiateDataPollers { force_refresh } = event;        if *force_refresh {            self.refresh_updated_objects(ctx);        }         self.start_polling_for_updated_objects(ctx);    }     fn handle_model_event(&mut self, event: &SyncQueueEvent, ctx: &mut ModelContext<Self>) {        match event {            SyncQueueEvent::ObjectCreationSuccessful {                server_creation_info,                client_id,                revision_and_editor,                metadata_ts,                initiated_by,            } => {                let server_id = &server_creation_info.server_id_and_type.id;                 // Update server ID in sqlite.                self.save_to_db([ModelEvent::UpdateObjectAfterServerCreation {                    client_id: client_id.sqlite_hash(),                    server_creation_info: server_creation_info.clone(),                }]);                 // Update in-memory model.                CloudModel::handle(ctx).update(ctx, |cloud_model, ctx| {                    cloud_model.update_object_after_server_creation(                        *client_id,                        server_creation_info.clone(),                        ctx,                    );                    if let Some(object) = cloud_model.get_mut_by_uid(&server_id.uid()) {                        let is_no_longer_in_flight = {                            let status_if_no_reqs = CloudObjectSyncStatus::NoLocalChanges;                            object.decrement_in_flight_request_count(status_if_no_reqs)                        };                         if is_no_longer_in_flight {                            // Update sync status in sqlite.                             self.save_to_db([ModelEvent::MarkObjectAsSynced {                                hashed_sqlite_id: server_creation_info                                    .server_id_and_type                                    .sqlite_type_and_uid_hash(),                                revision_and_editor: revision_and_editor.clone(),                                metadata_ts: Some(*metadata_ts),                            }]);                        }                         ctx.notify();                    }                     cloud_model.set_latest_revision_and_editor(                        &server_id.uid(),                        revision_and_editor.clone(),                        ctx,                    );                     // When an object is created and we get a successful server response, part of marking the object as synced is accepting the                    // canonical metadata_ts.                    cloud_model.update_object_metadata_last_updated_ts(                        &server_id.uid(),                        *metadata_ts,                        ctx,                    );                     // If we have created a GSO, we need to update the in-memory model for any dependent workflows.                    // Go through every workflow and try to replace the client ID with the new server ID.                    if server_creation_info.server_id_and_type.id_type                        == ObjectIdType::GenericStringObject                    {                        let client_id = SyncId::ClientId(*client_id);                        let server_id = SyncId::ServerId(*server_id);                         if cloud_model.get_workflow_enum(&server_id).is_some() {                            cloud_model                                .get_all_active_and_inactive_workflows_mut()                                .for_each(|workflow_object| {                                    let mut workflow = workflow_object.model().clone();                                    let updated_model =                                        workflow.data.replace_object_id(client_id, server_id);                                     // If we changed anything, then update the in-memory model, emit a CloudEvent, and update the DB                                    if updated_model {                                        workflow_object.set_model(workflow);                                         ctx.emit(CloudModelEvent::ObjectUpdated {                                            type_and_id: workflow_object.cloud_object_type_and_id(),                                            source: UpdateSource::Local,                                        });                                         self.save_to_db([workflow_object.upsert_event()]);                                    }                                });                        } else if cloud_model.get_ai_execution_profile(&server_id).is_some() {                            AIExecutionProfilesModel::handle(ctx).update(ctx, |model, ctx| {                                model.replace_client_id_with_server_id(server_id, client_id, ctx);                            });                        }                    }                });                 // Delete the actions on the client ID. Once we get a server ID for an object, we start dequeuing any pending object actions and those                // directly populate the ObjectActions model with the server ID, so we don't need to worry about any conversion or anything like that.                ObjectActions::handle(ctx).update(ctx, |object_actions, ctx| {                    object_actions.delete_actions_for_object(&client_id.to_string(), ctx);                });                self.sync_actions_for_objects_to_sqlite(vec![&client_id.to_string()], ctx);                 ctx.emit(UpdateManagerEvent::ObjectOperationComplete {                    result: ObjectOperationResult {                        success_type: OperationSuccessType::Success,                        operation: ObjectOperation::Create {                            initiated_by: *initiated_by,                        },                        client_id: Some(*client_id),                        server_id: Some(*server_id),                        num_objects: None,                    },                });            }            SyncQueueEvent::ObjectUpdateSuccessful {                server_id,                revision_and_editor,            } => {                CloudModel::handle(ctx).update(ctx, |cloud_model, ctx| {                    // Update the object's revision to the latest one from the server                    cloud_model.set_latest_revision_and_editor(                        &server_id.uid(),                        revision_and_editor.clone(),                        ctx,                    );                    // After we update the revision, check if we can now clear the conflicting object                    cloud_model.check_and_maybe_clear_current_conflict(&server_id.uid(), ctx);                     // Decrement the object's request count and save it to sqlite if it's sync'd                    if let Some(object) = cloud_model.get_mut_by_uid(&server_id.uid()) {                        let is_no_longer_in_flight = {                            object.decrement_in_flight_request_count(                                CloudObjectSyncStatus::NoLocalChanges,                            )                        };                         if is_no_longer_in_flight {                            self.save_to_db([ModelEvent::MarkObjectAsSynced {                                hashed_sqlite_id: server_id                                    .sqlite_type_and_uid_hash(object.object_type().into()),                                revision_and_editor: revision_and_editor.clone(),                                metadata_ts: None,                            }]);                        }                         ctx.notify();                    }                });                 ctx.emit(UpdateManagerEvent::ObjectOperationComplete {                    result: ObjectOperationResult {                        success_type: OperationSuccessType::Success,                        operation: ObjectOperation::Update,                        client_id: None,                        server_id: Some(*server_id),                        num_objects: None,                    },                });            }            SyncQueueEvent::ObjectCreationFailure {                reason: CreationFailureReason::UniqueKeyConflict { id, initiated_by },            } => {                self.handle_failure_response(id, true, ctx);                ctx.emit(UpdateManagerEvent::ObjectOperationComplete {                    result: ObjectOperationResult {                        success_type: OperationSuccessType::Failure,                        operation: ObjectOperation::Create {                            initiated_by: *initiated_by,                        },                        client_id: ClientId::from_hash(id),                        server_id: None,                        num_objects: None,                    },                });            }            SyncQueueEvent::ObjectCreationFailure {                reason: CreationFailureReason::Other { id, initiated_by },            } => {                self.handle_failure_response(id, false, ctx);                ctx.emit(UpdateManagerEvent::ObjectOperationComplete {

Lines 1–500 of 4833.