chore: init local ai when switching workspace

frontend/appflowy_flutter/lib/workspace/presentation/settings/pages/setting_ai_view/local_ai_setting.dart

@@ -90,14 +90,18 @@ class LocalAiSettingHeader extends StatelessWidget {
         ),
         Toggle(
           value: isEnabled,
-          onChanged: (_) => _onToggleChanged(context),
+          onChanged: (value) {
+            _onToggleChanged(value, context);
+          },
         ),
       ],
     );
   }
 
-  void _onToggleChanged(BuildContext context) {
-    if (isEnabled) {
+  void _onToggleChanged(bool value, BuildContext context) {
+    if (value) {
+      context.read<LocalAiPluginBloc>().add(const LocalAiPluginEvent.toggle());
+    } else {
       showConfirmDialog(
         context: context,
         title: LocaleKeys.settings_aiPage_keys_disableLocalAITitle.tr(),
@@ -110,8 +114,6 @@ class LocalAiSettingHeader extends StatelessWidget {
               .add(const LocalAiPluginEvent.toggle());
         },
       );
-    } else {
-      context.read<LocalAiPluginBloc>().add(const LocalAiPluginEvent.toggle());
     }
   }
 }

frontend/rust-lib/flowy-ai/src/ai_manager.rs

@@ -12,7 +12,7 @@ use dashmap::DashMap;
 use flowy_ai_pub::cloud::{
   AIModel, ChatCloudService, ChatSettings, UpdateChatParams, DEFAULT_AI_MODEL_NAME,
 };
-use flowy_error::{FlowyError, FlowyResult};
+use flowy_error::{ErrorCode, FlowyError, FlowyResult};
 use flowy_sqlite::kv::KVStorePreferences;
 
 use crate::notification::{chat_notification_builder, ChatNotification};
@@ -104,36 +104,86 @@ impl AIManager {
     }
   }
 
-  #[instrument(skip_all, err)]
-  pub async fn initialize(&self, _workspace_id: &str) -> Result<(), FlowyError> {
-    let local_ai = self.local_ai.clone();
-    tokio::spawn(async move {
-      if let Err(err) = local_ai.destroy_plugin().await {
-        error!("Failed to destroy plugin: {}", err);
+  async fn reload_with_workspace_id(&self, workspace_id: &str) {
+    // Check if local AI is enabled for this workspace and if we're in local mode
+    let result = self.user_service.is_local_model().await;
+    if let Err(err) = &result {
+      if matches!(err.code, ErrorCode::UserNotLogin) {
+        info!("[AI Manager] User not logged in, skipping local AI reload");
+        return;
       }
+    }
 
-      if let Err(err) = local_ai.reload().await {
-        error!("[AI Manager] failed to reload local AI: {:?}", err);
-      }
-    });
+    let is_local = result.unwrap_or(false);
+    let is_enabled = self.local_ai.is_enabled_on_workspace(workspace_id);
+    let is_running = self.local_ai.is_running();
+    info!(
+      "[AI Manager] Reloading workspace: {}, is_local: {}, is_enabled: {}, is_running: {}",
+      workspace_id, is_local, is_enabled, is_running
+    );
+
+    // Shutdown AI if it's running but shouldn't be (not enabled and not in local mode)
+    if is_running && !is_enabled && !is_local {
+      info!("[AI Manager] Local AI is running but not enabled, shutting it down");
+      let local_ai = self.local_ai.clone();
+      tokio::spawn(async move {
+        // Wait for 5 seconds to allow other services to initialize
+        // TODO: pick a right time to start plugin service. Maybe [UserStatusCallback::did_launch]
+        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+
+        if let Err(err) = local_ai.toggle_plugin(false).await {
+          error!("[AI Manager] failed to shutdown local AI: {:?}", err);
+        }
+      });
+      return;
+    }
+
+    // Start AI if it's enabled but not running
+    if is_enabled && !is_running {
+      info!("[AI Manager] Local AI is enabled but not running, starting it now");
+      let local_ai = self.local_ai.clone();
+      tokio::spawn(async move {
+        // Wait for 5 seconds to allow other services to initialize
+        // TODO: pick a right time to start plugin service. Maybe [UserStatusCallback::did_launch]
+        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+
+        if let Err(err) = local_ai.toggle_plugin(true).await {
+          error!("[AI Manager] failed to start local AI: {:?}", err);
+        }
+      });
+      return;
+    }
+
+    // Log status for other cases
+    if is_running {
+      info!("[AI Manager] Local AI is already running");
+    }
+  }
+
+  #[instrument(skip_all, err)]
+  pub async fn on_launch_if_authenticated(&self, workspace_id: &str) -> Result<(), FlowyError> {
+    self.reload_with_workspace_id(workspace_id).await;
+    Ok(())
+  }
+
+  pub async fn initialize_after_sign_in(&self, workspace_id: &str) -> Result<(), FlowyError> {
+    self.reload_with_workspace_id(workspace_id).await;
+    Ok(())
+  }
+
+  pub async fn initialize_after_sign_up(&self, workspace_id: &str) -> Result<(), FlowyError> {
+    self.reload_with_workspace_id(workspace_id).await;
     Ok(())
   }
 
   #[instrument(skip_all, err)]
   pub async fn initialize_after_open_workspace(
     &self,
-    _workspace_id: &Uuid,
+    workspace_id: &Uuid,
   ) -> Result<(), FlowyError> {
-    let local_ai = self.local_ai.clone();
-    tokio::spawn(async move {
-      if let Err(err) = local_ai.destroy_plugin().await {
-        error!("Failed to destroy plugin: {}", err);
-      }
-
-      if let Err(err) = local_ai.reload().await {
-        error!("[AI Manager] failed to reload local AI: {:?}", err);
-      }
-    });
+    self
+      .reload_with_workspace_id(&workspace_id.to_string())
+      .await;
     Ok(())
   }
 

frontend/rust-lib/flowy-ai/src/local_ai/controller.rs

@@ -99,7 +99,7 @@ impl LocalAIController {
           continue;
         };
 
-        let key = local_ai_enabled_key(&workspace_id);
+        let key = local_ai_enabled_key(&workspace_id.to_string());
         info!("[AI Plugin] state: {:?}", state);
 
         // Read whether plugin is enabled from store; default to true
@@ -157,14 +157,15 @@ impl LocalAIController {
   }
   #[instrument(level = "debug", skip_all)]
   pub async fn observe_plugin_resource(&self) {
-    debug!(
-      "[AI Plugin] init plugin when first run. thread: {:?}",
-      std::thread::current().id()
-    );
     let sys = get_operating_system();
     if !sys.is_desktop() {
       return;
     }
+
+    debug!(
+      "[AI Plugin] observer plugin state. thread: {:?}",
+      std::thread::current().id()
+    );
     async fn try_init_plugin(
       resource: &Arc<LocalAIResourceController>,
       ai_plugin: &Arc<OllamaAIPlugin>,
@@ -196,12 +197,6 @@ impl LocalAIController {
     });
   }
 
-  pub async fn reload(&self) -> FlowyResult<()> {
-    let is_enabled = self.is_enabled();
-    self.toggle_plugin(is_enabled).await?;
-    Ok(())
-  }
-
   fn upgrade_store_preferences(&self) -> FlowyResult<Arc<KVStorePreferences>> {
     self
       .store_preferences
@@ -211,9 +206,6 @@ impl LocalAIController {
 
   /// Indicate whether the local AI plugin is running.
   pub fn is_running(&self) -> bool {
-    if !self.is_enabled() {
-      return false;
-    }
     self.ai_plugin.get_plugin_running_state().is_running()
   }
 
@@ -225,20 +217,25 @@ impl LocalAIController {
       return false;
     }
 
-    if let Ok(key) = self
-      .user_service
-      .workspace_id()
-      .map(|workspace_id| local_ai_enabled_key(&workspace_id))
-    {
-      match self.upgrade_store_preferences() {
-        Ok(store) => store.get_bool(&key).unwrap_or(false),
-        Err(_) => false,
-      }
+    if let Ok(workspace_id) = self.user_service.workspace_id() {
+      self.is_enabled_on_workspace(&workspace_id.to_string())
     } else {
       false
     }
   }
 
+  pub fn is_enabled_on_workspace(&self, workspace_id: &str) -> bool {
+    let key = local_ai_enabled_key(workspace_id);
+    if !get_operating_system().is_desktop() {
+      return false;
+    }
+
+    match self.upgrade_store_preferences() {
+      Ok(store) => store.get_bool(&key).unwrap_or(false),
+      Err(_) => false,
+    }
+  }
+
   pub fn get_plugin_chat_model(&self) -> Option<String> {
     if !self.is_enabled() {
       return None;
@@ -298,7 +295,8 @@ impl LocalAIController {
     );
 
     if self.resource.set_llm_setting(setting).await.is_ok() {
-      self.reload().await?;
+      let is_enabled = self.is_enabled();
+      self.toggle_plugin(is_enabled).await?;
     }
     Ok(())
   }
@@ -373,7 +371,7 @@ impl LocalAIController {
 
   pub async fn toggle_local_ai(&self) -> FlowyResult<bool> {
     let workspace_id = self.user_service.workspace_id()?;
-    let key = local_ai_enabled_key(&workspace_id);
+    let key = local_ai_enabled_key(&workspace_id.to_string());
     let store_preferences = self.upgrade_store_preferences()?;
     let enabled = !store_preferences.get_bool(&key).unwrap_or(true);
     store_preferences.set_bool(&key, enabled)?;
@@ -482,7 +480,7 @@ impl LocalAIController {
   }
 
   #[instrument(level = "debug", skip_all)]
-  async fn toggle_plugin(&self, enabled: bool) -> FlowyResult<()> {
+  pub(crate) async fn toggle_plugin(&self, enabled: bool) -> FlowyResult<()> {
     info!(
       "[AI Plugin] enable: {}, thread id: {:?}",
       enabled,
@@ -618,6 +616,6 @@ impl LLMResourceService for LLMResourceServiceImpl {
 }
 
 const APPFLOWY_LOCAL_AI_ENABLED: &str = "appflowy_local_ai_enabled";
-fn local_ai_enabled_key(workspace_id: &Uuid) -> String {
+fn local_ai_enabled_key(workspace_id: &str) -> String {
   format!("{}:{}", APPFLOWY_LOCAL_AI_ENABLED, workspace_id)
 }

frontend/rust-lib/flowy-core/src/user_state_callback.rs

@@ -38,15 +38,6 @@ pub(crate) struct UserStatusCallbackImpl {
 }
 
 impl UserStatusCallbackImpl {
-  fn init_ai_component(&self, workspace_id: String) {
-    let cloned_ai_manager = self.ai_manager.clone();
-    self.runtime.spawn(async move {
-      if let Err(err) = cloned_ai_manager.initialize(&workspace_id).await {
-        error!("Failed to initialize AIManager: {:?}", err);
-      }
-    });
-  }
-
   async fn folder_init_data_source(
     &self,
     user_id: i64,
@@ -95,7 +86,6 @@ impl UserStatusCallback for UserStatusCallbackImpl {
     auth_type: &AuthType,
   ) -> FlowyResult<()> {
     let workspace_id = user_workspace.workspace_id()?;
-
     if let Some(cloud_config) = cloud_config {
       self
         .server_provider
@@ -124,7 +114,15 @@ impl UserStatusCallback for UserStatusCallbackImpl {
     self.document_manager.initialize(user_id).await?;
 
     let workspace_id = user_workspace.id.clone();
-    self.init_ai_component(workspace_id);
+    let cloned_ai_manager = self.ai_manager.clone();
+    self.runtime.spawn(async move {
+      if let Err(err) = cloned_ai_manager
+        .on_launch_if_authenticated(&workspace_id)
+        .await
+      {
+        error!("Failed to initialize AIManager: {:?}", err);
+      }
+    });
     Ok(())
   }
 
@@ -158,8 +156,11 @@ impl UserStatusCallback for UserStatusCallbackImpl {
       .initialize_after_sign_in(user_id)
       .await?;
 
-    let workspace_id = user_workspace.id.clone();
-    self.init_ai_component(workspace_id);
+    self
+      .ai_manager
+      .initialize_after_sign_in(&user_workspace.id)
+      .await?;
+
     Ok(())
   }
 
@@ -207,8 +208,10 @@ impl UserStatusCallback for UserStatusCallbackImpl {
       .await
       .context("DocumentManager error")?;
 
-    let workspace_id = user_workspace.id.clone();
-    self.init_ai_component(workspace_id);
+    self
+      .ai_manager
+      .initialize_after_sign_up(&user_workspace.id)
+      .await?;
     Ok(())
   }
 

frontend/rust-lib/flowy-user/src/event_map.rs

@@ -291,6 +291,11 @@ pub trait UserStatusCallback: Send + Sync + 'static {
   ) -> FlowyResult<()> {
     Ok(())
   }
+
+  async fn did_launch(&self) -> FlowyResult<()> {
+    Ok(())
+  }
+
   /// Fires right after the user successfully signs in.
   async fn on_sign_in(
     &self,